From f0a138bb2daf384ab9d1cb8c54a8ee3e6f9aa3d6 Mon Sep 17 00:00:00 2001
From: AllenFarcas <alfarcas@amd.com>
Date: Mon, 10 Aug 2026 01:19:53 +0000
Subject: [PATCH 1/5] Enable upstream v2.17 features on ROCm and fix latent
 defects

Fixed:
- Blockwise FP8 GEMM accumulate=True. -ffast-math folded fpext(fptrunc(x))
  back to x, so the deliberate rounding before the beta*C add emitted no code
  and the epilogue accumulated in fp32. Routed through round_to_out_dtype(),
  which is opaque to the optimizer. Test wired into ci/pytorch.sh.
- gfx1250 TMA path selection: sm_arch() returns 125, which satisfied the
  upstream sm >= 100 gate and wrongly selected the TMA path.
- Float8CurrentScaling amax scratch was returned to the caching allocator
  while its kernels were still enqueued. Stream-ordered reuse covers that in
  eager mode, but graph capture replaces stream order with a DAG, so the
  freed block was aliased in the graph mempool and replays read partials they
  never wrote. Now a persistent per-(device, stream) buffer.
- Hipify ran at configure time with no CMake dependency on its inputs, so
  editing a source under common/ silently compiled the previous revision.

Enabled:
- NVFP4 4over6 quantization (gfx950).
- nvte_multi_tensor_transpose_to_bhsd, on the two non-TMA paths.
- Standalone TopK (nvte_topk).
- MXFP8 grouped dequantization (needs NVTE_ROCM_ENABLE_MXFP8=1). Shapes
  outside the supported set -- per-tensor rows not a multiple of 32 with
  columnwise scaling, or last dims not a multiple of 16 -- are rejected with
  a clear error rather than silently mis-scaled.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 ci/pytorch.sh                                 |   2 +
 tests/cpp/operator/CMakeLists.txt             |   1 +
 .../cpp/operator/test_cast_nvfp4_transpose.cu | 236 ++++++++++-
 .../operator/test_dequantize_mxfp8_grouped.cu |  84 +++-
 .../test_multi_tensor_transpose_bhsd.cu       | 145 +++++++
 tests/jax/test_custom_call_compute.py         |   1 -
 tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py  |   8 +-
 .../nvfp4/test_nvfp4_quantize_exact.py        |  53 ++-
 tests/pytorch/test_fusible_ops.py             |   3 +-
 transformer_engine/common/CMakeLists.txt      |  24 +-
 .../common/cast/dispatch/dequantize.cuh       |  10 +-
 .../common/cast/dispatch/quantize.cuh         |  28 +-
 .../cast/mxfp8/group_dequantize_mxfp8.cuh     | 155 ++++++-
 .../mxfp8/rocm_group_dequantize_mxfp8.cuh     | 212 ++++++++++
 .../common/fused_attn/flash_attn.cu           |  29 +-
 .../common/gemm/kittens/CMakeLists.txt        |   3 +
 .../cdna3/blockwise_fp8_gemm_helper.cuh       |  25 +-
 .../cdna4/blockwise_fp8_gemm_helper.cuh       |  18 +-
 transformer_engine/common/gemm/rocm_gemm.cu   |   6 +-
 .../common/transpose/cast_transpose.h         |   5 +-
 ...quantize_transpose_vector_blockwise_fp4.cu | 393 +++++++++++++++---
 .../common/util/standalone_topk.cuh           |  17 +
 transformer_engine/common/util/topk.cu        |   8 +
 .../jax/csrc/extensions/pybind.cpp            |   4 -
 .../jax/csrc/extensions/topk.cpp              |   7 -
 transformer_engine/pytorch/csrc/common.cpp    |  29 +-
 26 files changed, 1362 insertions(+), 144 deletions(-)
 create mode 100644 tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu
 create mode 100644 transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh

diff --git a/ci/pytorch.sh b/ci/pytorch.sh
index 38cef08ce6..993d1cea95 100755
--- a/ci/pytorch.sh
+++ b/ci/pytorch.sh
@@ -66,6 +66,7 @@ run_test_config(){
     run_default_fa 1 test_float8_current_scaling_exact.py
     run_default_fa 1 test_float8blockwisetensor.py
     run_default_fa 1 test_float8_blockwise_scaling_exact.py
+    run_default_fa 1 test_float8_blockwise_gemm_exact.py
     run_default_fa 1 test_quantized_tensor.py
     test $_fus_attn = auto -o $_fus_attn = ck && run 1 test_cpu_offloading.py
     test $_fus_attn = auto -o $_fus_attn = ck -o $_fus_attn = aotriton && NVTE_FLASH_ATTN=0 NVTE_CPU_OFFLOAD_V1=1 run 3 test_cpu_offloading_v1.py
@@ -75,6 +76,7 @@ run_test_config(){
     run_default_fa 1 test_gemm_autotune.py
     run 1 test_gqa.py
     run 1 test_grouped_linear.py
+    NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 test_grouped_tensor.py
     run 1 test_jit.py
     NVTE_ROCM_ENABLE_MXFP8=1 run_default_fa 1 test_multi_tensor.py
     run 1 test_numerics.py
diff --git a/tests/cpp/operator/CMakeLists.txt b/tests/cpp/operator/CMakeLists.txt
index a13d4f37dd..ecaaa83e9d 100644
--- a/tests/cpp/operator/CMakeLists.txt
+++ b/tests/cpp/operator/CMakeLists.txt
@@ -46,6 +46,7 @@ if(USE_ROCM)
   list(REMOVE_ITEM test_cuda_sources
               test_grouped_gemm.cu)
   list(APPEND test_cuda_sources
+              test_multi_tensor_transpose_bhsd.cu
               test_dequantize_nvfp4.cu
               test_cublaslt_gemm.cu
               test_cast_mxfp4_transpose.cu
diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu
index 27eb7ad51e..30aea9d33b 100644
--- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu
+++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu
@@ -8,6 +8,8 @@
 
 #include <algorithm>
 #include <array>
+#include <cfloat>
+#include <cmath>
 #include <fstream>
 #include <iostream>
 #include <memory>
@@ -188,6 +190,146 @@ NVFP4FourOverSixQuantization quantize_4over6_pair(
   };
 }
 
+// CUDA keeps upstream's looser rule; its fast-math configs score in packed fp16.
+#ifdef __HIP_PLATFORM_AMD__
+constexpr bool kAcceptEitherCandidate = false;
+#else
+constexpr bool kAcceptEitherCandidate = true;
+#endif
+
+struct NVFP4FourOverSixDecision {
+    NVFP4FourOverSixCandidate candidate = NVFP4FourOverSixCandidate::Map6;
+    // Candidate errors too close for FP32 summation order to resolve.
+    bool ambiguous = false;
+};
+
+// Mirrors fp4_roundtrip_err() in quantize_transpose_vector_blockwise_fp4.cu,
+// operand order included.
+float fp4_roundtrip_err_ref(const float x, const float block_scale_inverse, const float sf,
+                            const float global_amax, const float err_denom, const bool use_mse) {
+    const float2 scaled = {x * block_scale_inverse, 0.0f};
+    const fp4e2m1x2 quantized(scaled);
+    const float dequant = static_cast<float>(cvt_fp4x2_to_double2(quantized).x);
+    const float val = dequant * sf * global_amax / err_denom;
+    const float diff = val - x;
+    return use_mse ? diff * diff : std::fabs(diff);
+}
+
+// Lower round-trip error wins; ties go to map-to-6.
+NVFP4FourOverSixDecision decide_4over6_candidate(const std::vector<float>& block,
+                                                 const float block_amax,
+                                                 const float S_enc,
+                                                 const int e4m3_max,
+                                                 const bool use_mse) {
+    const NVFP4FourOverSixQuantization quantization =
+        compute_4over6_quantization_scales(block_amax, S_enc);
+    const float err_denom = 6.0f * static_cast<float>(e4m3_max);
+    const float global_amax = err_denom / S_enc;
+
+    double err_map4 = 0.0;
+    double err_map6 = 0.0;
+    for (const float x : block) {
+        err_map4 += fp4_roundtrip_err_ref(x, quantization.reciprocal_map4,
+                                          static_cast<float>(quantization.scale_map4), global_amax,
+                                          err_denom, use_mse);
+        err_map6 += fp4_roundtrip_err_ref(x, quantization.reciprocal_map6,
+                                          static_cast<float>(quantization.scale_map6), global_amax,
+                                          err_denom, use_mse);
+    }
+
+    // Worst-case FP32 summation drift of the kernel's warp reduction.
+    const double tie_band = 4.0 * static_cast<double>(block.size()) *
+                            static_cast<double>(FLT_EPSILON) * std::max(err_map4, err_map6);
+
+    NVFP4FourOverSixDecision decision;
+    decision.ambiguous = std::fabs(err_map4 - err_map6) <= tie_band;
+    decision.candidate =
+        err_map4 < err_map6 ? NVFP4FourOverSixCandidate::Map4 : NVFP4FourOverSixCandidate::Map6;
+    return decision;
+}
+
+// Laid out like the scale tensor so the checker can index it by scale_idx.
+template <typename InputType>
+std::vector<NVFP4FourOverSixDecision> compute_4over6_expected_decisions(
+    float (*OP)(const float),
+    const InputType* const input,
+    const size_t rows,
+    const size_t cols,
+    const size_t scales_stride,
+    const float* const amax,
+    const bool use_fast_math,
+    const bool use_2d_quantization,
+    const bool row_scaled_nvfp4,
+    const bool use_mse,
+    const int e4m3_max) {
+
+    constexpr size_t block_size_Y = 16;
+    constexpr size_t block_size_X = 16;
+    const size_t blocks_X = divide_round_up(cols, block_size_X);
+    std::vector<NVFP4FourOverSixDecision> decisions(rows * scales_stride);
+
+    // Same numerical truncation the reference quantizers apply.
+    auto activated = [OP, input, cols](const size_t i, const size_t j) {
+        const float act_elt = OP(static_cast<float>(input[i * cols + j]));
+        return static_cast<float>(static_cast<InputType>(act_elt));
+    };
+
+    if (use_2d_quantization) {
+        const float S_enc = compute_global_encode_scaling_factor_FP4(*amax, use_fast_math,
+                                                                     e4m3_max);
+        const size_t blocks_Y = divide_round_up(rows, block_size_Y);
+        for (size_t block_Y = 0; block_Y < blocks_Y; ++block_Y) {
+            for (size_t block_X = 0; block_X < blocks_X; ++block_X) {
+                const size_t i_min = block_Y * block_size_Y;
+                const size_t i_max = std::min(i_min + block_size_Y, rows);
+                const size_t j_min = block_X * block_size_X;
+                const size_t j_max = std::min(j_min + block_size_X, cols);
+
+                std::vector<float> block;
+                block.reserve(block_size_Y * block_size_X);
+                float block_amax = 0.0f;
+                for (size_t i = i_min; i < i_max; ++i) {
+                    for (size_t j = j_min; j < j_max; ++j) {
+                        const float elt = activated(i, j);
+                        block.push_back(elt);
+                        block_amax = std::max(block_amax, std::abs(elt));
+                    }
+                }
+
+                const NVFP4FourOverSixDecision decision =
+                    decide_4over6_candidate(block, block_amax, S_enc, e4m3_max, use_mse);
+                // Block scale and its candidate replicate down the block's rows.
+                for (size_t i = i_min; i < i_max; ++i) {
+                    decisions[i * scales_stride + block_X] = decision;
+                }
+            }
+        }
+        return decisions;
+    }
+
+    for (size_t i = 0; i < rows; ++i) {
+        const float S_enc = compute_global_encode_scaling_factor_FP4(
+            row_scaled_nvfp4 ? amax[i] : *amax, use_fast_math, e4m3_max);
+        for (size_t block_X = 0; block_X < blocks_X; ++block_X) {
+            const size_t j_min = block_X * block_size_X;
+            const size_t j_max = std::min(j_min + block_size_X, cols);
+
+            std::vector<float> block;
+            block.reserve(block_size_X);
+            float block_amax = 0.0f;
+            for (size_t j = j_min; j < j_max; ++j) {
+                const float elt = activated(i, j);
+                block.push_back(elt);
+                block_amax = std::max(block_amax, std::abs(elt));
+            }
+
+            decisions[i * scales_stride + block_X] =
+                decide_4over6_candidate(block, block_amax, S_enc, e4m3_max, use_mse);
+        }
+    }
+    return decisions;
+}
+
 // 1D Scaling: Original implementation with 1x16 blocks
 template <typename InputType>
 void quantize_nvfp4_1d(float (*OP)(const float),
@@ -735,6 +877,7 @@ void compare_nvfp4_4over6_candidates(const std::string& name,
                                      const fp8e4m3* const ref_scales_map4,
                                      const fp4e2m1x2* const ref_data_map6,
                                      const fp8e4m3* const ref_scales_map6,
+                                     const std::vector<NVFP4FourOverSixDecision>& decisions,
                                      const size_t rows,
                                      const size_t cols,
                                      const size_t blocks_X,
@@ -742,6 +885,16 @@ void compare_nvfp4_4over6_candidates(const std::string& name,
     constexpr int max_mismatches_to_print = 3;
     const auto* const test_data_pairs = reinterpret_cast<const fp4e2m1x2*>(test_data);
     size_t total_mismatches = 0;
+    size_t wrong_candidate = 0;
+    size_t expected_map4 = 0;
+    size_t expected_map6 = 0;
+    // Blocks whose two candidates encode to the same bytes: a selection cannot be
+    // read back out of them, so they neither pin the kernel nor weaken the check.
+    size_t identical_candidates = 0;
+    // Blocks that do encode differently but whose candidate errors are too close
+    // for the host to call. These are the only ones where the kernel keeps any
+    // latitude, so their count is the coverage this check gives up.
+    size_t unpinned_blocks = 0;
 
     for (size_t row = 0; row < rows; ++row) {
         for (size_t block_x = 0; block_x < blocks_X; ++block_x) {
@@ -754,34 +907,72 @@ void compare_nvfp4_4over6_candidates(const std::string& name,
                 bitwise_equal(test_scales[scale_idx], ref_scales_map6[scale_idx]);
             const bool data_matches_map6 =
                 nvfp4_output_block_matches(test_data_pairs, ref_data_map6, row, cols, block_x);
+            const bool matches_map4 = scale_matches_map4 && data_matches_map4;
+            const bool matches_map6 = scale_matches_map6 && data_matches_map6;
+
+            const bool candidates_differ =
+                !bitwise_equal(ref_scales_map4[scale_idx], ref_scales_map6[scale_idx]) ||
+                !nvfp4_output_block_matches(ref_data_map4, ref_data_map6, row, cols, block_x);
+
+            const NVFP4FourOverSixDecision& decision = decisions[scale_idx];
+            bool matched = false;
+            const char* expectation = nullptr;
+            if (decision.ambiguous || kAcceptEitherCandidate) {
+                // Too close to call from the host; either encoding is a correct answer.
+                if (candidates_differ) {
+                    ++unpinned_blocks;
+                } else {
+                    ++identical_candidates;
+                }
+                matched = matches_map4 || matches_map6;
+                expectation = "map-to-4 or map-to-6 (candidate errors within the tie band)";
+            } else if (decision.candidate == NVFP4FourOverSixCandidate::Map4) {
+                ++expected_map4;
+                matched = matches_map4;
+                expectation = "map-to-4";
+            } else {
+                ++expected_map6;
+                matched = matches_map6;
+                expectation = "map-to-6";
+            }
 
-            if ((scale_matches_map4 && data_matches_map4) ||
-                (scale_matches_map6 && data_matches_map6)) {
+            if (matched) {
                 continue;
             }
 
             ++total_mismatches;
+            const bool matched_other = matches_map4 || matches_map6;
+            if (matched_other) {
+                ++wrong_candidate;
+            }
             if (total_mismatches <= max_mismatches_to_print) {
                 std::cout << "Error in tensor " << name << ": 4over6 block mismatch at row "
-                          << row << ", block_x " << block_x
-                          << ". The output did not match either map-to-4 or map-to-6 exactly."
-                          << std::endl;
+                          << row << ", block_x " << block_x << ". Expected " << expectation
+                          << "; the output "
+                          << (matched_other ? "matched the other candidate instead"
+                                            : "matched neither candidate exactly")
+                          << "." << std::endl;
             }
         }
     }
 
     std::cout << "=== SUMMARY for tensor " << name << " ===" << std::endl;
     std::cout << "Total 4over6 blocks checked: " << (rows * blocks_X) << std::endl;
+    std::cout << "Blocks pinned to map-to-4: " << expected_map4
+              << ", pinned to map-to-6: " << expected_map6
+              << ", left unpinned by the tie band: " << unpinned_blocks
+              << ", encoding identically either way: " << identical_candidates << std::endl;
     if (total_mismatches > 0) {
         std::cout << "STATUS: FAILED for output" << std::endl;
-        std::cout << "Total mismatched 4over6 blocks found: " << total_mismatches << std::endl;
+        std::cout << "Total mismatched 4over6 blocks found: " << total_mismatches
+                  << " (" << wrong_candidate << " picked the wrong candidate)" << std::endl;
         std::cout << "============================" << std::endl;
         GTEST_FAIL() << "Found " << total_mismatches << " 4over6 block mismatches in tensor "
-                     << name;
+                     << name << " (" << wrong_candidate << " picked the wrong candidate)";
     }
 
     std::cout << "STATUS: PASSED for output" << std::endl;
-    std::cout << "Each 4over6 block matched either map-to-4 or map-to-6 exactly" << std::endl;
+    std::cout << "Each 4over6 block matched the candidate the reference selected" << std::endl;
     std::cout << "============================" << std::endl;
 }
 
@@ -823,7 +1014,8 @@ void performTest(float (*OP)(const float),
 
 #ifdef __HIP_PLATFORM_AMD__
     if (te_fp8_fnuz()) GTEST_SKIP() << "NVFP4 not supported on gfx942 (fnuz)";
-    if (use_4over6) GTEST_SKIP() << "NVFP4 4over6 not supported on ROCm";
+    if (use_4over6 && use_4over6_err_use_fast_math)
+      GTEST_SKIP() << "NVFP4 4over6 fast-math error mode is not supported on ROCm";
 #endif
 
     const size_t rows = first_dimension(shape);
@@ -863,6 +1055,8 @@ void performTest(float (*OP)(const float),
     std::unique_ptr<fp4e2m1x2[]> ref_output_t_map6;
     std::unique_ptr<fp8e4m3[]> ref_scales_map6;
     std::unique_ptr<fp8e4m3[]> ref_scales_t_map6;
+    std::vector<NVFP4FourOverSixDecision> expected_decisions;
+    std::vector<NVFP4FourOverSixDecision> expected_decisions_t;
 
     fillCase<fp32>(&input, InputsFillCase::uniform);
 
@@ -974,6 +1168,21 @@ void performTest(float (*OP)(const float),
                                use_4over6,
                                e4m3_max,
                                NVFP4FourOverSixCandidate::Map6);
+
+        // The kernel is free to pick either candidate per block, but not freely:
+        // it must pick the one with the smaller round-trip error. Derive that
+        // choice here so the checker can hold it to exactly one reference.
+        const bool use_mse = mode == kNVTENVFP44Over6MinMSE;
+        expected_decisions = compute_4over6_expected_decisions<InputType>(
+            OP, input.rowwise_cpu_dptr<InputType>(), rows, cols, scales_stride, ref_amax.data(),
+            use_fast_math, is_2d_quantization, row_scaled_nvfp4, use_mse, e4m3_max);
+        if (!row_scaled_nvfp4) {
+            const std::vector<InputType> input_t =
+                create_transpose(input.rowwise_cpu_dptr<InputType>(), rows, cols);
+            expected_decisions_t = compute_4over6_expected_decisions<InputType>(
+                OP, input_t.data(), cols, rows, scales_stride_t, ref_amax.data(), use_fast_math,
+                is_2d_quantization, row_scaled_nvfp4, use_mse, e4m3_max);
+        }
     } else {
         compute_ref<InputType>(OP,
                                input.rowwise_cpu_dptr<InputType>(),
@@ -997,8 +1206,11 @@ void performTest(float (*OP)(const float),
     hipDeviceProp_t prop;
     hipGetDeviceProperties(&prop, 0);
     const bool is_gfx950 = prop.major == 9 && prop.minor == 5;
-    for (bool use_stochastic_rounding : (is_gfx950 ? std::vector<bool>{false, true}
-                                                   : std::vector<bool>{false})) {
+    // 4over6 excluded: the quantize dispatch refuses it combined with stochastic
+    // rounding, the same way this test already excludes it from fast math below.
+    for (bool use_stochastic_rounding : (is_gfx950 && !use_4over6
+                                             ? std::vector<bool>{false, true}
+                                             : std::vector<bool>{false})) {
 #endif
 
     // Initialize stochastic rounding
@@ -1056,6 +1268,7 @@ void performTest(float (*OP)(const float),
                                         ref_scales.get(),
                                         ref_output_map6.get(),
                                         ref_scales_map6.get(),
+                                        expected_decisions,
                                         rows,
                                         cols,
                                         unpadded_blocks_X,
@@ -1068,6 +1281,7 @@ void performTest(float (*OP)(const float),
                                             ref_scales_t.get(),
                                             ref_output_t_map6.get(),
                                             ref_scales_t_map6.get(),
+                                            expected_decisions_t,
                                             cols,
                                             rows,
                                             unpadded_blocks_X_t,
diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
index 4a18bb5891..e13eb7ab4d 100644
--- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
+++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.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.
@@ -229,6 +231,40 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
                                   &offsets_tensor, sizeof(offsets_tensor));
   }
 
+  bool bad_last_dim = false;
+  for (size_t t = 0; t < num_tensors; ++t) {
+    if (last_dims_h[t] % 16 != 0) bad_last_dim = true;
+  }
+#ifdef __HIP_PLATFORM_AMD__
+  auto free_device = [&]() {
+    cudaFree(in_data_d);
+    cudaFree(out_grouped_d);
+    cudaFree(in_scales_d);
+    cudaFree(first_dims_d);
+    cudaFree(last_dims_d);
+    cudaFree(offsets_d);
+  };
+
+  if (bad_last_dim) {
+    EXPECT_THROW(nvte_group_dequantize(in_group_tensor, out_group_tensor, 0), std::runtime_error);
+    free_device();
+    return;
+  }
+
+  const bool uniform_last_dim = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM);
+  bool ragged_colwise = false;
+  if (!rowwise && uniform_last_dim && num_tensors > 1) {
+    for (size_t t = 0; t < num_tensors; ++t) {
+      if (first_dims_h[t] % 32 != 0) ragged_colwise = true;
+    }
+  }
+  if (ragged_colwise) {
+    EXPECT_THROW(nvte_group_dequantize(in_group_tensor, out_group_tensor, 0), std::runtime_error);
+    free_device();
+    return;
+  }
+#endif
+
   // Run grouped dequantize
   nvte_group_dequantize(in_group_tensor, out_group_tensor, 0);
   cudaDeviceSynchronize();
@@ -310,15 +346,16 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
 
     int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset,
                         tensor_elts * sizeof(OutputType));
-    if (result != 0) {
-      // Find first mismatch for error reporting
-      for (size_t i = 0; i < tensor_elts; ++i) {
-        if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) {
-          GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i
-                       << " (global offset " << (data_offset + i) << "): grouped="
-                       << static_cast<float>(out_grouped_h[data_offset + i])
-                       << " vs reference=" << static_cast<float>(out_ref_h[data_offset + i]);
-        }
+    // memcmp is the assertion; the loop below only locates the mismatch. Comparing with != as
+    // the gate would let -0.0 vs +0.0 and distinct NaN payloads pass after memcmp caught them,
+    // and the e8m0 scale bytes are drawn across the full [0,255] range including NaN.
+    ASSERT_EQ(result, 0) << "Bitwise mismatch in tensor " << t;
+    for (size_t i = 0; i < tensor_elts; ++i) {
+      if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) {
+        GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i
+                     << " (global offset " << (data_offset + i) << "): grouped="
+                     << static_cast<float>(out_grouped_h[data_offset + i])
+                     << " vs reference=" << static_cast<float>(out_ref_h[data_offset + i]);
       }
     }
   }
@@ -349,6 +386,18 @@ std::vector<std::vector<size_t>> input_configs = {
     {VARYING_FIRST_DIM, 3, 768, 96, 256, 256, 256},
     {VARYING_LAST_DIM, 2, 160, 384, 128, 256},
     {VARYING_LAST_DIM, 3, 96, 512, 128, 128, 256},
+#ifdef __HIP_PLATFORM_AMD__
+    // Per-tensor rows not scale-block aligned: rejected on ROCm.
+    {VARYING_FIRST_DIM, 2, 96, 128, 48, 48},
+    // Last dim not 16-byte aligned: vectorized loads would shift the tile.
+    {SAME_BOTH_DIMS, 1, 128, 120},
+    // Launcher derives fewer blocks than the tile count, so the grid-stride loop must iterate.
+    {VARYING_LAST_DIM, 2, 256, 256, 96, 160},
+    // Rows are scale-block aligned but not chunk aligned: valid, must not be rejected.
+    {VARYING_FIRST_DIM, 2, 128, 128, 64, 64},
+    // Uniform shape whose per-tensor rows are not scale-block aligned.
+    {SAME_BOTH_DIMS, 2, 96, 128},
+#endif
 };
 
 std::vector<ScalingDirection> scaling_directions = {
@@ -366,10 +415,12 @@ class GroupedDequantizeMXFP8TestSuite
                                                  >> {};
 
 TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) {
+#ifndef __HIP_PLATFORM_AMD__
   // Skip tests for pre-Blackwell architectures
   if (getDeviceComputeCapability() < blackwellComputeCapability) {
     GTEST_SKIP();
   }
+#endif
 
   using namespace transformer_engine;
   using namespace test;
@@ -419,6 +470,19 @@ TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) {
         (shape_rep == VARYING_FIRST_DIM || shape_rep == VARYING_BOTH_DIMS);
     const bool last_dim_varies =
         (shape_rep == VARYING_LAST_DIM || shape_rep == VARYING_BOTH_DIMS);
+#ifdef __HIP_PLATFORM_AMD__
+    // ROCm rejects the shapes below instead of skipping them; see the EXPECT_THROW paths.
+    const bool uniform_last = (shape_rep == SAME_BOTH_DIMS || shape_rep == VARYING_FIRST_DIM);
+    if (first_dim_varies && shape_rep != VARYING_FIRST_DIM && (first_dims[t] % 128 != 0)) {
+      GTEST_SKIP();
+    }
+    if (last_dim_varies && shape_rep != VARYING_LAST_DIM && (last_dims[t] % 128 != 0)) {
+      GTEST_SKIP();
+    }
+    if (!rowwise && !uniform_last && (first_dims[t] % 32 != 0)) {
+      GTEST_SKIP();
+    }
+#else
     if (first_dim_varies && (first_dims[t] % 128 != 0)) {
       GTEST_SKIP();
     }
@@ -429,10 +493,10 @@ TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) {
     if (last_dims[t] % 16 != 0) {
       GTEST_SKIP();
     }
-    // For colwise: first dim must be divisible by 32
     if (!rowwise && (first_dims[t] % 32 != 0)) {
       GTEST_SKIP();
     }
+#endif
     // For rowwise: last dim must be divisible by 32
     if (rowwise && (last_dims[t] % 32 != 0)) {
       GTEST_SKIP();
diff --git a/tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu b/tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu
new file mode 100644
index 0000000000..b859108f71
--- /dev/null
+++ b/tests/cpp/operator/test_multi_tensor_transpose_bhsd.cu
@@ -0,0 +1,145 @@
+/*************************************************************************
+ * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * License for AMD contributions = MIT. See LICENSE for more information
+ ************************************************************************/
+
+#include <cuda_runtime.h>
+#include <gtest/gtest.h>
+#include <transformer_engine/fused_attn.h>
+
+#include <cstdint>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include "../test_common.h"
+
+using namespace transformer_engine;
+using namespace test;
+
+namespace {
+
+struct PermuteCase {
+  size_t B, S, H, D_in, D_out;
+};
+
+// out[b][h][s][d] = in[b][s][h][d] (BSHD) or in[s][b][h][d] (SBHD), zero-padded
+// past D_in when D_out is larger.
+template <typename T>
+void compute_ref(const T *in, T *out, const PermuteCase &c, bool is_bshd) {
+  for (size_t b = 0; b < c.B; ++b) {
+    for (size_t h = 0; h < c.H; ++h) {
+      for (size_t s = 0; s < c.S; ++s) {
+        const size_t in_off =
+            (is_bshd ? ((b * c.S + s) * c.H + h) : ((s * c.B + b) * c.H + h)) * c.D_in;
+        T *dst = out + ((b * c.H + h) * c.S + s) * c.D_out;
+        for (size_t d = 0; d < c.D_out; ++d) {
+          dst[d] = (d < c.D_in) ? in[in_off + d] : static_cast<T>(0);
+        }
+      }
+    }
+  }
+}
+
+// Shared-memory transpose needs 32*(32*D_pad+4) bytes; query the device budget.
+template <typename T>
+bool exceeds_smem_budget(const std::vector<PermuteCase> &cases) {
+  size_t d_in_max = 0;
+  for (const auto &c : cases) d_in_max = std::max(d_in_max, c.D_in);
+  const size_t d_bytes = d_in_max * sizeof(T);
+  if (d_bytes % 4 == 0) return false;  // vectorized path, no shared memory
+  const size_t d_pad = (d_bytes + 3u) & ~size_t(3);
+  const size_t needed = 32 * (32 * d_pad + 4);
+  int max_smem = 0;
+  NVTE_CHECK_CUDA(cudaDeviceGetAttribute(&max_smem, cudaDevAttrMaxSharedMemoryPerBlockOptin, 0));
+  return needed > static_cast<size_t>(max_smem);
+}
+
+// D_in * sizeof(T) % 4 selects vectorized vs shared-memory path; TMA is CUDA-only.
+template <typename T>
+void performTest(const std::vector<PermuteCase> &cases, bool is_bshd, DType dtype) {
+  if (exceeds_smem_budget<T>(cases)) {
+    GTEST_SKIP() << "shared-memory transpose exceeds this device's shared memory";
+  }
+  const size_t n = cases.size();
+  std::vector<std::unique_ptr<Tensor>> ins, outs;
+  std::vector<NVTETensor> in_h, out_h;
+  std::vector<std::vector<T>> refs(n);
+
+  for (size_t i = 0; i < n; ++i) {
+    const auto &c = cases[i];
+    const std::vector<size_t> in_shape = is_bshd ? std::vector<size_t>{c.B, c.S, c.H, c.D_in}
+                                                 : std::vector<size_t>{c.S, c.B, c.H, c.D_in};
+    const std::vector<size_t> out_shape{c.B, c.H, c.S, c.D_out};
+
+    auto in = std::make_unique<Tensor>("in_" + std::to_string(i), in_shape, dtype);
+    auto out = std::make_unique<Tensor>("out_" + std::to_string(i), out_shape, dtype);
+    fillUniform(in.get());
+    // Poison output so unwritten positions cannot match the expected zero pad.
+    NVTE_CHECK_CUDA(cudaMemset(out->rowwise_dptr(), 0xCD, c.B * c.H * c.S * c.D_out * sizeof(T)));
+
+    in->to_cpu();
+    refs[i].resize(c.B * c.H * c.S * c.D_out);
+    compute_ref<T>(in->rowwise_cpu_dptr<T>(), refs[i].data(), c, is_bshd);
+
+    in_h.push_back(in->data());
+    out_h.push_back(out->data());
+    ins.emplace_back(std::move(in));
+    outs.emplace_back(std::move(out));
+  }
+
+  nvte_multi_tensor_transpose_to_bhsd(
+      in_h.data(), out_h.data(), n,
+      is_bshd ? NVTE_QKV_Format::NVTE_BSHD : NVTE_QKV_Format::NVTE_SBHD, 0);
+  NVTE_CHECK_CUDA(cudaDeviceSynchronize());
+  NVTE_CHECK_CUDA(cudaGetLastError());
+
+  // A pure layout permute, so byte equality is the right assertion.
+  for (size_t i = 0; i < n; ++i) {
+    outs[i]->to_cpu();
+    compareResults("bhsd_" + std::to_string(i),
+                   reinterpret_cast<const uint8_t *>(outs[i]->rowwise_cpu_dptr<T>()),
+                   reinterpret_cast<const uint8_t *>(refs[i].data()), refs[i].size() * sizeof(T));
+  }
+}
+
+class MultiTensorTransposeBhsdTestSuite
+    : public ::testing::TestWithParam<std::tuple<std::vector<PermuteCase>, bool>> {};
+
+TEST_P(MultiTensorTransposeBhsdTestSuite, TestFp16) {
+  const auto cases = std::get<0>(GetParam());
+  const bool is_bshd = std::get<1>(GetParam());
+  performTest<fp16>(cases, is_bshd, DType::kFloat16);
+}
+
+TEST_P(MultiTensorTransposeBhsdTestSuite, TestByte) {
+  const auto cases = std::get<0>(GetParam());
+  const bool is_bshd = std::get<1>(GetParam());
+  performTest<byte>(cases, is_bshd, DType::kByte);
+}
+
+INSTANTIATE_TEST_SUITE_P(
+    OperatorTest, MultiTensorTransposeBhsdTestSuite,
+    ::testing::Combine(::testing::Values(
+                           // Single tensor, D_in*elem a multiple of 4 -> vectorized fallback.
+                           std::vector<PermuteCase>{{2, 64, 4, 64, 64}},
+                           // Odd D_in with 1-byte elements: shared-memory transpose.
+                           std::vector<PermuteCase>{{2, 48, 3, 33, 33}},
+                           // D_out > D_in exercises the zero pad on the vectorized path.
+                           std::vector<PermuteCase>{{1, 32, 2, 40, 64}},
+                           // Same pad via the shared-memory transpose, within gfx942 budget.
+                           std::vector<PermuteCase>{{1, 40, 2, 21, 32}},
+                           // Multiple tensors per launch; B must match across the group.
+                           std::vector<PermuteCase>{
+                               {2, 64, 4, 64, 64}, {2, 33, 2, 16, 24}, {2, 16, 1, 8, 8}},
+                           // S below one tile and H = 1, the degenerate grid.
+                           std::vector<PermuteCase>{{1, 1, 1, 16, 16}}),
+                       ::testing::Bool()),
+    [](const testing::TestParamInfo<MultiTensorTransposeBhsdTestSuite::ParamType> &info) {
+      return "case" + std::to_string(info.index) + "_n" +
+             std::to_string(std::get<0>(info.param).size()) +
+             (std::get<1>(info.param) ? "_bshd" : "_sbhd");
+    });
+
+}  // namespace
diff --git a/tests/jax/test_custom_call_compute.py b/tests/jax/test_custom_call_compute.py
index b86091df4d..6144a484fc 100644
--- a/tests/jax/test_custom_call_compute.py
+++ b/tests/jax/test_custom_call_compute.py
@@ -2108,7 +2108,6 @@ def f(x):
         assert_allclose(actual, expected, dtype=dtype)
 
 
-@pytest.mark.skipif(is_hip_extension(), reason="Standalone TopK (nvte_topk) is not supported on ROCm")
 @pytest.mark.parametrize("dtype", [jnp.bfloat16, jnp.float32])
 @pytest.mark.parametrize(
     "problem_size",
diff --git a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py
index d7e02c55c7..b3ef196bf3 100644
--- a/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py
+++ b/tests/pytorch/nvfp4/test_nvfp4_gemm_exact.py
@@ -39,8 +39,6 @@ def check_nvfp4_gemm_versus_reference(
 ):
     if nvfp4_e4m3_max != 448 and not use_4over6:
         pytest.skip("E4M3 max 256 is only meaningful for 4over6")
-    if IS_HIP_EXTENSION and use_4over6:
-        pytest.skip("NVFP4 4over6 is not supported on ROCm")
     te_dtype = te.DType.kFloat4E2M1
 
     # Setup device and random seed
@@ -467,7 +465,7 @@ def check_nvfp4_row_scaled_gemm_matches_emulated(
     ids=["rowxrow", "colxrow", "colxcol"],
 )
 @pytest.mark.parametrize("row_scaled_nvfp4", [False, True], ids=["nvfp4", "nvfp4_row_scaled"])
-@pytest.mark.parametrize("use_4over6", [False] if IS_HIP_EXTENSION else [False, True], ids=["default"] if IS_HIP_EXTENSION else ["default", "4over6"])
+@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"])
 @pytest.mark.parametrize("nvfp4_e4m3_max", [448, 256], ids=["e4m3_448", "e4m3_256"])
 @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"])
 def test_nvfp4_gemm_versus_reference(
@@ -530,7 +528,7 @@ def test_nvfp4_gemm_versus_reference(
 @pytest.mark.parametrize("out_dtype", [torch.float32, torch.bfloat16], ids=str)
 @pytest.mark.parametrize("use_bias", [False, True], ids=["no_bias", "bias"])
 @pytest.mark.parametrize("single_output", [False, True], ids=["list_output", "single_output"])
-@pytest.mark.parametrize("use_4over6", [False] if IS_HIP_EXTENSION else [False, True], ids=["default"] if IS_HIP_EXTENSION else ["default", "4over6"])
+@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"])
 @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"])
 @pytest.mark.skipif(IS_HIP_EXTENSION, reason="Grouped NVFP4 GEMM is not supported on ROCm")
 def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm(
@@ -579,7 +577,7 @@ def test_nvfp4_row_scaled_grouped_gemm_matches_per_gemm(
 @pytest.mark.parametrize("x_dtype", [torch.float32, torch.bfloat16], ids=str)
 @pytest.mark.parametrize("w_dtype", [torch.float32, torch.bfloat16], ids=str)
 @pytest.mark.parametrize("out_dtype", [torch.bfloat16, torch.float32], ids=str)
-@pytest.mark.parametrize("use_4over6", [False] if IS_HIP_EXTENSION else [False, True], ids=["default"] if IS_HIP_EXTENSION else ["default", "4over6"])
+@pytest.mark.parametrize("use_4over6", [False, True], ids=["default", "4over6"])
 @pytest.mark.parametrize("nvfp4_4over6_err_mode", ["MAE", "MSE"], ids=["mae_err", "mse_err"])
 def test_nvfp4_row_scaled_gemm_matches_emulated(
     M: int,
diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
index d60e48f752..0a01f08116 100644
--- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
+++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
@@ -18,7 +18,6 @@
 from transformer_engine.pytorch.custom_recipes import utils
 from transformer_engine.common.recipe import NVFP4BlockScaling
 
-
 recipe_available, reason_for_no_recipe = te.is_nvfp4_available(return_reason=True)
 
 
@@ -84,8 +83,6 @@ def maybe_skip_row_scaled_unsupported_quantization(
     M: int | None = None,
     N: int | None = None,
 ) -> None:
-    if IS_HIP_EXTENSION and use_4over6:
-        pytest.skip("NVFP4 4over6 is not supported on ROCm")
     if use_4over6 and with_2d_quantization:
         if x_dtype != torch.bfloat16 or M is None or N is None or M % 32 != 0 or N % 32 != 0:
             pytest.skip("NVFP4 2D 4over6 exact tests require the optimized BF16 kernel path")
@@ -151,6 +148,8 @@ def check_quantization_nvfp4_versus_reference(
     )
 
     if use_4over6:
+        if IS_HIP_EXTENSION and nvfp4_4over6_err_use_fast_math:
+            pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm")
         with nvfp4_4over6_err_fast_math(nvfp4_4over6_err_use_fast_math):
             if use_cpp_allocator:
                 x_nvfp4_sut = nvfp4_quantizer(x)
@@ -363,6 +362,8 @@ def test_nvfp4_quantization_extrema_versus_reference(
     )
 
     if nvfp4_4over6_config.use_4over6:
+        if IS_HIP_EXTENSION and nvfp4_4over6_config.err_use_fast_math:
+            pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm")
         with nvfp4_4over6_err_fast_math(nvfp4_4over6_config.err_use_fast_math):
             if use_cpp_allocator:
                 x_nvfp4_sut = nvfp4_quantizer(x)
@@ -508,6 +509,8 @@ def test_nvfp4_quantization_boundary_values(
     )
 
     if nvfp4_4over6_config.use_4over6:
+        if IS_HIP_EXTENSION and nvfp4_4over6_config.err_use_fast_math:
+            pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm")
         with nvfp4_4over6_err_fast_math(nvfp4_4over6_config.err_use_fast_math):
             if use_cpp_allocator:
                 x_nvfp4_sut = nvfp4_quantizer(x)
@@ -639,6 +642,8 @@ def test_nvfp4_quantization_noncontiguous_inputs(
     )
 
     if nvfp4_4over6_config.use_4over6:
+        if IS_HIP_EXTENSION and nvfp4_4over6_config.err_use_fast_math:
+            pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm")
         with nvfp4_4over6_err_fast_math(nvfp4_4over6_config.err_use_fast_math):
             if use_cpp_allocator:
                 x_nvfp4_sut = nvfp4_quantizer(x_nc)
@@ -811,3 +816,45 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer:
 
     # Sanity: column-only path must not allocate a rowwise output.
     assert out_col_only._rowwise_data is None
+
+
+@pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
+@pytest.mark.parametrize("M, N", [(128, 128), (256, 256), (512, 1024), (320, 256)])
+@pytest.mark.parametrize(
+    "with_2d_quantization", [True, False], ids=["2d_quantization", "1d_quantization"]
+)
+def test_nvfp4_4over6_columnwise_only_matches_both_directions(
+    M: int,
+    N: int,
+    with_2d_quantization: bool,
+):
+    """Columnwise-only 4over6 must match the columnwise half of both-directions.
+
+    Without a rowwise pass the candidate errors have to be reduced by the
+    columnwise-only path itself; if that reduction is skipped, selection reads an
+    uninitialised buffer and silently falls back to map-to-6.
+    """
+    if with_2d_quantization and not IS_HIP_EXTENSION:
+        pytest.skip("CUDA routes 4over6 to a kernel that requires rowwise output for 2D")
+    device = "cuda"
+    torch.manual_seed(0)
+    torch.cuda.manual_seed(0)
+    x = torch.randn((M, N), dtype=torch.bfloat16, device=device)
+
+    def _quantize(*, rowwise: bool):
+        return NVFP4Quantizer(
+            fp4_dtype=tex.DType.kFloat4E2M1,
+            rowwise=rowwise,
+            columnwise=True,
+            with_2d_quantization=with_2d_quantization,
+            nvfp4_use_4over6=True,
+        )(x)
+
+    both = _quantize(rowwise=True)
+    col_only = _quantize(rowwise=False)
+
+    torch.testing.assert_close(col_only._columnwise_data, both._columnwise_data, atol=0, rtol=0)
+    torch.testing.assert_close(
+        col_only._columnwise_scale_inv, both._columnwise_scale_inv, atol=0, rtol=0
+    )
+    assert col_only._rowwise_data is None
diff --git a/tests/pytorch/test_fusible_ops.py b/tests/pytorch/test_fusible_ops.py
index 83d54ab095..4ca7691a4a 100644
--- a/tests/pytorch/test_fusible_ops.py
+++ b/tests/pytorch/test_fusible_ops.py
@@ -76,8 +76,7 @@
     _quantization_list.append("mxfp8")
 if nvfp4_available:
     _quantization_list.append("nvfp4")
-    if not IS_HIP_EXTENSION:  # NVFP4 4over6 is not supported on ROCm
-        _quantization_list.append("nvfp4_4over6")
+    _quantization_list.append("nvfp4_4over6")
 
 
 @pytest.fixture(autouse=True, scope="function")
diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt
index e646b72c38..5fbb9a0146 100644
--- a/transformer_engine/common/CMakeLists.txt
+++ b/transformer_engine/common/CMakeLists.txt
@@ -310,8 +310,7 @@ set(cuda_only_cuda_sources
      fused_attn/fused_attn_f16_arbitrary_seqlen.cu
      fused_attn/fused_attn_fp8.cu
      fused_attn/utils.cu
-     swizzle/swizzle_block_scaling.cu
-     util/topk.cu) #CUDA-only
+     swizzle/swizzle_block_scaling.cu) #CUDA-only
 
 list(APPEND transformer_engine_cuda_arch_specific_sources
      fused_attn/flash_attn.cu
@@ -443,6 +442,27 @@ else() #USE_ROCM
   # process source code files
   include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/hipify/hipify.cmake")
   TE_Hipify(${CMAKE_CURRENT_SOURCE_DIR})
+
+  # TE_Hipify runs at configure time via execute_process, so CMake holds no dependency on the
+  # CUDA sources it reads. Without registering them, editing a source here leaves the previously
+  # generated _hip.* in place and the build silently compiles the old revision. Generated outputs
+  # are excluded: they are rewritten during configure, so listing them would re-trigger forever.
+  set(_hipify_deps)
+  foreach(_src ${transformer_engine_SOURCES})
+    get_filename_component(_abs "${CMAKE_CURRENT_SOURCE_DIR}/${_src}" ABSOLUTE)
+    list(APPEND _hipify_deps "${_abs}")
+  endforeach()
+  file(GLOB_RECURSE _hipify_hdrs
+       "${CMAKE_CURRENT_SOURCE_DIR}/*.h"
+       "${CMAKE_CURRENT_SOURCE_DIR}/*.cuh"
+       "${CMAKE_CURRENT_SOURCE_DIR}/*.hpp")
+  foreach(_hdr ${_hipify_hdrs})
+    if(NOT _hdr MATCHES "_hip\\.(h|cuh|hpp)$")
+      list(APPEND _hipify_deps "${_hdr}")
+    endif()
+  endforeach()
+  set_property(DIRECTORY APPEND PROPERTY CMAKE_CONFIGURE_DEPENDS ${_hipify_deps})
+
   TE_GetHipifiedSources("${transformer_engine_SOURCES}" ${CMAKE_CURRENT_SOURCE_DIR} te_hip_sources)
 
   message("${message_line}")
diff --git a/transformer_engine/common/cast/dispatch/dequantize.cuh b/transformer_engine/common/cast/dispatch/dequantize.cuh
index 40ceb59d72..ca4d0e7adc 100644
--- a/transformer_engine/common/cast/dispatch/dequantize.cuh
+++ b/transformer_engine/common/cast/dispatch/dequantize.cuh
@@ -43,13 +43,13 @@ inline void dequantize_helper(const Tensor &input, Tensor *output, cudaStream_t
     case NVTE_MXFP8_1D_SCALING: {
 #ifndef __HIP_PLATFORM_AMD__
       if (is_supported_by_CC_100()) {
-#endif //#ifndef __HIP_PLATFORM_AMD__
+#endif  //#ifndef __HIP_PLATFORM_AMD__
         mxfp8::dequantize(input, output, stream);
 #ifndef __HIP_PLATFORM_AMD__
       } else {
         NVTE_ERROR("MXFP8 Dequantization is NOT supported by architectures < 10.0");
       }
-#endif //#ifndef __HIP_PLATFORM_AMD__
+#endif  //#ifndef __HIP_PLATFORM_AMD__
       break;
     }
     case NVTE_NVFP4_1D_SCALING: {
@@ -70,13 +70,13 @@ inline void group_dequantize_helper(const GroupedTensor &input, GroupedTensor *o
     case NVTE_MXFP8_1D_SCALING: {
 #ifndef __HIP_PLATFORM_AMD__
       if (is_supported_by_CC_100()) {
+#endif  //#ifndef __HIP_PLATFORM_AMD__
         mxfp8::group_dequantize(&input, output, stream);
+#ifndef __HIP_PLATFORM_AMD__
       } else {
         NVTE_ERROR("MXFP8 Grouped Dequantization is NOT supported by architectures < 10.0");
       }
-#else
-      NVTE_ERROR("MXFP8 Grouped Dequantization is not supported on ROCm.");
-#endif
+#endif  //#ifndef __HIP_PLATFORM_AMD__
       break;
     }
     default:
diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh
index 4b31f1d188..bba249dd79 100644
--- a/transformer_engine/common/cast/dispatch/quantize.cuh
+++ b/transformer_engine/common/cast/dispatch/quantize.cuh
@@ -139,6 +139,12 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
                  "Non-4over6 NVFP4 quantization requires E4M3 max 448.");
       NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding,
                  "NVFP4 4over6 quantization does not support stochastic rounding.");
+#ifdef __HIP_PLATFORM_AMD__
+      // The fast-math error path is an optimisation of candidate selection, not a
+      // capability; refuse it rather than silently scoring with the exact one.
+      NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math,
+                 "NVFP4 4over6 fast-math error mode is not supported on ROCm.");
+#endif
       if (row_scaled_nvfp4) {
         NVTE_CHECK(!quant_config_cpp.nvfp4_2d_quantization,
                    "Row-scaled NVFP4 quantization does not support 2D quantization.");
@@ -154,9 +160,9 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
           (output_tensor->has_data() ||
            (output_tensor->has_columnwise_data() && quant_config_cpp.nvfp4_2d_quantization));
 
-      // Launch NVFP4 quantize kernel. 4over6 and the optimized quantize_transpose kernels are
-      // CUDA-only (Blackwell); ROCm falls through to the portable blockwise path below, which
-      // supports row-scaled NVFP4.
+      // Launch NVFP4 quantize kernel. Upstream's dedicated 4over6 and optimized
+      // quantize_transpose kernels are CUDA-only (Blackwell); ROCm falls through to the portable
+      // blockwise path below, which implements 4over6 and row-scaled NVFP4 itself.
 #ifndef __HIP_PLATFORM_AMD__
       if (nvfp4_use_4over6) {
         if (quant_config_cpp.nvfp4_2d_quantization) {
@@ -191,6 +197,8 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
             /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization,
             /*row_scaled_nvfp4=*/row_scaled_nvfp4,
             /*noop_tensor=*/noop_tensor->data,
+            /*nvfp4_e4m3_max=*/output_tensor->nvfp4_e4m3_max,
+            /*nvfp4_4over6_mode=*/quant_config_cpp.nvfp4_4over6_mode,
             /*stream=*/stream);
 #ifndef __HIP_PLATFORM_AMD__
       }
@@ -318,6 +326,12 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
                  "Non-4over6 NVFP4 quantization requires E4M3 max 448.");
       NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding,
                  "NVFP4 4over6 quantization does not support stochastic rounding.");
+#ifdef __HIP_PLATFORM_AMD__
+      // The fast-math error path is an optimisation of candidate selection, not a
+      // capability; refuse it rather than silently scoring with the exact one.
+      NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math,
+                 "NVFP4 4over6 fast-math error mode is not supported on ROCm.");
+#endif
       NVTE_CHECK(!output_tensor->row_scaled_nvfp4,
                  "Backward NVFP4 quantization does not support row-scaled outputs.");
       // Columnwise-only is supported on the optimized path only for 2D scaling; rowwise-only and
@@ -328,9 +342,9 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
           (output_tensor->has_data() ||
            (output_tensor->has_columnwise_data() && quant_config_cpp.nvfp4_2d_quantization));
 
-      // Launch NVFP4 quantize kernel. 4over6 and the optimized quantize_transpose kernels are
-      // CUDA-only (Blackwell); ROCm falls through to the portable blockwise path below, which
-      // supports row-scaled NVFP4.
+      // Launch NVFP4 quantize kernel. Upstream's dedicated 4over6 and optimized
+      // quantize_transpose kernels are CUDA-only (Blackwell); ROCm falls through to the portable
+      // blockwise path below, which implements 4over6 and row-scaled NVFP4 itself.
 #ifndef __HIP_PLATFORM_AMD__
       if (nvfp4_use_4over6) {
         if (quant_config_cpp.nvfp4_2d_quantization) {
@@ -365,6 +379,8 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
             /*use_2d_quantization=*/quant_config_cpp.nvfp4_2d_quantization,
             /*row_scaled_nvfp4=*/false,
             /*noop_tensor=*/noop_tensor->data,
+            /*nvfp4_e4m3_max=*/output_tensor->nvfp4_e4m3_max,
+            /*nvfp4_4over6_mode=*/quant_config_cpp.nvfp4_4over6_mode,
             /*stream=*/stream);
 #ifndef __HIP_PLATFORM_AMD__
       }
diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
index fbe1985cba..62a7f61e52 100644
--- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
@@ -26,10 +26,27 @@
 #include "../../utils.cuh"
 #include "group_quantize_mxfp8.cuh"
 
+#else
+
+#include <cuda_runtime.h>
+#include <transformer_engine/transformer_engine.h>
+
+#include <vector>
+
+#include "../../common.h"
+#include "../../util/ptx.cuh"
+#include "../../utils.cuh"
+#include "./rocm_vectorized_2d.cuh"
+
+#endif  // #ifndef __HIP_PLATFORM_AMD__
+
 namespace transformer_engine {
 namespace dispatch {
 namespace mxfp8 {
 namespace group_dequantize_kernel {
+#ifdef __HIP_PLATFORM_AMD__
+#include "rocm_group_dequantize_mxfp8.cuh"
+#else
 
 constexpr int MAX_SUPPORTED_TENSOR_DESCRIPTORS = 64;
 __device__ alignas(128) CUtensorMap g_tensor_maps_input[MAX_SUPPORTED_TENSOR_DESCRIPTORS];
@@ -373,10 +390,145 @@ __global__ void __launch_bounds__(128)
   destroy_barriers<ITERATIONS>(mbar, is_master_thread);
 #endif  // #if (defined __CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
 }
+#endif  // #ifdef __HIP_PLATFORM_AMD__
 }  // namespace group_dequantize_kernel
 
 inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
                              cudaStream_t stream) {
+#ifdef __HIP_PLATFORM_AMD__
+  using namespace group_dequantize_kernel;
+
+  const bool use_rowwise_scaling = input->has_data();
+  const bool use_colwise_scaling = input->has_columnwise_data();
+  NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling,
+             "Input tensor must have either rowwise or columnwise data.");
+  NVTE_CHECK(!(use_rowwise_scaling && use_colwise_scaling),
+             "Dequantize only supports rowwise or columnwise scaling, not both simultaneously.");
+
+  NVTE_CHECK(!input->with_gemm_swizzled_scales, "Input must have scales in compact format.");
+  NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision.");
+  NVTE_CHECK(!is_fp4_dtype(output->dtype()), "Output must not be FP4.");
+  NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input must have FP8 type.");
+
+  NVTE_CHECK(input->num_tensors == output->num_tensors,
+             "Number of input and output tensors must be same.");
+
+  ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;
+  if (input->all_same_shape()) {
+    shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;
+  } else if (input->all_same_first_dim()) {
+    shape_rep = ShapeRepresentation::VARYING_LAST_DIM;
+  } else if (input->all_same_last_dim()) {
+    shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;
+  } else if (input->varying_both_dims()) {
+    shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS;
+  }
+
+  const bool is_single_tensor = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS ||
+                                 shape_rep == ShapeRepresentation::VARYING_FIRST_DIM);
+
+  const size_t first_logical_dim = input->logical_shape.data[0];
+  const size_t last_logical_dim = input->logical_shape.data[1];
+  const size_t elts_total = first_logical_dim * last_logical_dim;
+  if (elts_total == 0) return;
+
+  const size_t num_tensors = input->num_tensors;
+
+  size_t blocks_X = 0;
+  size_t blocks_Y = 1;
+  if (is_single_tensor) {
+    blocks_Y = DIVUP(first_logical_dim, CHUNK_DIM_Y);
+    blocks_X = DIVUP(last_logical_dim, CHUNK_DIM_X);
+  } else {
+    NVTE_CHECK(last_logical_dim % CHUNK_DIM_X == 0,
+               "Last dimension of a grouped tensor should be divisible by 128.");
+    if (shape_rep == ShapeRepresentation::VARYING_LAST_DIM) {
+      blocks_X = DIVUP(first_logical_dim, CHUNK_DIM_Y) * (last_logical_dim / CHUNK_DIM_X);
+    } else {
+      blocks_X = DIVUP(elts_total, CHUNK_DIM_Y * CHUNK_DIM_X);
+    }
+  }
+
+  const dim3 grid(blocks_X, blocks_Y);
+  const dim3 block(THREADS_PER_CHUNK);
+
+  const int64_t *const offsets_ptr = reinterpret_cast<const int64_t *>(input->tensor_offsets.dptr);
+  const int64_t *const first_dims_ptr = reinterpret_cast<const int64_t *>(input->first_dims.dptr);
+  const int64_t *const last_dims_ptr = reinterpret_cast<const int64_t *>(input->last_dims.dptr);
+
+  // Only the uniform-last-dim path indexes colwise scales from the group's global row.
+  if (use_colwise_scaling && last_dims_ptr == nullptr && num_tensors > 1) {
+    constexpr size_t COLWISE_SCALE_ROWS = 32;
+    std::vector<int64_t> host_first_dims(num_tensors);
+    if (first_dims_ptr != nullptr) {
+      NVTE_CHECK_CUDA(cudaMemcpyAsync(host_first_dims.data(), first_dims_ptr,
+                                      num_tensors * sizeof(int64_t), cudaMemcpyDeviceToHost,
+                                      stream));
+      NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+    } else {
+      NVTE_CHECK(first_logical_dim % num_tensors == 0, "Uniform grouped tensor has ",
+                 first_logical_dim, " rows, which is not divisible by ", num_tensors, " tensors.");
+      host_first_dims.assign(num_tensors, static_cast<int64_t>(first_logical_dim / num_tensors));
+    }
+    for (size_t t = 0; t < num_tensors; t++) {
+      NVTE_CHECK(static_cast<size_t>(host_first_dims[t]) % COLWISE_SCALE_ROWS == 0,
+                 "Columnwise grouped dequantize requires each per-tensor first dimension to be "
+                 "divisible by ",
+                 COLWISE_SCALE_ROWS, "; tensor ", t, " has ", host_first_dims[t], ".");
+    }
+  }
+
+  // Row strides must be a whole number of vectors; VECTOR_WIDTH is 16 elements.
+  constexpr size_t VECTOR_ELEMS = 16;
+  if (last_dims_ptr == nullptr) {
+    NVTE_CHECK(last_logical_dim % VECTOR_ELEMS == 0,
+               "Grouped dequantize requires the last dimension to be divisible by ", VECTOR_ELEMS,
+               "; got ", last_logical_dim, ".");
+  } else {
+    std::vector<int64_t> host_last_dims(num_tensors);
+    NVTE_CHECK_CUDA(cudaMemcpyAsync(host_last_dims.data(), last_dims_ptr,
+                                    num_tensors * sizeof(int64_t), cudaMemcpyDeviceToHost, stream));
+    NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+    for (size_t t = 0; t < num_tensors; t++) {
+      NVTE_CHECK(static_cast<size_t>(host_last_dims[t]) % VECTOR_ELEMS == 0,
+                 "Grouped dequantize requires each per-tensor last dimension to be divisible by ",
+                 VECTOR_ELEMS, "; tensor ", t, " has ", host_last_dims[t], ".");
+    }
+  }
+
+  const e8m0_t *const scales_ptr =
+      use_rowwise_scaling ? reinterpret_cast<e8m0_t *>(input->scale_inv.dptr)
+                          : reinterpret_cast<e8m0_t *>(input->columnwise_scale_inv.dptr);
+
+  const SimpleTensor &input_data = use_rowwise_scaling ? input->data : input->columnwise_data;
+
+  const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1;
+  const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1;
+
+  TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH(
+      scale_dim_Y_colwise, SCALE_DIM_Y,
+      TRANSFORMER_ENGINE_MX_SCALE_DIM_SWITCH(
+          scale_dim_X_rowwise, SCALE_DIM_X,
+          TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(
+              input->dtype(), IType,
+              TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
+                  output->dtype(), OType,
+                  // Only the single-tensor layout has one row stride, so only it can be
+                  // proven vector-aligned on the host; groups take the general path.
+                  TRANSFORMER_ENGINE_SWITCH_CONDITION(
+                      is_single_tensor && !(last_logical_dim % (32 * sizeof(OType))), IS_ALIGNED,
+                      grouped_dequantize_mxfp8_kernel<IType, OType, SCALE_DIM_Y, SCALE_DIM_X,
+                                                      IS_ALIGNED><<<grid, block, 0, stream>>>(
+                          reinterpret_cast<const IType *>(input_data.dptr),
+                          reinterpret_cast<OType *>(output->data.dptr), scales_ptr,
+                          first_logical_dim, last_logical_dim, num_tensors, offsets_ptr,
+                          first_dims_ptr, last_dims_ptr););  // NOLINT(*)
+              );                                             // NOLINT(*)
+          );                                                 // NOLINT(*)
+      );                                                     // NOLINT(*)
+  );                                                         // NOLINT(*)
+  NVTE_CHECK_CUDA(cudaGetLastError());
+#else
   using namespace group_dequantize_kernel;
 
   checkCuDriverContext(stream);
@@ -490,12 +642,11 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
           });  // NOLINT(*)
   );           // NOLINT(*)
   NVTE_CHECK_CUDA(cudaGetLastError());
+#endif  // #ifdef __HIP_PLATFORM_AMD__
 }
 
 }  // namespace mxfp8
 }  // namespace dispatch
 }  // namespace transformer_engine
 
-#endif  // __HIP_PLATFORM_AMD__
-
 #endif  // TRANSFORMER_ENGINE_GROUP_DEQUANTIZE_MXFP8_CUH_
diff --git a/transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh
new file mode 100644
index 0000000000..5392cb21a9
--- /dev/null
+++ b/transformer_engine/common/cast/mxfp8/rocm_group_dequantize_mxfp8.cuh
@@ -0,0 +1,212 @@
+/*************************************************************************
+ * Copyright (c) 2026, Advanced Micro Devices, Inc. All rights reserved.
+ *
+ * License for AMD contributions = MIT. See LICENSE for more information
+ ************************************************************************/
+
+#pragma once
+// drop-in rocm replacement for mxfp8 grouped dequantize kernel
+//#include "hip/hip_runtime.h" //dummy include to prevent hipification adding this header
+
+constexpr size_t CHUNK_DIM_Y = 128;
+constexpr size_t CHUNK_DIM_X = 128;
+constexpr size_t THREADS_PER_CHUNK = 128;
+
+constexpr size_t ELEMS_PER_THREAD = 16;
+constexpr size_t BUFFER_DIM_Y = 16;
+constexpr size_t BUFFER_DIM_X = CHUNK_DIM_X;  // 128
+constexpr size_t SHMEM_DIM_Y = BUFFER_DIM_Y;  // 16
+constexpr size_t SHMEM_DIM_X = BUFFER_DIM_X;  // 128
+
+constexpr size_t THREADS_PER_CHUNK_X_ROWWISE = CHUNK_DIM_X / ELEMS_PER_THREAD;  //  8 = 128 / 16
+constexpr size_t THREADS_PER_CHUNK_X_COLWISE = CHUNK_DIM_X;                     //  128
+constexpr size_t ITERATIONS = CHUNK_DIM_Y / BUFFER_DIM_Y;                       //    8 = 128 / 16
+static_assert(ITERATIONS >= 1);
+
+template <typename IType, typename OType, size_t SCALE_DIM_Y, size_t SCALE_DIM_X, bool IS_ALIGNED>
+__global__ void __launch_bounds__(THREADS_PER_CHUNK)
+    grouped_dequantize_mxfp8_kernel(const IType *input_ptr, OType *output_ptr,
+                                    const e8m0_t *scales_ptr, const size_t first_logical_dim,
+                                    const size_t last_logical_dim, const size_t num_tensors,
+                                    const int64_t *const offsets_ptr,
+                                    const int64_t *const first_dims_ptr,
+                                    const int64_t *const last_dims_ptr) {
+  constexpr bool USE_ROWWISE_SCALING = SCALE_DIM_X > 1;
+
+  constexpr size_t SCALES_ROWWISE_PER_CHUNK_Y = CHUNK_DIM_Y;                //  128
+  constexpr size_t SCALES_ROWWISE_PER_CHUNK_X = CHUNK_DIM_X / SCALE_DIM_X;  //    4 = 128 / 32
+
+  constexpr size_t SCALES_COLWISE_PER_CHUNK_Y = CHUNK_DIM_Y / SCALE_DIM_Y;  //    4 = 128 / 32
+  constexpr size_t SCALES_COLWISE_PER_CHUNK_X = CHUNK_DIM_X;                //  128
+
+  constexpr size_t THREADS_PER_SCALE_X_ROWWISE =
+      DIVUP(SCALE_DIM_X, ELEMS_PER_THREAD);  // 2 = 32 / 16
+  constexpr size_t VECTOR_WIDTH = IS_ALIGNED ? 8 : 16;
+
+  const int tid_rowwise_Y = threadIdx.x / THREADS_PER_CHUNK_X_ROWWISE;
+  const int tid_rowwise_X = threadIdx.x % THREADS_PER_CHUNK_X_ROWWISE;
+  const int tid_colwise_X = threadIdx.x % THREADS_PER_CHUNK_X_COLWISE;
+
+  const int thread_offset_Y = tid_rowwise_Y;
+  const int thread_offset_X_rowwise = tid_rowwise_X * ELEMS_PER_THREAD;
+
+  alignas(128) __shared__ IType in_sh[SHMEM_DIM_Y][SHMEM_DIM_X];
+  alignas(128) __shared__ OType out_sh[SHMEM_DIM_Y][SHMEM_DIM_X];
+
+  // Mirrors the launcher's is_single_tensor: null last_dims means a 2D grid, not a tile stride.
+  const bool uniform_last_dim = (last_dims_ptr == nullptr);
+  size_t total_tiles = 1;
+  size_t tiles_stride = 1;
+  if (!uniform_last_dim) {
+    total_tiles = 0;
+    for (size_t t = 0; t < num_tensors; t++) {
+      const size_t t_rows =
+          (first_dims_ptr != nullptr) ? static_cast<size_t>(first_dims_ptr[t]) : first_logical_dim;
+      const size_t t_cols = static_cast<size_t>(last_dims_ptr[t]);
+      total_tiles += DIVUP(t_rows, CHUNK_DIM_Y) * DIVUP(t_cols, CHUNK_DIM_X);
+    }
+    tiles_stride = gridDim.x;
+  }
+
+  for (size_t block_tile = uniform_last_dim ? 0 : blockIdx.x; block_tile < total_tiles;
+       block_tile += tiles_stride) {
+    size_t rows = 0;
+    size_t cols = 0;
+    int block_id_Y = 0;
+    int block_id_X = 0;
+    size_t tensor_base_elts = 0;
+    size_t scales_base_offset = 0;
+
+    if (uniform_last_dim) {
+      // Uniform last dim: the group is one contiguous 2D tensor and, because ROCm scale
+      // tensors are unpadded, the per-tensor scale blocks stack into one contiguous array.
+      rows = first_logical_dim;
+      cols = last_logical_dim;
+      block_id_Y = blockIdx.y;
+      block_id_X = blockIdx.x;
+    } else {
+      // Find the tensor owning this tile.
+      size_t tiles_before = 0;
+      size_t tensor_id = 0;
+      for (; tensor_id < num_tensors; tensor_id++) {
+        const size_t t_rows = (first_dims_ptr != nullptr)
+                                  ? static_cast<size_t>(first_dims_ptr[tensor_id])
+                                  : first_logical_dim;
+        const size_t t_cols = static_cast<size_t>(last_dims_ptr[tensor_id]);
+        const size_t t_tiles = DIVUP(t_rows, CHUNK_DIM_Y) * DIVUP(t_cols, CHUNK_DIM_X);
+        if (block_tile < tiles_before + t_tiles) {
+          const size_t local_tile = block_tile - tiles_before;
+          const size_t tiles_x = DIVUP(t_cols, CHUNK_DIM_X);
+          rows = t_rows;
+          cols = t_cols;
+          block_id_Y = static_cast<int>(local_tile / tiles_x);
+          block_id_X = static_cast<int>(local_tile % tiles_x);
+          break;
+        }
+        tiles_before += t_tiles;
+        tensor_base_elts += t_rows * t_cols;
+        scales_base_offset += USE_ROWWISE_SCALING ? t_rows * DIVUP(t_cols, SCALE_DIM_X)
+                                                  : DIVUP(t_rows, SCALE_DIM_Y) * t_cols;
+      }
+      if (tensor_id >= num_tensors) continue;
+      if (offsets_ptr != nullptr) tensor_base_elts = static_cast<size_t>(offsets_ptr[tensor_id]);
+    }
+
+    if (rows == 0 || cols == 0) continue;
+
+    // Per-tile views of the group buffers; the kernel arguments stay untouched so that the
+    // bases do not accumulate across tiles.
+    const IType *const tile_input_ptr = input_ptr + tensor_base_elts;
+    OType *const tile_output_ptr = output_ptr + tensor_base_elts;
+    const e8m0_t *const tile_scales_ptr = scales_ptr + scales_base_offset;
+
+    const size_t scales_stride = USE_ROWWISE_SCALING ? DIVUP(cols, SCALE_DIM_X) : cols;
+
+    const int chunk_offset_Y = block_id_Y * CHUNK_DIM_Y;
+    const int chunk_offset_X = block_id_X * CHUNK_DIM_X;
+
+    const int scales_rowwise_chunk_offset_Y = block_id_Y * SCALES_ROWWISE_PER_CHUNK_Y;
+    const int scales_rowwise_chunk_offset_X = block_id_X * SCALES_ROWWISE_PER_CHUNK_X;
+    const int scales_colwise_chunk_offset_Y = block_id_Y * SCALES_COLWISE_PER_CHUNK_Y;
+    const int scales_colwise_chunk_offset_X = block_id_X * SCALES_COLWISE_PER_CHUNK_X;
+
+    for (int iter = 0; iter < ITERATIONS; iter++) {
+      const int chunk_it_offset_y = chunk_offset_Y + iter * BUFFER_DIM_Y;
+      const int chunk_it_offset_x = chunk_offset_X;
+
+      copy_2d_to_shared<IType, VECTOR_WIDTH, IS_ALIGNED>(&in_sh[0][0], tile_input_ptr,
+                                                         chunk_it_offset_x, chunk_it_offset_y, cols,
+                                                         SHMEM_DIM_Y, SHMEM_DIM_X, rows, cols);
+      __syncthreads();
+
+      const int scale_offset_Y =
+          USE_ROWWISE_SCALING
+              ? (scales_rowwise_chunk_offset_Y + iter * BUFFER_DIM_Y + tid_rowwise_Y)
+              : (scales_colwise_chunk_offset_Y + (iter * BUFFER_DIM_Y) / SCALE_DIM_Y);
+
+      const int scale_offset_X =
+          USE_ROWWISE_SCALING
+              ? (scales_rowwise_chunk_offset_X + tid_rowwise_X / THREADS_PER_SCALE_X_ROWWISE)
+              : (scales_colwise_chunk_offset_X + tid_colwise_X);
+
+      const size_t scales_rows = USE_ROWWISE_SCALING ? rows : DIVUP(rows, SCALE_DIM_Y);
+      const size_t scales_cols = USE_ROWWISE_SCALING ? DIVUP(cols, SCALE_DIM_X) : cols;
+
+      e8m0_t biased_exponent = static_cast<e8m0_t>(127);
+      if (static_cast<size_t>(scale_offset_Y) < scales_rows &&
+          static_cast<size_t>(scale_offset_X) < scales_cols) {
+        const size_t scale_idx = scale_offset_Y * scales_stride + scale_offset_X;
+        biased_exponent = tile_scales_ptr[scale_idx];
+      }
+      const float block_scale = ptx::exp2f(biased_exponent);
+
+      if constexpr (USE_ROWWISE_SCALING) {
+        Vec<IType, ELEMS_PER_THREAD> in;
+        Vec<OType, ELEMS_PER_THREAD> out;
+
+        const int shmem_offset_y = thread_offset_Y;
+        const int shmem_offset_x = thread_offset_X_rowwise;
+        in.load_from(&in_sh[shmem_offset_y][shmem_offset_x]);
+
+#pragma unroll
+        for (int j = 0; j < ELEMS_PER_THREAD; j++) {
+#if defined(__gfx1250__)
+          // FIXME: Force E4M3 OCP interpretation because HIP headers do not declare
+          // which type gfx1250 supports. This can be removed once HIP headers are updated.
+          const float elt =
+              std::is_same_v<IType, fp8e4m3>
+                  ? static_cast<float>(*reinterpret_cast<__hip_fp8_e4m3 *>(&in.data.elt[j]))
+                  : static_cast<float>(in.data.elt[j]);
+          out.data.elt[j] = static_cast<OType>(block_scale * elt);
+#else
+          out.data.elt[j] = static_cast<OType>(block_scale * static_cast<float>(in.data.elt[j]));
+#endif
+        }
+        out.store_to(&out_sh[shmem_offset_y][shmem_offset_x]);
+      } else {
+#pragma unroll
+        for (int i = 0; i < BUFFER_DIM_Y; i++) {
+#if defined(__gfx1250__)
+          // FIXME: Force E4M3 OCP interpretation because HIP headers do not declare
+          // which type gfx1250 supports. This can be removed once HIP headers are updated.
+          const float elt = std::is_same_v<IType, fp8e4m3>
+                                ? static_cast<float>(
+                                      *reinterpret_cast<__hip_fp8_e4m3 *>(&in_sh[i][tid_colwise_X]))
+                                : static_cast<float>(in_sh[i][tid_colwise_X]);
+#else
+          const float elt = static_cast<float>(in_sh[i][tid_colwise_X]);
+#endif
+          out_sh[i][tid_colwise_X] = static_cast<OType>(block_scale * elt);
+        }
+      }
+
+      __syncthreads();
+
+      bulk_tensor_2d_shared_to_global<OType, VECTOR_WIDTH, IS_ALIGNED>(
+          &out_sh[0][0], tile_output_ptr, chunk_it_offset_x, chunk_it_offset_y, cols, SHMEM_DIM_Y,
+          SHMEM_DIM_X, rows, cols);
+
+      __syncthreads();
+    }
+  }
+}
diff --git a/transformer_engine/common/fused_attn/flash_attn.cu b/transformer_engine/common/fused_attn/flash_attn.cu
index 14704d8cd9..296c4b0add 100644
--- a/transformer_engine/common/fused_attn/flash_attn.cu
+++ b/transformer_engine/common/fused_attn/flash_attn.cu
@@ -161,7 +161,6 @@ void prepare_flash_attn_bwd(Tensor q, Tensor k, Tensor v, Tensor qkv, cudaStream
 // multi_tensor_transpose_to_bhsd: BSHD/SBHD -> BHSD
 // ============================================================================
 
-#ifndef __HIP_PLATFORM_AMD__  // Disabled on ROCm
 namespace multi_tensor_transpose_to_bhsd {
 
 using flash_attention::Vec;
@@ -178,9 +177,11 @@ struct PermuteParams {
   PermuteSlot slots[kMaxPermuteTensors];
 };
 
+#ifndef __HIP_PLATFORM_AMD__  // CUtensorMap has no ROCm counterpart
 struct TmaMapParams {
   CUtensorMap maps[kMaxPermuteTensors];
 };
+#endif
 
 // ---------- path 3: fallback_not_vec_aligned ----------
 
@@ -369,6 +370,8 @@ __launch_bounds__(fallback_permute_threads) __global__
 
 // ---------- path 1: TMA ----------
 
+#ifndef __HIP_PLATFORM_AMD__  // TMA: CUtensorMap, cp.async.bulk, mbarrier, SM 10.0+
+
 constexpr int tma_permute_threads = 128;
 constexpr int tma_permute_s_tile_default = 32;
 
@@ -589,6 +592,8 @@ static void create_strided_tensor_map(CUtensorMap &map, void *ptr, DType dtype,
   }
 }
 
+#endif  // !__HIP_PLATFORM_AMD__ (path 1: TMA)
+
 void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_tensors,
                                     NVTE_QKV_Format original_format, cudaStream_t stream) {
   using namespace transformer_engine;
@@ -605,7 +610,15 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_
   size_t s_max = 0, h_max = 0, s_min = SIZE_MAX;
   size_t d_in_max = 0, d_out_max = 0;
   bool any_not_vec_aligned = false;
+#ifdef __HIP_PLATFORM_AMD__
+  // Both readers below are compiled out on ROCm, so this only keeps the variable
+  // defined and states the intent: had they been reachable, the sm >= 100 test
+  // would not have excluded us -- sm_arch() returns the raw gfx number, so
+  // gfx1250 reports 125 and satisfies it.
+  bool all_tma_ok = false;
+#else
   bool all_tma_ok = true;
+#endif
 
   for (size_t i = 0; i < num_tensors; ++i) {
     const size_t H = outputs[i].shape()[1];
@@ -623,6 +636,7 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_
     if (inner < 32 || inner % 16 != 0) all_tma_ok = false;
   }
 
+#ifndef __HIP_PLATFORM_AMD__
   if (all_tma_ok) {
     const int sm = cuda::sm_arch(cuda::current_device());
     if (sm < 100) {
@@ -641,12 +655,14 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_
       }
     }
   }
+#endif  // !__HIP_PLATFORM_AMD__
 
   // Dispatch order:
   //  1. TMA path: SM 10.0+, D_in*elem >= 32 && 16-aligned, supported dtype,
   //     and s_tile*D_in*elem is uint4-aligned.
   //  2. Fallback path (vec-aligned): vectorized loads/stores when D_in*elem % 4 == 0.
   //  3. Fallback path (not-vec-aligned): shared-memory transpose when D_in*elem % 4 != 0.
+#ifndef __HIP_PLATFORM_AMD__
   if (all_tma_ok) {
     const size_t s_tile = std::min(static_cast<size_t>(tma_permute_s_tile_default), s_min);
     bool tma_aligned = true;
@@ -686,6 +702,7 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_
       return;
     }
   }
+#endif  // !__HIP_PLATFORM_AMD__ (TMA dispatch)
 
   if (!any_not_vec_aligned) {
     const unsigned int permute_s_splits = std::max(
@@ -730,7 +747,6 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_
 }
 
 }  // namespace multi_tensor_transpose_to_bhsd
-#endif
 
 // ===================================================================================
 // multi_tensor_pad_last_dim: pad the last dim of multiple tensors to certain alignment
@@ -872,7 +888,6 @@ void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs
                                          size_t num_tensors, NVTE_QKV_Format original_format,
                                          cudaStream_t stream) {
   NVTE_API_CALL(nvte_multi_tensor_transpose_to_bhsd);
-#ifndef __HIP_PLATFORM_AMD__
   NVTE_CHECK(original_format == NVTE_QKV_Format::NVTE_BSHD ||
                  original_format == NVTE_QKV_Format::NVTE_SBHD,
              "nvte_multi_tensor_transpose_to_bhsd: only BSHD/SBHD -> BHSD is currently "
@@ -890,14 +905,6 @@ void nvte_multi_tensor_transpose_to_bhsd(NVTETensor *inputs, NVTETensor *outputs
     multi_tensor_transpose_to_bhsd::multi_tensor_transpose_to_bhsd(
         in_vec.data() + offset, out_vec.data() + offset, batch, original_format, stream);
   }
-#else
-  (void)inputs;
-  (void)outputs;
-  (void)num_tensors;
-  (void)original_format;
-  (void)stream;
-  NVTE_ERROR("nvte_multi_tensor_transpose_to_bhsd is not supported on ROCm.");
-#endif
 }
 
 void nvte_multi_tensor_pad_last_dim(NVTETensor *inputs, NVTETensor *outputs, size_t num_tensors,
diff --git a/transformer_engine/common/gemm/kittens/CMakeLists.txt b/transformer_engine/common/gemm/kittens/CMakeLists.txt
index f09a7c7ca5..6fee3125a8 100644
--- a/transformer_engine/common/gemm/kittens/CMakeLists.txt
+++ b/transformer_engine/common/gemm/kittens/CMakeLists.txt
@@ -57,6 +57,9 @@ else()
         add_library(kittens_gemm_cdna4 OBJECT cdna4/blockwise_fp8_gemm.cpp cdna4/mxfp8_gemm.cpp)
         set_source_files_properties(cdna4/mxfp8_gemm.cpp
             PROPERTIES LANGUAGE HIP COMPILE_FLAGS "-gline-tables-only")
+        # -ffast-math folds fpext(fptrunc(x)) back to x, so any deliberate narrow-then-widen
+        # in this file has to be made opaque to the optimizer -- see round_to_out_dtype() in
+        # cdna4/blockwise_fp8_gemm_helper.cuh.
         set_source_files_properties(cdna4/blockwise_fp8_gemm.cpp
             PROPERTIES LANGUAGE HIP COMPILE_FLAGS "-gline-tables-only -ffast-math")
         set_target_properties(kittens_gemm_cdna4 PROPERTIES
diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
index 413b126c89..0ae6c4eb2b 100644
--- a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
+++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
@@ -94,20 +94,27 @@ __device__ inline float read_elem(const void *p, int dtype, int idx) {
 }
 
 template <typename OType>
-__device__ inline float rtne_cast_roundtrip(float v) {
-    if constexpr (std::is_same_v<OType, float>) {
-        return v;
+__device__ inline OType convert_out(float v) {
+    if constexpr (std::is_same_v<OType, kittens::bf16>) {
+        return kittens::base_types::convertor<OType, float>::convert(rtne_bias(v));
     } else {
-        return static_cast<float>(kittens::base_types::convertor<OType, float>::convert(rtne_bias(v)));
+        return kittens::base_types::convertor<OType, float>::convert(v);
     }
 }
 
+// cuBLAS converts the GEMM result to the output type before accumulating beta*C, so the
+// blockwise reference does the same (see qgemm() in
+// tests/pytorch/references/blockwise_fp8_gemm_reference.py). Keeping the accumulator in
+// fp32 across the beta add is more accurate but does not reproduce that result, so the
+// round is deliberate rather than an oversight. It has to go through convert_out so that
+// the rtne_bias correction is applied exactly where the store applies it -- fp16 must not
+// pick up a bf16-granularity bias.
 template <typename OType>
-__device__ inline OType convert_out(float v) {
-    if constexpr (std::is_same_v<OType, kittens::bf16>) {
-        return kittens::base_types::convertor<OType, float>::convert(rtne_bias(v));
+__device__ inline float round_to_out_dtype(float v) {
+    if constexpr (std::is_same_v<OType, float>) {
+        return v;
     } else {
-        return kittens::base_types::convertor<OType, float>::convert(v);
+        return kittens::base_types::convertor<float, OType>::convert(convert_out<OType>(v));
     }
 }
 
@@ -184,7 +191,7 @@ __device__ inline void apply_epilogue(
                 float x = v[r];
                 if constexpr (HAS_BIAS) x += bias_v;
                 if constexpr (HAS_BETA) {
-                    x = rtne_cast_roundtrip<OType>(x);
+                    x = round_to_out_dtype<OType>(x);
                     x += beta * static_cast<float>(c_in[m_g * N + col]);
                 }
                 if constexpr (HAS_GELU) {
diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
index 7207a97fd7..f173dcc49c 100644
--- a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
+++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
@@ -83,14 +83,28 @@ __device__ inline float read_elem(const void *p, int dtype, int idx) {
     return reinterpret_cast<const float *>(p)[idx];
 }
 
+// cuBLAS converts the GEMM result to the output type before accumulating beta*C, so the
+// blockwise reference does the same (see qgemm() in
+// tests/pytorch/references/blockwise_fp8_gemm_reference.py). Keeping the accumulator in
+// fp32 across the beta add is more accurate but does not reproduce that result, so the
+// round is deliberate rather than an oversight.
+//
+// This translation unit is built with -ffast-math (see gemm/kittens/CMakeLists.txt), which
+// folds fpext(fptrunc(x)) straight back to x -- written plainly, the round-trip below emits
+// no conversion at all and the whole epilogue silently reverts to an fp32 accumulate. The
+// empty asm makes the narrowed value opaque so the conversion has to happen.
 template <typename OType>
 __device__ inline float round_to_out_dtype(float v) {
     if constexpr (std::is_same_v<OType, float>) {
         return v;
     } else if constexpr (std::is_same_v<OType, kittens::bf16>) {
-        return __bfloat162float(__float2bfloat16(v));
+        uint16_t bits = __builtin_bit_cast(uint16_t, __float2bfloat16(v));
+        asm volatile("" : "+v"(bits));
+        return __bfloat162float(__builtin_bit_cast(kittens::bf16, bits));
     } else {
-        return __half2float(__float2half(v));
+        uint16_t bits = __builtin_bit_cast(uint16_t, __float2half(v));
+        asm volatile("" : "+v"(bits));
+        return __half2float(__builtin_bit_cast(__half, bits));
     }
 }
 
diff --git a/transformer_engine/common/gemm/rocm_gemm.cu b/transformer_engine/common/gemm/rocm_gemm.cu
index 082c96ee7a..e200a41bfe 100644
--- a/transformer_engine/common/gemm/rocm_gemm.cu
+++ b/transformer_engine/common/gemm/rocm_gemm.cu
@@ -535,8 +535,10 @@ static void dequant_fp4_gemm_inputs(
     const void** alpha_ptr_out, hipStream_t stream) {
 
   const float fp4_max = 6.0f;
-  const float fp8_max = te_fp8_fnuz() ? 240.0f : 448.0f;
-  const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max * fp8_max);
+  // Read the bound per operand as upstream does: 4over6 tensors carry 256, not 448.
+  const float fp8_max_A = te_fp8_fnuz() ? 240.0f : static_cast<float>(inputA.nvfp4_e4m3_max);
+  const float fp8_max_B = te_fp8_fnuz() ? 240.0f : static_cast<float>(inputB.nvfp4_e4m3_max);
+  const float factor_inv = 1.0f / (fp4_max * fp4_max * fp8_max_A * fp8_max_B);
 
   const float* amax_A = (transa == CUBLAS_OP_T)
       ? reinterpret_cast<const float*>(inputA.amax.dptr)
diff --git a/transformer_engine/common/transpose/cast_transpose.h b/transformer_engine/common/transpose/cast_transpose.h
index c462b30147..bbc5e6cd8b 100644
--- a/transformer_engine/common/transpose/cast_transpose.h
+++ b/transformer_engine/common/transpose/cast_transpose.h
@@ -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.
@@ -68,7 +70,8 @@ void quantize_transpose_vector_blockwise_fp4(
     const bool return_identity, const bool return_transpose, const bool pow2_scale,
     const bool swizzled_scale, const bool use_stochastic_rounding,
     const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4,
-    const SimpleTensor &noop_tensor, cudaStream_t stream);
+    const SimpleTensor &noop_tensor, const int nvfp4_e4m3_max,
+    const NVTENVFP44Over6Mode nvfp4_4over6_mode, cudaStream_t stream);
 
 }  // namespace transformer_engine::detail
 
diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
index e6481e09a2..a8dadf228e 100644
--- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
+++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
@@ -191,6 +191,33 @@ constexpr int smem_size_for_tile(int tile_dim) {
 static __device__ constexpr unsigned int WARP_REDUCE_AMAX_GROUP_MASKS[8] = {
     0x01010101, 0x02020202, 0x04040404, 0x08080808, 0x10101010, 0x20202020, 0x40404040, 0x80808080};
 
+#ifdef __HIP_PLATFORM_AMD__
+// Sum the 16 row errors of a 2D block in upstream's tree order: pair rows 8
+// apart, then 4, then 2, then 1 (quantize_4over6_nvfp4.cuh::reduce_group_sum_16).
+// fp32 addition is not associative and 4over6 selection compares two sums that
+// are frequently equal to within an ulp, so the association is part of the
+// contract rather than an implementation detail. Upstream can run this tree in
+// warp shuffles because one warp there owns 16 consecutive rows; here a warp
+// owns only 4, so the rows paired first sit in different warps and the whole
+// set has to be staged in shared memory and reduced by one thread.
+template <int kBlockSize, typename CType>
+__device__ __forceinline__ float reduceBlockError(const CType* staged) {
+  float v[kBlockSize];
+#pragma unroll
+  for (int i = 0; i < kBlockSize; ++i) {
+    v[i] = static_cast<float>(staged[i]);
+  }
+#pragma unroll
+  for (int stride = kBlockSize / 2; stride > 0; stride /= 2) {
+#pragma unroll
+    for (int i = 0; i < stride; ++i) {
+      v[i] = __fadd_rn(v[i], v[i + stride]);
+    }
+  }
+  return v[0];
+}
+#endif
+
 // max for every group_size elements in warp
 template <int group_size, int shfl_down_stride>
 __device__ __forceinline__ float groupMax(float val, unsigned int groupMask) {
@@ -220,18 +247,100 @@ __device__ __forceinline__ float ComputeEncodeScaleFP4(ScaleType decode_scale,
                TypeExtrema<float>::max);
 }
 
+#ifdef __HIP_PLATFORM_AMD__
+// Round-trip one value through FP4 at a given scale and return its contribution
+// to the block error.
+// E2M1 holds 16 values, so decode by table rather than through the packed-pair
+// conversion intrinsic, whose element ordering is easy to get wrong here.
+__device__ __forceinline__ float fp4_decode_e2m1(const __hip_fp4_storage_t q) {
+  constexpr float kMagnitude[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f};
+  const float m = kMagnitude[q & 0x7];
+  return (q & 0x8) ? -m : m;
+}
+
+// Rounding is pinned with the _rn intrinsics and the operand order upstream uses
+// (quantize_4over6_nvfp4.cuh:199). Selection turns on comparing two sums that are
+// often close, so an FMA contraction here flips near-ties and the result stops
+// being bit-comparable with the reference.
+__device__ __forceinline__ float fp4_roundtrip_err(const float x, const float block_scale_inverse,
+                                                   const float sf, const float global_amax,
+                                                   const float err_denom, const bool use_mse) {
+  const __hip_fp4_storage_t q =
+      __hip_cvt_float_to_fp4(__fmul_rn(x, block_scale_inverse), __HIP_E2M1, hipRoundNearest);
+  const float dequant = fp4_decode_e2m1(q);
+  const float val = __fdiv_rn(__fmul_rn(__fmul_rn(dequant, sf), global_amax), err_denom);
+  const float diff = __fsub_rn(val, x);
+  return use_mse ? __fmul_rn(diff, diff) : fabsf(diff);
+}
+
+// 4over6: a block scale expanded by 1.5x maps the block max onto FP4 code 4
+// rather than code 6, spacing the codes more finely for blocks whose values sit
+// away from the top of the range. Score both encodings of the block and keep the
+// one that reconstructs it with less error.
+//
+// Takes the block as a plain array: the rowwise and columnwise call sites walk
+// shared memory differently, so a shared loop cannot serve both.
+//
+// Split in two so a caller that has already reduced the block error can rebuild the
+// two scales without walking the block again.
+template <typename ScaleType>
+__device__ __forceinline__ void Compute4Over6Scales(const float amax,
+                                                    const float global_decode_scale,
+                                                    const float global_encode_scale,
+                                                    ScaleType* sf6_out, ScaleType* sf4_out,
+                                                    float* enc6_out, float* enc4_out) {
+  // Both candidates come off upstream's base expression, amax/6*S expanded by 1.5 for
+  // map-to-4 (quantize_4over6_nvfp4.cuh:129). Folding the /6 into the multiplier the way
+  // the non-4over6 path does rounds differently and moves blocks across E4M3 codes.
+  const float base = __fmul_rn(__fdiv_rn(amax, 6.0f), global_encode_scale);
+  const ScaleType sf6 = static_cast<ScaleType>(fminf(base, TypeExtrema<float>::max));
+  const ScaleType sf4 =
+      static_cast<ScaleType>(fminf(__fmul_rn(base, 1.5f), TypeExtrema<float>::max));
+  *sf6_out = sf6;
+  *sf4_out = sf4;
+  *enc6_out = ComputeEncodeScaleFP4<ScaleType>(sf6, global_decode_scale);
+  *enc4_out = ComputeEncodeScaleFP4<ScaleType>(sf4, global_decode_scale);
+}
+
+template <typename ScaleType>
+__device__ __forceinline__ void Compute4Over6Candidates(
+    const float* x, const int n, const float amax, const float global_decode_scale,
+    const bool use_mse, const float global_encode_scale, const float true_global_amax,
+    const float err_denom, ScaleType* sf6_out, ScaleType* sf4_out, float* enc6_out, float* enc4_out,
+    float* err6_out, float* err4_out) {
+  Compute4Over6Scales<ScaleType>(amax, global_decode_scale, global_encode_scale, sf6_out, sf4_out,
+                                 enc6_out, enc4_out);
+  const ScaleType sf6 = *sf6_out;
+  const ScaleType sf4 = *sf4_out;
+  const float enc6 = *enc6_out;
+  const float enc4 = *enc4_out;
+  // Dequantisation is val = dequant * sf * global_amax / (fp4_max * e4m3_max). Take the
+  // amax the encode scale was derived from; dividing err_denom back out costs a bit, and
+  // selection compares sums that are often equal to within one.
+  const float global_amax = true_global_amax;
+
+  float err6 = 0.0f;
+  float err4 = 0.0f;
+  for (int i = 0; i < n; ++i) {
+    err6 = __fadd_rn(err6, fp4_roundtrip_err(x[i], enc6, static_cast<float>(sf6), global_amax,
+                                             err_denom, use_mse));
+    err4 = __fadd_rn(err4, fp4_roundtrip_err(x[i], enc4, static_cast<float>(sf4), global_amax,
+                                             err_denom, use_mse));
+  }
+  *err6_out = err6;
+  *err4_out = err4;
+}
+
+#endif  // __HIP_PLATFORM_AMD__
+
 template <typename IType, typename ScaleType>
 __device__ __forceinline__ float ComputeOutputFP4(IType input, float encode_scale) {
   return static_cast<float>(input) * encode_scale;
 }
 
-__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax) {
-#if defined(__HIP_PLATFORM_AMD__) && !defined(__HIP_DEVICE_COMPILE__)
-  // On AMD host, TypeExtrema<fp8e4m3>::max is non-constexpr (runtime FNUZ detection)
-  const float fp8_max = TypeExtrema<fp8e4m3>::max;
-#else
-  constexpr float fp8_max = TypeExtrema<fp8e4m3>::max;
-#endif
+__device__ __forceinline__ float ComputeGlobalEncodeScaleFP4(const float global_amax,
+                                                             const float e4m3_max) {
+  const float fp8_max = e4m3_max;
   constexpr float fp4_max = TypeExtrema<fp4e2m1>::max;
   float global_encode_scale = fp8_max * fp4_max / global_amax;
   // If scale is infinity, return max value of float32
@@ -433,14 +542,15 @@ __device__ __forceinline__ __nv_fp4x4_e2m1 cvt_fp32_to_fp4_4x(const float2 in01,
 
 template <bool kReturnIdentity, bool kReturnTranspose, bool kIsE8Scaling, bool kAligned,
           typename CType, typename IType, typename OType, typename ScaleType, bool kSwizzledScale,
-          bool kApplyStochasticRounding, bool kIs2DBlockScaling, bool kRowScaledNVFP4>
+          bool kApplyStochasticRounding, bool kIs2DBlockScaling, bool kRowScaledNVFP4,
+          bool kUse4Over6>
 __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpose_kernel(
     const IType* const input, const float* global_amax, OType* const output_c,
     OType* const output_t, ScaleType* const tile_scales_inv_c, ScaleType* const tile_scales_inv_t,
     const size_t row_length, const size_t num_rows, const size_t scale_stride_x,
     const size_t scale_stride_y, const size_t scale_t_stride_x, const size_t scale_t_stride_y,
     const size_t kScaleBlockDim, const float epsilon, const size_t* rng_state,
-    const float* noop_ptr) {
+    const float* noop_ptr, const float e4m3_max, const bool use_mse) {
   constexpr int kNVecContainer = kNVecOut / kNFP4PerContainer;
   using SMemVec = Vec<IType, kNVecSMem>;
   using OVec = Vec<OType, kNVecContainer>;
@@ -478,12 +588,21 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
   // for 128x128 block, 2D block scaling means there will be 8x8 amax values for nvfp4, 4x4 for 2D mxfp4
   // use constexpr to define the size, when not using 2D, use minimal size 1x1
   constexpr int kFP4BlockScalingSize = 16;
+  constexpr float kFp4Max = 6.0f;  // TypeExtrema<fp4e2m1>::max
   constexpr int k2DBlockAmaxDim = kIs2DBlockScaling ? (kTileDim / kFP4BlockScalingSize) : 1;
   constexpr int kNumRowsPerWarp = kThreadsPerWarp / kNumThreadsStore;  // 4
   constexpr int k2DBlockAmaxReduceDim =
       kIs2DBlockScaling ? (kFP4BlockScalingSize / kNumRowsPerWarp) : 1;
   __shared__ CType amax_smem_red[k2DBlockAmaxDim][k2DBlockAmaxDim][k2DBlockAmaxReduceDim];
   __shared__ CType amax_smem[k2DBlockAmaxDim][k2DBlockAmaxDim];
+  // 4over6 scores its two candidates over the same block the scale serves, so the
+  // error is reduced over the same 16x16 block the 2D amax is. Unlike the amax it
+  // stages all 16 rows: max is associative, the fp32 sum is not, so the rows cannot
+  // be pre-combined per warp without diverging from upstream's tree.
+  constexpr int k4Over6Dim = kUse4Over6 ? k2DBlockAmaxDim : 1;
+  constexpr int k4Over6RowDim = kUse4Over6 ? kFP4BlockScalingSize : 1;
+  __shared__ CType err_smem_red[2][k4Over6Dim][k4Over6Dim][k4Over6RowDim];
+  __shared__ CType err_smem[2][k4Over6Dim][k4Over6Dim];
 
   // Step 1: Load input to shared memory
   {
@@ -532,12 +651,17 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
 
   const int kNumThreadsReduce = kScaleBlockDim / kNVecOut;
 #ifdef __HIP_PLATFORM_AMD__
-  const float global_encode_scale =
-      (kIsE8Scaling || global_amax == nullptr) ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]);
+  const float global_encode_scale = (kIsE8Scaling || global_amax == nullptr)
+                                        ? 1.0f
+                                        : ComputeGlobalEncodeScaleFP4(global_amax[0], e4m3_max);
 #else
   const float global_encode_scale =
-      kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0]);
+      kIsE8Scaling ? 1.0f : ComputeGlobalEncodeScaleFP4(global_amax[0], e4m3_max);
 #endif
+  // Mirrors the global_encode_scale guard exactly: where that falls back to 1.0f the
+  // dequantisation denominator is err_denom, so the equivalent amax is err_denom itself.
+  const float block_true_global_amax =
+      (kIsE8Scaling || global_amax == nullptr) ? kFp4Max * e4m3_max : global_amax[0];
   constexpr float fp4_max_inv = 1.0f / TypeExtrema<fp4e2m1>::max;
   const float global_encode_scale_multiplier = global_encode_scale * fp4_max_inv;
   const float global_decode_scale = 1.0 / global_encode_scale;
@@ -632,15 +756,68 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
       float row_global_encode_scale = global_encode_scale;
       if constexpr (kRowScaledNVFP4) {
         row_global_encode_scale =
-            row_idx < num_rows ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx]) : 1.0f;
+            row_idx < num_rows ? ComputeGlobalEncodeScaleFP4(global_amax[row_idx], e4m3_max) : 1.0f;
+      }
+      float row_true_global_amax = block_true_global_amax;
+      if constexpr (kRowScaledNVFP4) {
+        row_true_global_amax = row_idx < num_rows ? global_amax[row_idx] : kFp4Max * e4m3_max;
       }
       const float row_global_encode_scale_multiplier =
           kRowScaledNVFP4 ? row_global_encode_scale * fp4_max_inv : global_encode_scale_multiplier;
       const float row_global_decode_scale =
           kRowScaledNVFP4 ? 1.0f / row_global_encode_scale : global_decode_scale;
-      ScaleType scale_inv =
-          ComputeDecodeScaleFP4<ScaleType>(amax, row_global_encode_scale_multiplier);
-      float encode_scale = ComputeEncodeScaleFP4<ScaleType>(scale_inv, row_global_decode_scale);
+      ScaleType scale_inv;
+      float encode_scale;
+#ifdef __HIP_PLATFORM_AMD__
+      // 4over6 lives here only on ROCm; CUDA dispatches it to upstream's own kernel,
+      // so the helpers below are not defined there.
+      if constexpr (kUse4Over6) {
+        float blk[kNVecOut];
+#pragma unroll
+        for (int i = 0; i < kNVecOut / kNVecSMem; ++i) {
+#pragma unroll
+          for (int j = 0; j < kNVecSMem; ++j) {
+            blk[i * kNVecSMem + j] = static_cast<float>(smem_vec[i].data.elt[j]);
+          }
+        }
+        ScaleType sf6, sf4;
+        float enc6, enc4, err6, err4;
+        Compute4Over6Candidates<ScaleType>(
+            blk, kNVecOut, amax, row_global_decode_scale, use_mse, row_global_encode_scale,
+            row_true_global_amax, kFp4Max * e4m3_max, &sf6, &sf4, &enc6, &enc4, &err6, &err4);
+        if constexpr (kIs2DBlockScaling) {
+          // The scale spans a 16x16 block, so the error must too. Stage this row's
+          // error and let one thread run upstream's tree over the 16; indices
+          // recomputed rather than hoisted so the non-4over6 code is untouched.
+          constexpr int kNumRowsPerIter2 = kThreadsPerBlock / kNumThreadsStore;
+          const int warp_i = threadIdx.x / kThreadsPerWarp;
+          const int tx = threadIdx.x % kNumThreadsStore;
+          const int ty = (threadIdx.x / kNumThreadsStore) % kNumRowsPerWarp;
+          const int row_i = iter * kNumRowsPerIter2 + warp_i * kNumRowsPerWarp + ty;
+          const int blk_i = row_i / kFP4BlockScalingSize;
+          const int row_in_blk = row_i % kFP4BlockScalingSize;
+          err_smem_red[0][blk_i][tx][row_in_blk] = err6;
+          err_smem_red[1][blk_i][tx][row_in_blk] = err4;
+          __syncthreads();
+          if (row_in_blk == 0) {
+            err_smem[0][blk_i][tx] =
+                reduceBlockError<kFP4BlockScalingSize>(err_smem_red[0][blk_i][tx]);
+            err_smem[1][blk_i][tx] =
+                reduceBlockError<kFP4BlockScalingSize>(err_smem_red[1][blk_i][tx]);
+          }
+          __syncthreads();
+          err6 = err_smem[0][blk_i][tx];
+          err4 = err_smem[1][blk_i][tx];
+        }
+        const bool pick4 = err4 < err6;  // tie keeps map-to-6, matching upstream
+        scale_inv = pick4 ? sf4 : sf6;
+        encode_scale = pick4 ? enc4 : enc6;
+      } else  // NOLINT(readability/braces)
+#endif
+      {
+        scale_inv = ComputeDecodeScaleFP4<ScaleType>(amax, row_global_encode_scale_multiplier);
+        encode_scale = ComputeEncodeScaleFP4<ScaleType>(scale_inv, row_global_decode_scale);
+      }
       // Step 2.5: Write scale_inv
       bool write_scale_inv = is_src_lane;
       if constexpr (!kAligned) {
@@ -752,6 +929,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
         amax_smem[data_row_idx / kFP4BlockScalingSize][tid_in_warp_x] = amax_2d;
       }
       __syncthreads();
+      // No 4over6 work here: Step 3 scores the columnwise candidates on their own
+      // column groups, so it only needs amax_smem from this pass.
       r_s += r_stride;
     }
   }
@@ -779,6 +958,55 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
     // This mask represents which threads should do the reduction together.
     const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane;
     const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0;
+#ifdef __HIP_PLATFORM_AMD__
+    // Step 3.0: 4over6 block errors for the whole tile.
+    // Upstream scores the columnwise candidates on their own column groups
+    // (quantize_4over6_nvfp4.cuh::quantize_stage_colwise) instead of reusing what the
+    // rowwise pass computed. Both passes cover the same 16x16 block but sum its errors
+    // along different axes, and fp32 addition is not associative, so the two can land on
+    // different candidates for one block; that is the algorithm, not a bug.
+    //
+    // Done here rather than inside the loop below because a block's 16 columns are split
+    // across both smem_idx halves, so no single pass of that loop sees a whole block.
+    // Keeping it separate also leaves nothing live across the barriers.
+    if constexpr (kUse4Over6 && kIs2DBlockScaling) {
+      const int row_blk = threadIdx.x % kNumThreadsStore;
+      int c = c_s;
+#pragma unroll
+      for (int iter = 0; iter < num_iterations; ++iter, c += c_stride) {
+#pragma unroll
+        for (int smem_idx = 0; smem_idx < kNVecSMem; ++smem_idx) {
+          const int col_in_tile = c * kNVecSMem + smem_idx;
+          float blk[kNVecOut];
+#pragma unroll
+          for (int i = 0; i < kNVecOut; ++i) {
+            blk[i] = static_cast<float>(smem[(r_s + i) * kSMemCol + c].data.elt[smem_idx]);
+          }
+          ScaleType sf6, sf4;
+          float enc6, enc4, err6, err4;
+          Compute4Over6Candidates<ScaleType>(
+              blk, kNVecOut, amax_smem[row_blk][col_in_tile / kFP4BlockScalingSize],
+              global_decode_scale, use_mse, global_encode_scale, block_true_global_amax,
+              kFp4Max * e4m3_max, &sf6, &sf4, &enc6, &enc4, &err6, &err4);
+          err_smem_red[0][row_blk][col_in_tile / kFP4BlockScalingSize]
+                      [col_in_tile % kFP4BlockScalingSize] = err6;
+          err_smem_red[1][row_blk][col_in_tile / kFP4BlockScalingSize]
+                      [col_in_tile % kFP4BlockScalingSize] = err4;
+        }
+      }
+      __syncthreads();
+      // One thread per (row block, column block) runs upstream's tree over the 16 staged
+      // column errors; the columns paired first are 8 apart and so live in different
+      // warps, exactly as the rows do in the rowwise pass.
+      for (int idx = threadIdx.x; idx < k2DBlockAmaxDim * k2DBlockAmaxDim; idx += blockDim.x) {
+        const int rb = idx / k2DBlockAmaxDim;
+        const int cb = idx % k2DBlockAmaxDim;
+        err_smem[0][rb][cb] = reduceBlockError<kFP4BlockScalingSize>(err_smem_red[0][rb][cb]);
+        err_smem[1][rb][cb] = reduceBlockError<kFP4BlockScalingSize>(err_smem_red[1][rb][cb]);
+      }
+      __syncthreads();
+    }
+#endif
 #pragma unroll
     for (int iter = 0; iter < num_iterations; ++iter) {
       SMemVec smem_vec[kNVecOut];
@@ -823,9 +1051,39 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
           amax = __shfl_sync(mask, amax, src_lane);
         }
         // Step 3.4: Compute scale
-        ScaleType scale_inv =
-            ComputeDecodeScaleFP4<ScaleType>(amax, global_encode_scale_multiplier);
-        float encode_scale = ComputeEncodeScaleFP4<ScaleType>(scale_inv, global_decode_scale);
+        ScaleType scale_inv;
+        float encode_scale;
+#ifdef __HIP_PLATFORM_AMD__
+        if constexpr (kUse4Over6) {
+          ScaleType sf6, sf4;
+          float enc6, enc4, err6, err4;
+          if constexpr (kIs2DBlockScaling) {
+            // Step 3.0 already scored this block; only the two scales are needed here.
+            const int row_blk = threadIdx.x % kNumThreadsStore;
+            const int col_blk = (c_s * kNVecSMem + smem_idx) / kFP4BlockScalingSize;
+            err6 = err_smem[0][row_blk][col_blk];
+            err4 = err_smem[1][row_blk][col_blk];
+            Compute4Over6Scales<ScaleType>(amax, global_decode_scale, global_encode_scale, &sf6,
+                                           &sf4, &enc6, &enc4);
+          } else {
+            float blk[kNVecOut];
+#pragma unroll
+            for (int i = 0; i < kNVecOut; ++i) {
+              blk[i] = static_cast<float>(smem_vec[i].data.elt[smem_idx]);
+            }
+            Compute4Over6Candidates<ScaleType>(
+                blk, kNVecOut, amax, global_decode_scale, use_mse, global_encode_scale,
+                block_true_global_amax, kFp4Max * e4m3_max, &sf6, &sf4, &enc6, &enc4, &err6, &err4);
+          }
+          const bool pick4 = err4 < err6;  // tie keeps map-to-6, matching upstream
+          scale_inv = pick4 ? sf4 : sf6;
+          encode_scale = pick4 ? enc4 : enc6;
+        } else  // NOLINT(readability/braces)
+#endif
+        {
+          scale_inv = ComputeDecodeScaleFP4<ScaleType>(amax, global_encode_scale_multiplier);
+          encode_scale = ComputeEncodeScaleFP4<ScaleType>(scale_inv, global_decode_scale);
+        }
         // Step 3.5: Write scale_inv_t
         bool write_scale_inv = is_src_lane;
         if constexpr (!kAligned) {
@@ -886,7 +1144,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
       }
     }
   }
-}
+}  // NOLINT(readability/fn_size) 4over6 grows upstream's kernel past 500 lines
 
 }  // namespace
 }  // namespace quantize_transpose_nvfp4
@@ -900,7 +1158,8 @@ void quantize_transpose_vector_blockwise_fp4(
     const bool return_identity, const bool return_transpose, const bool pow2_scale,
     const bool swizzled_scale, const bool use_stochastic_rounding,
     const NVTETensor rng_state_tensor, const bool use_2d_quantization, const bool row_scaled_nvfp4,
-    const SimpleTensor& noop_tensor, cudaStream_t stream) {
+    const SimpleTensor& noop_tensor, const int nvfp4_e4m3_max,
+    const NVTENVFP44Over6Mode nvfp4_4over6_mode, cudaStream_t stream) {
   NVTE_API_CALL(quantize_transpose_vector_blockwise_fp4);
 #if defined(__HIP_PLATFORM_AMD__) || CUDA_VERSION >= 12080
 
@@ -929,6 +1188,14 @@ void quantize_transpose_vector_blockwise_fp4(
     return;
   }
 
+  // gfx942 stores E4M3 as FNUZ (max 240), gfx950 as OCP (max 448). The quantizer must use the
+  // same bound the GEMM's dequant does (rocm_gemm.cu:539), or the two disagree by 448/240.
+#ifdef __HIP_PLATFORM_AMD__
+  const float e4m3_max_for_arch = te_fp8_fnuz() ? 240.0f : static_cast<float>(nvfp4_e4m3_max);
+#else
+  const float e4m3_max_for_arch = static_cast<float>(nvfp4_e4m3_max);
+#endif
+
   size_t scale_stride_x = 0;
   size_t scale_stride_y = 0;
 
@@ -1009,48 +1276,60 @@ void quantize_transpose_vector_blockwise_fp4(
                                   TRANSFORMER_ENGINE_SWITCH_CONDITION(
                                       row_scaled_nvfp4, kRowScaledNVFP4,
 
+  // 4over6 only has a body on ROCm, so pin the switch
+  // false on CUDA rather than instantiate every
+  // specialization twice for a branch it cannot take.
+#ifdef __HIP_PLATFORM_AMD__
+                                      TRANSFORMER_ENGINE_SWITCH_CONDITION(
+                                          nvfp4_4over6_mode != kNVTENVFP44Over6Disabled, kUse4Over6,
+#else
+                                      TRANSFORMER_ENGINE_SWITCH_CONDITION(
+                                          false, kUse4Over6,
+#endif
+
 #ifdef __HIP_PLATFORM_AMD__
-                                      size_t smem_bytes =
-                                          smem_size_for_tile(tile_dim) * sizeof(InputType);
+                                          size_t smem_bytes =
+                                              smem_size_for_tile(tile_dim) * sizeof(InputType);
 #else
-                                      size_t smem_bytes = kSMemSize * sizeof(InputType);
+                                          size_t smem_bytes = kSMemSize * sizeof(InputType);
 #endif
-                                      auto kernel = block_scaled_1d_cast_transpose_kernel<
-                                          kReturnIdentity, kReturnTranspose, kPow2Scale, kAligned,
-                                          float, InputType, OutputType, ScaleType, kSwizzledScale,
-                                          kApplyStochasticRounding, kIs2DBlockScaling,
-                                          kRowScaledNVFP4>;
-                                      if (smem_bytes >= 48 * 1024) {
-                                        cudaError_t err = cudaFuncSetAttribute(
-                                            kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
-                                            smem_bytes);
-                                        NVTE_CHECK(err == cudaSuccess,
-                                                   "Failed to set dynamic shared memory size.");
+                                          auto kernel = block_scaled_1d_cast_transpose_kernel<
+                                              kReturnIdentity, kReturnTranspose, kPow2Scale,
+                                              kAligned, float, InputType, OutputType, ScaleType,
+                                              kSwizzledScale, kApplyStochasticRounding,
+                                              kIs2DBlockScaling, kRowScaledNVFP4, kUse4Over6>;
+                                          if (smem_bytes >= 48 * 1024) {
+                                            cudaError_t err = cudaFuncSetAttribute(
+                                                kernel, cudaFuncAttributeMaxDynamicSharedMemorySize,
+                                                smem_bytes);
+                                            NVTE_CHECK(err == cudaSuccess,
+                                                       "Failed to set dynamic shared memory size.");
 #ifdef __HIP_PLATFORM_AMD__
-                                      } kernel<<<grid, tile_dim * 2, smem_bytes,
-                                                 stream>>>(
+                                          } kernel<<<grid, tile_dim * 2, smem_bytes,
+                                                     stream>>>(
 #else
-                                      } kernel<<<grid, kThreadsPerBlock, smem_bytes,
-                                                 stream>>>(
+                                          } kernel<<<grid, kThreadsPerBlock, smem_bytes,
+                                                     stream>>>(
 #endif
-                                          reinterpret_cast<const InputType*>(input.dptr),
-                                          reinterpret_cast<const float*>(global_amax.dptr),
-                                          reinterpret_cast<OutputType*>(output.dptr),
-                                          reinterpret_cast<OutputType*>(output_t.dptr),
-                                          reinterpret_cast<ScaleType*>(scale_inv.dptr),
-                                          reinterpret_cast<ScaleType*>(scale_inv_t.dptr),
-                                          row_length, num_rows, scale_stride_x, scale_stride_y,
-                                          scale_t_stride_x, scale_t_stride_y, kScaleBlockDim,
-                                          epsilon, rng_state,
-                                          noop_ptr);)  // kRowScaledNVFP4
-                                  )                    // kIs2DBlockScaling
-                              )                        // kApplyStochasticRounding
-                          )                            // kSwizzledScale
-                      )                                // kAligned
-                  )                                    // kReturnTranspose
-              )                                        // kReturnIdentity
-          )                                            // OutputType
-      )                                                // InputType
+                                              reinterpret_cast<const InputType*>(input.dptr),
+                                              reinterpret_cast<const float*>(global_amax.dptr),
+                                              reinterpret_cast<OutputType*>(output.dptr),
+                                              reinterpret_cast<OutputType*>(output_t.dptr),
+                                              reinterpret_cast<ScaleType*>(scale_inv.dptr),
+                                              reinterpret_cast<ScaleType*>(scale_inv_t.dptr),
+                                              row_length, num_rows, scale_stride_x, scale_stride_y,
+                                              scale_t_stride_x, scale_t_stride_y, kScaleBlockDim,
+                                              epsilon, rng_state, noop_ptr, e4m3_max_for_arch,
+                                              nvfp4_4over6_mode ==
+                                                  kNVTENVFP44Over6MinMSE);))  // kUse4Over6 / kRowScaledNVFP4
+                                  )                                           // kIs2DBlockScaling
+                              )  // kApplyStochasticRounding
+                          )      // kSwizzledScale
+                      )          // kAligned
+                  )              // kReturnTranspose
+              )                  // kReturnIdentity
+          )                      // OutputType
+      )                          // InputType
 
   NVTE_CHECK_CUDA(cudaGetLastError());
 #else
diff --git a/transformer_engine/common/util/standalone_topk.cuh b/transformer_engine/common/util/standalone_topk.cuh
index 3d19cbfcf2..190027acc0 100644
--- a/transformer_engine/common/util/standalone_topk.cuh
+++ b/transformer_engine/common/util/standalone_topk.cuh
@@ -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.
@@ -7,7 +9,12 @@
 #pragma once
 
 #include <cooperative_groups.h>
+#ifndef __HIP_PLATFORM_AMD__
+// Only scan_warp_version needs cg::reduce, and that is excluded below on ROCm;
+// HIP keeps the reduce overloads behind an amd_detail header the umbrella one
+// does not pull in.
 #include <cooperative_groups/reduce.h>
+#endif
 #include <cuda_bf16.h>
 
 #include <cub/device/device_radix_sort.cuh>
@@ -15,6 +22,12 @@
 #include <vector>
 namespace cg = cooperative_groups;
 
+#ifdef __HIP_PLATFORM_AMD__
+// hipify rewrites the CUB class templates but not Traits or the BLOCK_*
+// enumerators, so resolve the rest through the namespace.
+namespace cub = hipcub;
+#endif
+
 // Workspace pointer-alignment helpers.
 inline size_t calc_aligned_size(const std::vector<size_t> &sizes) {
   const size_t ALIGN_BYTES = 256;
@@ -485,6 +498,9 @@ __device__ void choose_bucket(Counter<T, IdxT> *counter, const IdxT *histogram,
   }
 }
 
+// Dead: the call site below is commented out upstream. Excluded on ROCm rather
+// than ported -- it is a wave32 reduction over a 32-bit mask, which HIP rejects.
+#ifndef __HIP_PLATFORM_AMD__
 template <typename T, typename IdxT, int BitsPerPass, int BlockSize>
 __device__ void scan_warp_version(cg::thread_block_tile<WARP_SIZE> const &warp,
                                   volatile IdxT *histogram, Counter<T, IdxT> *counter, const IdxT k,
@@ -569,6 +585,7 @@ __device__ void scan_warp_version(cg::thread_block_tile<WARP_SIZE> const &warp,
     // }
   }
 }
+#endif  // !__HIP_PLATFORM_AMD__ (scan_warp_version)
 // For one-block version, last_filter() could be called when pass < num_passes
 // - 1. So `pass` could not be constexpr
 template <typename T, typename IdxT, int BitsPerPass, bool store_out>
diff --git a/transformer_engine/common/util/topk.cu b/transformer_engine/common/util/topk.cu
index 21018a4948..33f03fc910 100644
--- a/transformer_engine/common/util/topk.cu
+++ b/transformer_engine/common/util/topk.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.
@@ -49,7 +51,13 @@ void nvte_topk(cudaStream_t stream, const NVTETensor keys_in, const NVTETensor l
   } while (0)
 
   if (dtype == DType::kBFloat16) {
+#ifdef __HIP_PLATFORM_AMD__
+    // hipCUB specializes NumericTraits for hip_bfloat16, not __hip_bfloat16,
+    // and the radix select needs those traits to twiddle the key bits.
+    DISPATCH_TOPK(hip_bfloat16);
+#else
     DISPATCH_TOPK(__nv_bfloat16);
+#endif
   } else if (dtype == DType::kFloat32) {
     DISPATCH_TOPK(float);
   } else {
diff --git a/transformer_engine/jax/csrc/extensions/pybind.cpp b/transformer_engine/jax/csrc/extensions/pybind.cpp
index 3fc054303d..725ac26e80 100644
--- a/transformer_engine/jax/csrc/extensions/pybind.cpp
+++ b/transformer_engine/jax/csrc/extensions/pybind.cpp
@@ -151,9 +151,7 @@ pybind11::dict Registrations() {
 #endif  // NVTE_WITH_NCCL_EP
 
   // TopK
-#ifndef USE_ROCM  // Disabled on ROCm
   dict["te_topk_ffi"] = EncapsulateFFI(TopkHandler);
-#endif
 
   return dict;
 }
@@ -177,9 +175,7 @@ PYBIND11_MODULE(transformer_engine_jax, m) {
   m.def("get_norm_bwd_workspace_sizes", &GetNormBackwardWorkspaceSizes);
   m.def("get_fused_attn_fwd_workspace_sizes", &GetFusedAttnForwardWorkspaceSizes);
   m.def("get_fused_attn_bwd_workspace_sizes", &GetFusedAttnBackwardWorkspaceSizes);
-#ifndef USE_ROCM
   m.def("get_topk_workspace_sizes", &GetTopkWorkspaceSizes);
-#endif
   m.def("nvte_get_qkv_format", &nvte_get_qkv_format);
   m.def("is_non_nt_fp8_gemm_supported", &nvte_is_non_tn_fp8_gemm_supported);
 #ifndef USE_ROCM
diff --git a/transformer_engine/jax/csrc/extensions/topk.cpp b/transformer_engine/jax/csrc/extensions/topk.cpp
index 0cc68a45af..450ff08b38 100644
--- a/transformer_engine/jax/csrc/extensions/topk.cpp
+++ b/transformer_engine/jax/csrc/extensions/topk.cpp
@@ -1,6 +1,4 @@
 /*************************************************************************
- * 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.
@@ -14,9 +12,6 @@
 namespace transformer_engine {
 namespace jax {
 
-// Disabled on ROCm
-#ifndef USE_ROCM
-
 // ---------------------------------------------------------------------------
 // JAX FFI handler
 // ---------------------------------------------------------------------------
@@ -105,7 +100,5 @@ pybind11::tuple GetTopkWorkspaceSizes(int batch_size, int seq_len, int k) {
   return pybind11::make_tuple(std::make_pair(work_shape, workspace_tensor.dtype()));
 }
 
-#endif
-
 }  // namespace jax
 }  // namespace transformer_engine
diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index 0a47e600da..c6b312434c 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -13,6 +13,12 @@
 #include "transformer_engine/transformer_engine.h"
 
 #ifdef USE_ROCM
+#include <ATen/cuda/CUDAContext.h>
+
+#include <map>
+#include <mutex>
+#include <utility>
+
 #include "common/common.h"
 #endif
 
@@ -367,10 +373,25 @@ at::Tensor allocate_amax_workspace(const TensorWrapper& input_tensor) {
     return at::empty(0, at::CUDA(at::kFloat));
   }
 
-  const auto N = input_tensor.numel();
-  size_t workspace_blocks = nvte_amax_workspace_num_blocks(N);
-
-  return at::empty(workspace_blocks, at::CUDA(at::kFloat));
+  // A per-call tensor returns to the caching allocator at the caller's closing brace while
+  // amax_kernel and amax_final_reduce are still only enqueued. Stream-ordered reuse covers
+  // that in eager mode, but graph capture replaces stream order with a DAG, so the freed
+  // block is aliased in the graph mempool and replays read partials they never wrote.
+  // One buffer per (device, stream), sized to the block cap: same stream serializes.
+  static std::mutex amax_ws_mutex;
+  static std::map<std::pair<int, cudaStream_t>, at::Tensor> amax_ws_cache;
+
+  const int device_id = at::cuda::current_device();
+  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+  std::lock_guard<std::mutex> lock(amax_ws_mutex);
+  auto& ws = amax_ws_cache[std::make_pair(device_id, stream)];
+  if (!ws.defined()) {
+    // Any N past the block cap yields the cap itself; avoid SIZE_MAX so DIVUP cannot overflow.
+    ws = at::empty(nvte_amax_workspace_num_blocks(static_cast<size_t>(1) << 40),
+                   at::CUDA(at::kFloat));
+  }
+  return ws;
 }
 
 #endif

From 9aec031f3f7a838f3e3be7e272cfbb981dafdf14 Mon Sep 17 00:00:00 2001
From: AllenFarcas <alfarcas@amd.com>
Date: Mon, 10 Aug 2026 01:33:01 +0000
Subject: [PATCH 2/5] Trim explanatory comments to the constraints they record

Comments that reconstructed the reasoning behind a fix rather than stating the
constraint it depends on. Cut to one or two lines each, keeping only what
silently breaks if removed -- notably the -ffast-math note on cdna4, without
which the empty asm looks removable and the epilogue reverts to an fp32
accumulate.

Comment-only; no code changed.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 tests/cpp/operator/test_cast_nvfp4_transpose.cu    | 14 ++++----------
 .../cpp/operator/test_dequantize_mxfp8_grouped.cu  |  4 +---
 transformer_engine/common/CMakeLists.txt           |  6 ++----
 .../common/cast/dispatch/quantize.cuh              | 14 ++++----------
 .../kittens/cdna3/blockwise_fp8_gemm_helper.cuh    |  9 ++-------
 .../kittens/cdna4/blockwise_fp8_gemm_helper.cuh    | 13 +++----------
 6 files changed, 16 insertions(+), 44 deletions(-)

diff --git a/tests/cpp/operator/test_cast_nvfp4_transpose.cu b/tests/cpp/operator/test_cast_nvfp4_transpose.cu
index 30aea9d33b..40c40d86f9 100644
--- a/tests/cpp/operator/test_cast_nvfp4_transpose.cu
+++ b/tests/cpp/operator/test_cast_nvfp4_transpose.cu
@@ -888,12 +888,9 @@ void compare_nvfp4_4over6_candidates(const std::string& name,
     size_t wrong_candidate = 0;
     size_t expected_map4 = 0;
     size_t expected_map6 = 0;
-    // Blocks whose two candidates encode to the same bytes: a selection cannot be
-    // read back out of them, so they neither pin the kernel nor weaken the check.
+    // Both candidates encode identically; no selection is observable.
     size_t identical_candidates = 0;
-    // Blocks that do encode differently but whose candidate errors are too close
-    // for the host to call. These are the only ones where the kernel keeps any
-    // latitude, so their count is the coverage this check gives up.
+    // Candidate errors too close for the host to call.
     size_t unpinned_blocks = 0;
 
     for (size_t row = 0; row < rows; ++row) {
@@ -1169,9 +1166,7 @@ void performTest(float (*OP)(const float),
                                e4m3_max,
                                NVFP4FourOverSixCandidate::Map6);
 
-        // The kernel is free to pick either candidate per block, but not freely:
-        // it must pick the one with the smaller round-trip error. Derive that
-        // choice here so the checker can hold it to exactly one reference.
+        // Reference choice: the candidate with the smaller round-trip error.
         const bool use_mse = mode == kNVTENVFP44Over6MinMSE;
         expected_decisions = compute_4over6_expected_decisions<InputType>(
             OP, input.rowwise_cpu_dptr<InputType>(), rows, cols, scales_stride, ref_amax.data(),
@@ -1206,8 +1201,7 @@ void performTest(float (*OP)(const float),
     hipDeviceProp_t prop;
     hipGetDeviceProperties(&prop, 0);
     const bool is_gfx950 = prop.major == 9 && prop.minor == 5;
-    // 4over6 excluded: the quantize dispatch refuses it combined with stochastic
-    // rounding, the same way this test already excludes it from fast math below.
+    // The quantize dispatch refuses 4over6 combined with stochastic rounding.
     for (bool use_stochastic_rounding : (is_gfx950 && !use_4over6
                                              ? std::vector<bool>{false, true}
                                              : std::vector<bool>{false})) {
diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
index e13eb7ab4d..bfec25328e 100644
--- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
+++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
@@ -346,9 +346,7 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
 
     int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset,
                         tensor_elts * sizeof(OutputType));
-    // memcmp is the assertion; the loop below only locates the mismatch. Comparing with != as
-    // the gate would let -0.0 vs +0.0 and distinct NaN payloads pass after memcmp caught them,
-    // and the e8m0 scale bytes are drawn across the full [0,255] range including NaN.
+    // memcmp is the assertion; != would let -0.0 vs +0.0 and NaN payloads through.
     ASSERT_EQ(result, 0) << "Bitwise mismatch in tensor " << t;
     for (size_t i = 0; i < tensor_elts; ++i) {
       if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) {
diff --git a/transformer_engine/common/CMakeLists.txt b/transformer_engine/common/CMakeLists.txt
index 5fbb9a0146..0c9b3b506a 100644
--- a/transformer_engine/common/CMakeLists.txt
+++ b/transformer_engine/common/CMakeLists.txt
@@ -443,10 +443,8 @@ else() #USE_ROCM
   include("${CMAKE_CURRENT_SOURCE_DIR}/../../build_tools/hipify/hipify.cmake")
   TE_Hipify(${CMAKE_CURRENT_SOURCE_DIR})
 
-  # TE_Hipify runs at configure time via execute_process, so CMake holds no dependency on the
-  # CUDA sources it reads. Without registering them, editing a source here leaves the previously
-  # generated _hip.* in place and the build silently compiles the old revision. Generated outputs
-  # are excluded: they are rewritten during configure, so listing them would re-trigger forever.
+  # TE_Hipify runs at configure time, so its inputs must be registered or edits go unnoticed.
+  # Generated _hip.* are excluded to avoid re-triggering configure on their own output.
   set(_hipify_deps)
   foreach(_src ${transformer_engine_SOURCES})
     get_filename_component(_abs "${CMAKE_CURRENT_SOURCE_DIR}/${_src}" ABSOLUTE)
diff --git a/transformer_engine/common/cast/dispatch/quantize.cuh b/transformer_engine/common/cast/dispatch/quantize.cuh
index bba249dd79..4afaf29a59 100644
--- a/transformer_engine/common/cast/dispatch/quantize.cuh
+++ b/transformer_engine/common/cast/dispatch/quantize.cuh
@@ -140,8 +140,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
       NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding,
                  "NVFP4 4over6 quantization does not support stochastic rounding.");
 #ifdef __HIP_PLATFORM_AMD__
-      // The fast-math error path is an optimisation of candidate selection, not a
-      // capability; refuse it rather than silently scoring with the exact one.
+      // Refuse the fast-math error path rather than silently scoring with the exact one.
       NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math,
                  "NVFP4 4over6 fast-math error mode is not supported on ROCm.");
 #endif
@@ -160,9 +159,7 @@ void quantize_fwd_helper(const NVTETensor input, NVTETensor output,
           (output_tensor->has_data() ||
            (output_tensor->has_columnwise_data() && quant_config_cpp.nvfp4_2d_quantization));
 
-      // Launch NVFP4 quantize kernel. Upstream's dedicated 4over6 and optimized
-      // quantize_transpose kernels are CUDA-only (Blackwell); ROCm falls through to the portable
-      // blockwise path below, which implements 4over6 and row-scaled NVFP4 itself.
+      // Launch NVFP4 quantize kernel. ROCm uses the portable blockwise path below.
 #ifndef __HIP_PLATFORM_AMD__
       if (nvfp4_use_4over6) {
         if (quant_config_cpp.nvfp4_2d_quantization) {
@@ -327,8 +324,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
       NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.stochastic_rounding,
                  "NVFP4 4over6 quantization does not support stochastic rounding.");
 #ifdef __HIP_PLATFORM_AMD__
-      // The fast-math error path is an optimisation of candidate selection, not a
-      // capability; refuse it rather than silently scoring with the exact one.
+      // Refuse the fast-math error path rather than silently scoring with the exact one.
       NVTE_CHECK(!nvfp4_use_4over6 || !quant_config_cpp.nvfp4_4over6_err_use_fast_math,
                  "NVFP4 4over6 fast-math error mode is not supported on ROCm.");
 #endif
@@ -342,9 +338,7 @@ void quantize_bwd_helper(const NVTETensor grad, const NVTETensor input, NVTETens
           (output_tensor->has_data() ||
            (output_tensor->has_columnwise_data() && quant_config_cpp.nvfp4_2d_quantization));
 
-      // Launch NVFP4 quantize kernel. Upstream's dedicated 4over6 and optimized
-      // quantize_transpose kernels are CUDA-only (Blackwell); ROCm falls through to the portable
-      // blockwise path below, which implements 4over6 and row-scaled NVFP4 itself.
+      // Launch NVFP4 quantize kernel. ROCm uses the portable blockwise path below.
 #ifndef __HIP_PLATFORM_AMD__
       if (nvfp4_use_4over6) {
         if (quant_config_cpp.nvfp4_2d_quantization) {
diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
index 0ae6c4eb2b..086419a65d 100644
--- a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
+++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
@@ -102,13 +102,8 @@ __device__ inline OType convert_out(float v) {
     }
 }
 
-// cuBLAS converts the GEMM result to the output type before accumulating beta*C, so the
-// blockwise reference does the same (see qgemm() in
-// tests/pytorch/references/blockwise_fp8_gemm_reference.py). Keeping the accumulator in
-// fp32 across the beta add is more accurate but does not reproduce that result, so the
-// round is deliberate rather than an oversight. It has to go through convert_out so that
-// the rtne_bias correction is applied exactly where the store applies it -- fp16 must not
-// pick up a bf16-granularity bias.
+// Deliberate: the reference rounds to the output type before the beta*C add
+// (blockwise_fp8_gemm_reference.py::qgemm). convert_out keeps rtne_bias at store granularity.
 template <typename OType>
 __device__ inline float round_to_out_dtype(float v) {
     if constexpr (std::is_same_v<OType, float>) {
diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
index f173dcc49c..bc9be22847 100644
--- a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
+++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
@@ -83,16 +83,9 @@ __device__ inline float read_elem(const void *p, int dtype, int idx) {
     return reinterpret_cast<const float *>(p)[idx];
 }
 
-// cuBLAS converts the GEMM result to the output type before accumulating beta*C, so the
-// blockwise reference does the same (see qgemm() in
-// tests/pytorch/references/blockwise_fp8_gemm_reference.py). Keeping the accumulator in
-// fp32 across the beta add is more accurate but does not reproduce that result, so the
-// round is deliberate rather than an oversight.
-//
-// This translation unit is built with -ffast-math (see gemm/kittens/CMakeLists.txt), which
-// folds fpext(fptrunc(x)) straight back to x -- written plainly, the round-trip below emits
-// no conversion at all and the whole epilogue silently reverts to an fp32 accumulate. The
-// empty asm makes the narrowed value opaque so the conversion has to happen.
+// Deliberate: the reference rounds to the output type before the beta*C add
+// (blockwise_fp8_gemm_reference.py::qgemm). This TU is built with -ffast-math, which folds
+// fpext(fptrunc(x)) back to x, so the empty asm is required to keep the conversion.
 template <typename OType>
 __device__ inline float round_to_out_dtype(float v) {
     if constexpr (std::is_same_v<OType, float>) {

From 14ce6cadf0d8e49d3a8d93c83b5624daecac0073 Mon Sep 17 00:00:00 2001
From: AllenFarcas <alfarcas@amd.com>
Date: Mon, 10 Aug 2026 22:11:30 +0000
Subject: [PATCH 3/5] Address PR review findings

- Collapse the CUDA arm of the 4over6 switch to a single instantiation.
  TRANSFORMER_ENGINE_SWITCH_CONDITION is a plain if/else, so both arms were
  always instantiated; pinning the condition to false did not prevent the dead
  copy. NVTE_SWITCH_4OVER6 expands to one arm on CUDA.

- Keep the grouped dequantize shape guards outside graph capture only where
  they need a host sync. The checks on last_logical_dim and
  first_logical_dim % num_tensors read host fields and are now unconditional;
  the two that read device-resident per-tensor dims raise a clear error under
  capture rather than silently skipping validation.

- Restore the AMD copyright line on jax/csrc/extensions/topk.cpp.

- Limit the blockwise FP8 accumulator conversion to the output dtypes the
  exactness suite covers. fp16 is not parametrized there and cdna3 has never
  been executed, so both arches keep their previous fp16 behaviour.

- Leak the amax workspace cache. It holds device memory and must not be
  destroyed during static teardown.

- Gate the 4over6 columnwise-only test to ROCm; it exercises the portable
  blockwise path and has never been run on CUDA.

- Move a ROCm skip to the top of check_quantization_nvfp4_versus_reference,
  alongside the other early skips.

- Report the mismatching element before asserting in the grouped dequantize
  test; ASSERT_EQ returns first, so the locator loop was unreachable.

- Scope bad_last_dim to the ROCm arm and trim comments that recorded reasoning
  rather than a constraint.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .../operator/test_dequantize_mxfp8_grouped.cu | 20 ++---
 .../nvfp4/test_nvfp4_quantize_exact.py        |  5 +-
 .../cast/mxfp8/group_dequantize_mxfp8.cuh     | 35 +++++----
 .../common/fused_attn/flash_attn.cu           |  4 -
 .../cdna3/blockwise_fp8_gemm_helper.cuh       |  6 +-
 .../cdna4/blockwise_fp8_gemm_helper.cuh       |  5 +-
 ...quantize_transpose_vector_blockwise_fp4.cu | 75 +++++++------------
 .../jax/csrc/extensions/topk.cpp              |  2 +
 transformer_engine/pytorch/csrc/common.cpp    | 12 ++-
 9 files changed, 75 insertions(+), 89 deletions(-)

diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
index bfec25328e..bc42ed2cc5 100644
--- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
+++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
@@ -231,11 +231,11 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
                                   &offsets_tensor, sizeof(offsets_tensor));
   }
 
+#ifdef __HIP_PLATFORM_AMD__
   bool bad_last_dim = false;
   for (size_t t = 0; t < num_tensors; ++t) {
     if (last_dims_h[t] % 16 != 0) bad_last_dim = true;
   }
-#ifdef __HIP_PLATFORM_AMD__
   auto free_device = [&]() {
     cudaFree(in_data_d);
     cudaFree(out_grouped_d);
@@ -346,16 +346,18 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
 
     int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset,
                         tensor_elts * sizeof(OutputType));
-    // memcmp is the assertion; != would let -0.0 vs +0.0 and NaN payloads through.
-    ASSERT_EQ(result, 0) << "Bitwise mismatch in tensor " << t;
-    for (size_t i = 0; i < tensor_elts; ++i) {
-      if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) {
-        GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i
-                     << " (global offset " << (data_offset + i) << "): grouped="
-                     << static_cast<float>(out_grouped_h[data_offset + i])
-                     << " vs reference=" << static_cast<float>(out_ref_h[data_offset + i]);
+    if (result != 0) {
+      for (size_t i = 0; i < tensor_elts; ++i) {
+        if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) {
+          GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i
+                       << " (global offset " << (data_offset + i) << "): grouped="
+                       << static_cast<float>(out_grouped_h[data_offset + i])
+                       << " vs reference=" << static_cast<float>(out_ref_h[data_offset + i]);
+        }
       }
     }
+    // memcmp is the assertion; != would let -0.0 vs +0.0 and NaN payloads through.
+    ASSERT_EQ(result, 0) << "Bitwise mismatch in tensor " << t;
   }
 
   // Cleanup
diff --git a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
index 0a01f08116..0eebe5d331 100644
--- a/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
+++ b/tests/pytorch/nvfp4/test_nvfp4_quantize_exact.py
@@ -117,6 +117,8 @@ def check_quantization_nvfp4_versus_reference(
 ) -> None:
     if nvfp4_e4m3_max != 448 and not use_4over6:
         pytest.skip("E4M3 max 256 is only meaningful for 4over6")
+    if IS_HIP_EXTENSION and use_4over6 and nvfp4_4over6_err_use_fast_math:
+        pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm")
     maybe_skip_row_scaled_unsupported_quantization(
         row_scaled_nvfp4, return_transpose, with_2d_quantization, use_4over6, x_dtype, M, N
     )
@@ -148,8 +150,6 @@ def check_quantization_nvfp4_versus_reference(
     )
 
     if use_4over6:
-        if IS_HIP_EXTENSION and nvfp4_4over6_err_use_fast_math:
-            pytest.skip("NVFP4 4over6 fast-math error mode is not supported on ROCm")
         with nvfp4_4over6_err_fast_math(nvfp4_4over6_err_use_fast_math):
             if use_cpp_allocator:
                 x_nvfp4_sut = nvfp4_quantizer(x)
@@ -818,6 +818,7 @@ def _make_quantizer(*, rowwise: bool, columnwise: bool) -> NVFP4Quantizer:
     assert out_col_only._rowwise_data is None
 
 
+@pytest.mark.skipif(not IS_HIP_EXTENSION, reason="Covers the ROCm blockwise 4over6 path")
 @pytest.mark.skipif(not recipe_available, reason=reason_for_no_recipe)
 @pytest.mark.parametrize("M, N", [(128, 128), (256, 256), (512, 1024), (320, 256)])
 @pytest.mark.parametrize(
diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
index 62a7f61e52..90da4d217c 100644
--- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
@@ -456,25 +456,35 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
   const int64_t *const first_dims_ptr = reinterpret_cast<const int64_t *>(input->first_dims.dptr);
   const int64_t *const last_dims_ptr = reinterpret_cast<const int64_t *>(input->last_dims.dptr);
 
+  // Per-tensor dims are device-resident, so reading them needs a host sync that capture forbids.
+  auto read_dims_to_host = [&](const int64_t *dev_dims) {
+    cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
+    NVTE_CHECK_CUDA(cudaStreamIsCapturing(stream, &capture_status));
+    NVTE_CHECK(capture_status == cudaStreamCaptureStatusNone,
+               "Grouped dequantize cannot validate per-tensor dimensions during graph capture.");
+    std::vector<int64_t> host_dims(num_tensors);
+    NVTE_CHECK_CUDA(cudaMemcpyAsync(host_dims.data(), dev_dims, num_tensors * sizeof(int64_t),
+                                    cudaMemcpyDeviceToHost, stream));
+    NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+    return host_dims;
+  };
+
   // Only the uniform-last-dim path indexes colwise scales from the group's global row.
   if (use_colwise_scaling && last_dims_ptr == nullptr && num_tensors > 1) {
     constexpr size_t COLWISE_SCALE_ROWS = 32;
-    std::vector<int64_t> host_first_dims(num_tensors);
+    std::vector<int64_t> first_dims;
     if (first_dims_ptr != nullptr) {
-      NVTE_CHECK_CUDA(cudaMemcpyAsync(host_first_dims.data(), first_dims_ptr,
-                                      num_tensors * sizeof(int64_t), cudaMemcpyDeviceToHost,
-                                      stream));
-      NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+      first_dims = read_dims_to_host(first_dims_ptr);
     } else {
       NVTE_CHECK(first_logical_dim % num_tensors == 0, "Uniform grouped tensor has ",
                  first_logical_dim, " rows, which is not divisible by ", num_tensors, " tensors.");
-      host_first_dims.assign(num_tensors, static_cast<int64_t>(first_logical_dim / num_tensors));
+      first_dims.assign(num_tensors, static_cast<int64_t>(first_logical_dim / num_tensors));
     }
     for (size_t t = 0; t < num_tensors; t++) {
-      NVTE_CHECK(static_cast<size_t>(host_first_dims[t]) % COLWISE_SCALE_ROWS == 0,
+      NVTE_CHECK(static_cast<size_t>(first_dims[t]) % COLWISE_SCALE_ROWS == 0,
                  "Columnwise grouped dequantize requires each per-tensor first dimension to be "
                  "divisible by ",
-                 COLWISE_SCALE_ROWS, "; tensor ", t, " has ", host_first_dims[t], ".");
+                 COLWISE_SCALE_ROWS, "; tensor ", t, " has ", first_dims[t], ".");
     }
   }
 
@@ -485,14 +495,11 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
                "Grouped dequantize requires the last dimension to be divisible by ", VECTOR_ELEMS,
                "; got ", last_logical_dim, ".");
   } else {
-    std::vector<int64_t> host_last_dims(num_tensors);
-    NVTE_CHECK_CUDA(cudaMemcpyAsync(host_last_dims.data(), last_dims_ptr,
-                                    num_tensors * sizeof(int64_t), cudaMemcpyDeviceToHost, stream));
-    NVTE_CHECK_CUDA(cudaStreamSynchronize(stream));
+    const std::vector<int64_t> last_dims = read_dims_to_host(last_dims_ptr);
     for (size_t t = 0; t < num_tensors; t++) {
-      NVTE_CHECK(static_cast<size_t>(host_last_dims[t]) % VECTOR_ELEMS == 0,
+      NVTE_CHECK(static_cast<size_t>(last_dims[t]) % VECTOR_ELEMS == 0,
                  "Grouped dequantize requires each per-tensor last dimension to be divisible by ",
-                 VECTOR_ELEMS, "; tensor ", t, " has ", host_last_dims[t], ".");
+                 VECTOR_ELEMS, "; tensor ", t, " has ", last_dims[t], ".");
     }
   }
 
diff --git a/transformer_engine/common/fused_attn/flash_attn.cu b/transformer_engine/common/fused_attn/flash_attn.cu
index 296c4b0add..2fc02431aa 100644
--- a/transformer_engine/common/fused_attn/flash_attn.cu
+++ b/transformer_engine/common/fused_attn/flash_attn.cu
@@ -611,10 +611,6 @@ void multi_tensor_transpose_to_bhsd(Tensor *inputs, Tensor *outputs, size_t num_
   size_t d_in_max = 0, d_out_max = 0;
   bool any_not_vec_aligned = false;
 #ifdef __HIP_PLATFORM_AMD__
-  // Both readers below are compiled out on ROCm, so this only keeps the variable
-  // defined and states the intent: had they been reachable, the sm >= 100 test
-  // would not have excluded us -- sm_arch() returns the raw gfx number, so
-  // gfx1250 reports 125 and satisfies it.
   bool all_tma_ok = false;
 #else
   bool all_tma_ok = true;
diff --git a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
index 086419a65d..4b27825266 100644
--- a/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
+++ b/transformer_engine/common/gemm/kittens/cdna3/blockwise_fp8_gemm_helper.cuh
@@ -108,8 +108,12 @@ template <typename OType>
 __device__ inline float round_to_out_dtype(float v) {
     if constexpr (std::is_same_v<OType, float>) {
         return v;
-    } else {
+    } else if constexpr (std::is_same_v<OType, kittens::bf16>) {
         return kittens::base_types::convertor<float, OType>::convert(convert_out<OType>(v));
+    } else {
+        // fp16 is not parametrized in test_float8_blockwise_gemm_exact.py; left unchanged.
+        return static_cast<float>(
+            kittens::base_types::convertor<OType, float>::convert(rtne_bias(v)));
     }
 }
 
diff --git a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
index bc9be22847..2a5ef60824 100644
--- a/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
+++ b/transformer_engine/common/gemm/kittens/cdna4/blockwise_fp8_gemm_helper.cuh
@@ -95,9 +95,8 @@ __device__ inline float round_to_out_dtype(float v) {
         asm volatile("" : "+v"(bits));
         return __bfloat162float(__builtin_bit_cast(kittens::bf16, bits));
     } else {
-        uint16_t bits = __builtin_bit_cast(uint16_t, __float2half(v));
-        asm volatile("" : "+v"(bits));
-        return __half2float(__builtin_bit_cast(__half, bits));
+        // fp16 is not parametrized in test_float8_blockwise_gemm_exact.py; left unchanged.
+        return __half2float(__float2half(v));
     }
 }
 
diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
index a8dadf228e..ccc9d4a82d 100644
--- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
+++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
@@ -27,6 +27,19 @@
 #endif
 
 namespace transformer_engine {
+// The 4over6 body exists only on ROCm; expanding both arms on CUDA would instantiate every
+// specialization twice for a branch it cannot take.
+#ifdef __HIP_PLATFORM_AMD__
+#define NVTE_SWITCH_4OVER6(CONDITION, FLAG, ...) \
+  TRANSFORMER_ENGINE_SWITCH_CONDITION(CONDITION, FLAG, __VA_ARGS__)
+#else
+#define NVTE_SWITCH_4OVER6(CONDITION, FLAG, ...) \
+  {                                              \
+    constexpr bool FLAG = false;                 \
+    { __VA_ARGS__ }                              \
+  }
+#endif
+
 
 #if defined(__HIP_PLATFORM_AMD__) || CUDA_VERSION >= 12080
 namespace quantize_transpose_nvfp4 {
@@ -192,14 +205,9 @@ static __device__ constexpr unsigned int WARP_REDUCE_AMAX_GROUP_MASKS[8] = {
     0x01010101, 0x02020202, 0x04040404, 0x08080808, 0x10101010, 0x20202020, 0x40404040, 0x80808080};
 
 #ifdef __HIP_PLATFORM_AMD__
-// Sum the 16 row errors of a 2D block in upstream's tree order: pair rows 8
-// apart, then 4, then 2, then 1 (quantize_4over6_nvfp4.cuh::reduce_group_sum_16).
-// fp32 addition is not associative and 4over6 selection compares two sums that
-// are frequently equal to within an ulp, so the association is part of the
-// contract rather than an implementation detail. Upstream can run this tree in
-// warp shuffles because one warp there owns 16 consecutive rows; here a warp
-// owns only 4, so the rows paired first sit in different warps and the whole
-// set has to be staged in shared memory and reduced by one thread.
+// Upstream's tree order: pair rows 8 apart, then 4, 2, 1
+// (quantize_4over6_nvfp4.cuh::reduce_group_sum_16). Staged in shared memory because a
+// ROCm warp owns 4 rows, not 16.
 template <int kBlockSize, typename CType>
 __device__ __forceinline__ float reduceBlockError(const CType* staged) {
   float v[kBlockSize];
@@ -248,20 +256,15 @@ __device__ __forceinline__ float ComputeEncodeScaleFP4(ScaleType decode_scale,
 }
 
 #ifdef __HIP_PLATFORM_AMD__
-// Round-trip one value through FP4 at a given scale and return its contribution
-// to the block error.
-// E2M1 holds 16 values, so decode by table rather than through the packed-pair
-// conversion intrinsic, whose element ordering is easy to get wrong here.
+// Round-trip one value through FP4 at a given scale, returning its block-error term.
 __device__ __forceinline__ float fp4_decode_e2m1(const __hip_fp4_storage_t q) {
   constexpr float kMagnitude[8] = {0.0f, 0.5f, 1.0f, 1.5f, 2.0f, 3.0f, 4.0f, 6.0f};
   const float m = kMagnitude[q & 0x7];
   return (q & 0x8) ? -m : m;
 }
 
-// Rounding is pinned with the _rn intrinsics and the operand order upstream uses
-// (quantize_4over6_nvfp4.cuh:199). Selection turns on comparing two sums that are
-// often close, so an FMA contraction here flips near-ties and the result stops
-// being bit-comparable with the reference.
+// _rn intrinsics and upstream's operand order (quantize_4over6_nvfp4.cuh:199): an FMA
+// contraction here flips near-ties.
 __device__ __forceinline__ float fp4_roundtrip_err(const float x, const float block_scale_inverse,
                                                    const float sf, const float global_amax,
                                                    const float err_denom, const bool use_mse) {
@@ -273,16 +276,8 @@ __device__ __forceinline__ float fp4_roundtrip_err(const float x, const float bl
   return use_mse ? __fmul_rn(diff, diff) : fabsf(diff);
 }
 
-// 4over6: a block scale expanded by 1.5x maps the block max onto FP4 code 4
-// rather than code 6, spacing the codes more finely for blocks whose values sit
-// away from the top of the range. Score both encodings of the block and keep the
-// one that reconstructs it with less error.
-//
-// Takes the block as a plain array: the rowwise and columnwise call sites walk
-// shared memory differently, so a shared loop cannot serve both.
-//
-// Split in two so a caller that has already reduced the block error can rebuild the
-// two scales without walking the block again.
+// 4over6: score a block scale expanded by 1.5x (block max onto FP4 code 4) against the
+// plain one, and keep whichever reconstructs the block with less error.
 template <typename ScaleType>
 __device__ __forceinline__ void Compute4Over6Scales(const float amax,
                                                     const float global_decode_scale,
@@ -595,10 +590,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
       kIs2DBlockScaling ? (kFP4BlockScalingSize / kNumRowsPerWarp) : 1;
   __shared__ CType amax_smem_red[k2DBlockAmaxDim][k2DBlockAmaxDim][k2DBlockAmaxReduceDim];
   __shared__ CType amax_smem[k2DBlockAmaxDim][k2DBlockAmaxDim];
-  // 4over6 scores its two candidates over the same block the scale serves, so the
-  // error is reduced over the same 16x16 block the 2D amax is. Unlike the amax it
-  // stages all 16 rows: max is associative, the fp32 sum is not, so the rows cannot
-  // be pre-combined per warp without diverging from upstream's tree.
+  // Unlike the amax, all 16 rows are staged: the fp32 sum is not associative.
   constexpr int k4Over6Dim = kUse4Over6 ? k2DBlockAmaxDim : 1;
   constexpr int k4Over6RowDim = kUse4Over6 ? kFP4BlockScalingSize : 1;
   __shared__ CType err_smem_red[2][k4Over6Dim][k4Over6Dim][k4Over6RowDim];
@@ -959,16 +951,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock) block_scaled_1d_cast_transpo
     const unsigned mask = ((1 << kNumThreadsReduce) - 1) << src_lane;
     const bool is_src_lane = (threadIdx.x % kNumThreadsReduce) == 0;
 #ifdef __HIP_PLATFORM_AMD__
-    // Step 3.0: 4over6 block errors for the whole tile.
-    // Upstream scores the columnwise candidates on their own column groups
-    // (quantize_4over6_nvfp4.cuh::quantize_stage_colwise) instead of reusing what the
-    // rowwise pass computed. Both passes cover the same 16x16 block but sum its errors
-    // along different axes, and fp32 addition is not associative, so the two can land on
-    // different candidates for one block; that is the algorithm, not a bug.
-    //
-    // Done here rather than inside the loop below because a block's 16 columns are split
-    // across both smem_idx halves, so no single pass of that loop sees a whole block.
-    // Keeping it separate also leaves nothing live across the barriers.
+    // Step 3.0: 4over6 block errors for the whole tile. Scored on their own column groups
+    // per upstream (quantize_4over6_nvfp4.cuh::quantize_stage_colwise); kept out of the loop
+    // below because a block's 16 columns span both smem_idx halves.
     if constexpr (kUse4Over6 && kIs2DBlockScaling) {
       const int row_blk = threadIdx.x % kNumThreadsStore;
       int c = c_s;
@@ -1276,16 +1261,8 @@ void quantize_transpose_vector_blockwise_fp4(
                                   TRANSFORMER_ENGINE_SWITCH_CONDITION(
                                       row_scaled_nvfp4, kRowScaledNVFP4,
 
-  // 4over6 only has a body on ROCm, so pin the switch
-  // false on CUDA rather than instantiate every
-  // specialization twice for a branch it cannot take.
-#ifdef __HIP_PLATFORM_AMD__
-                                      TRANSFORMER_ENGINE_SWITCH_CONDITION(
+                                      NVTE_SWITCH_4OVER6(
                                           nvfp4_4over6_mode != kNVTENVFP44Over6Disabled, kUse4Over6,
-#else
-                                      TRANSFORMER_ENGINE_SWITCH_CONDITION(
-                                          false, kUse4Over6,
-#endif
 
 #ifdef __HIP_PLATFORM_AMD__
                                           size_t smem_bytes =
diff --git a/transformer_engine/jax/csrc/extensions/topk.cpp b/transformer_engine/jax/csrc/extensions/topk.cpp
index 450ff08b38..3322e6b5c4 100644
--- a/transformer_engine/jax/csrc/extensions/topk.cpp
+++ b/transformer_engine/jax/csrc/extensions/topk.cpp
@@ -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.
diff --git a/transformer_engine/pytorch/csrc/common.cpp b/transformer_engine/pytorch/csrc/common.cpp
index c6b312434c..2f5bd9b16c 100644
--- a/transformer_engine/pytorch/csrc/common.cpp
+++ b/transformer_engine/pytorch/csrc/common.cpp
@@ -373,13 +373,11 @@ at::Tensor allocate_amax_workspace(const TensorWrapper& input_tensor) {
     return at::empty(0, at::CUDA(at::kFloat));
   }
 
-  // A per-call tensor returns to the caching allocator at the caller's closing brace while
-  // amax_kernel and amax_final_reduce are still only enqueued. Stream-ordered reuse covers
-  // that in eager mode, but graph capture replaces stream order with a DAG, so the freed
-  // block is aliased in the graph mempool and replays read partials they never wrote.
-  // One buffer per (device, stream), sized to the block cap: same stream serializes.
-  static std::mutex amax_ws_mutex;
-  static std::map<std::pair<int, cudaStream_t>, at::Tensor> amax_ws_cache;
+  // Persistent: a per-call tensor is reclaimed while its kernels are still enqueued, which
+  // graph capture turns into an aliased read. One buffer per (device, stream).
+  // Leaked deliberately: these hold device memory and must not be destroyed at static teardown.
+  static std::mutex &amax_ws_mutex = *new std::mutex;
+  static auto &amax_ws_cache = *new std::map<std::pair<int, cudaStream_t>, at::Tensor>;
 
   const int device_id = at::cuda::current_device();
   const cudaStream_t stream = at::cuda::getCurrentCUDAStream();

From 01214dfcc169707b934b39967cf7b3b8be18a617 Mon Sep 17 00:00:00 2001
From: AllenFarcas <alfarcas@amd.com>
Date: Wed, 12 Aug 2026 12:54:40 +0000
Subject: [PATCH 4/5] Address second round of PR review comments

Hoist the shared prologue of group_dequantize out of the platform guard.
The ROCm arm was added by wrapping the whole function, which duplicated the
validation checks, shape derivation and pointer casts into both arms. Only
the divergent code stays guarded now: checkCuDriverContext and the TMA
dispatch on CUDA, the empty-input return, 2D grid and dimension validation
on ROCm.

Report the grouped dequantize mismatch once after locating it, rather than
inside the search loop. The locator scans bytes so it always finds the
offset memcmp reported; comparing elements could match on -0.0 vs +0.0 or
NaN payloads and fall through, which is why the redundant ASSERT_EQ was
there. It is no longer reachable.

Restore the colwise divisibility comment on the CUDA arm and split the
kUse4Over6 and kRowScaledNVFP4 switch terminators onto separate lines.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 .../operator/test_dequantize_mxfp8_grouped.cu | 22 +++---
 .../cast/mxfp8/group_dequantize_mxfp8.cuh     | 78 ++++---------------
 ...quantize_transpose_vector_blockwise_fp4.cu |  5 +-
 3 files changed, 33 insertions(+), 72 deletions(-)

diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
index bc42ed2cc5..a716569539 100644
--- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
+++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
@@ -347,17 +347,20 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
     int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset,
                         tensor_elts * sizeof(OutputType));
     if (result != 0) {
-      for (size_t i = 0; i < tensor_elts; ++i) {
-        if (out_grouped_h[data_offset + i] != out_ref_h[data_offset + i]) {
-          GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i
-                       << " (global offset " << (data_offset + i) << "): grouped="
-                       << static_cast<float>(out_grouped_h[data_offset + i])
-                       << " vs reference=" << static_cast<float>(out_ref_h[data_offset + i]);
-        }
+      const auto *grouped_bytes =
+          reinterpret_cast<const unsigned char *>(out_grouped_h.data() + data_offset);
+      const auto *ref_bytes =
+          reinterpret_cast<const unsigned char *>(out_ref_h.data() + data_offset);
+      size_t byte = 0;
+      while (grouped_bytes[byte] == ref_bytes[byte]) {
+        ++byte;
       }
+      const size_t i = byte / sizeof(OutputType);
+      GTEST_FAIL() << "Bitwise mismatch at tensor " << t << " element " << i << " (global offset "
+                   << (data_offset + i)
+                   << "): grouped=" << static_cast<float>(out_grouped_h[data_offset + i])
+                   << " vs reference=" << static_cast<float>(out_ref_h[data_offset + i]);
     }
-    // memcmp is the assertion; != would let -0.0 vs +0.0 and NaN payloads through.
-    ASSERT_EQ(result, 0) << "Bitwise mismatch in tensor " << t;
   }
 
   // Cleanup
@@ -493,6 +496,7 @@ TEST_P(GroupedDequantizeMXFP8TestSuite, TestGroupedDequantizeMXFP8) {
     if (last_dims[t] % 16 != 0) {
       GTEST_SKIP();
     }
+    // For colwise: first dim must be divisible by 32
     if (!rowwise && (first_dims[t] % 32 != 0)) {
       GTEST_SKIP();
     }
diff --git a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
index 90da4d217c..d87a3207bc 100644
--- a/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
+++ b/transformer_engine/common/cast/mxfp8/group_dequantize_mxfp8.cuh
@@ -395,9 +395,12 @@ __global__ void __launch_bounds__(128)
 
 inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
                              cudaStream_t stream) {
-#ifdef __HIP_PLATFORM_AMD__
   using namespace group_dequantize_kernel;
 
+#ifndef __HIP_PLATFORM_AMD__
+  checkCuDriverContext(stream);
+#endif
+
   const bool use_rowwise_scaling = input->has_data();
   const bool use_colwise_scaling = input->has_columnwise_data();
   NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling,
@@ -430,10 +433,22 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
   const size_t first_logical_dim = input->logical_shape.data[0];
   const size_t last_logical_dim = input->logical_shape.data[1];
   const size_t elts_total = first_logical_dim * last_logical_dim;
-  if (elts_total == 0) return;
 
   const size_t num_tensors = input->num_tensors;
 
+  const int64_t *const offsets_ptr = reinterpret_cast<const int64_t *>(input->tensor_offsets.dptr);
+  const int64_t *const first_dims_ptr = reinterpret_cast<const int64_t *>(input->first_dims.dptr);
+  const int64_t *const last_dims_ptr = reinterpret_cast<const int64_t *>(input->last_dims.dptr);
+
+  const e8m0_t *const scales_ptr =
+      use_rowwise_scaling ? reinterpret_cast<e8m0_t *>(input->scale_inv.dptr)
+                          : reinterpret_cast<e8m0_t *>(input->columnwise_scale_inv.dptr);
+
+  const SimpleTensor &input_data = use_rowwise_scaling ? input->data : input->columnwise_data;
+
+#ifdef __HIP_PLATFORM_AMD__
+  if (elts_total == 0) return;
+
   size_t blocks_X = 0;
   size_t blocks_Y = 1;
   if (is_single_tensor) {
@@ -452,10 +467,6 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
   const dim3 grid(blocks_X, blocks_Y);
   const dim3 block(THREADS_PER_CHUNK);
 
-  const int64_t *const offsets_ptr = reinterpret_cast<const int64_t *>(input->tensor_offsets.dptr);
-  const int64_t *const first_dims_ptr = reinterpret_cast<const int64_t *>(input->first_dims.dptr);
-  const int64_t *const last_dims_ptr = reinterpret_cast<const int64_t *>(input->last_dims.dptr);
-
   // Per-tensor dims are device-resident, so reading them needs a host sync that capture forbids.
   auto read_dims_to_host = [&](const int64_t *dev_dims) {
     cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
@@ -503,12 +514,6 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
     }
   }
 
-  const e8m0_t *const scales_ptr =
-      use_rowwise_scaling ? reinterpret_cast<e8m0_t *>(input->scale_inv.dptr)
-                          : reinterpret_cast<e8m0_t *>(input->columnwise_scale_inv.dptr);
-
-  const SimpleTensor &input_data = use_rowwise_scaling ? input->data : input->columnwise_data;
-
   const size_t scale_dim_X_rowwise = use_rowwise_scaling ? 32 : 1;
   const size_t scale_dim_Y_colwise = use_colwise_scaling ? 32 : 1;
 
@@ -536,45 +541,6 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
   );                                                         // NOLINT(*)
   NVTE_CHECK_CUDA(cudaGetLastError());
 #else
-  using namespace group_dequantize_kernel;
-
-  checkCuDriverContext(stream);
-
-  const bool use_rowwise_scaling = input->has_data();
-  const bool use_colwise_scaling = input->has_columnwise_data();
-  NVTE_CHECK(use_rowwise_scaling || use_colwise_scaling,
-             "Input tensor must have either rowwise or columnwise data.");
-  NVTE_CHECK(!(use_rowwise_scaling && use_colwise_scaling),
-             "Dequantize only supports rowwise or columnwise scaling, not both simultaneously.");
-
-  NVTE_CHECK(!input->with_gemm_swizzled_scales, "Input must have scales in compact format.");
-  NVTE_CHECK(!is_fp8_dtype(output->dtype()), "Output must be in higher precision.");
-  NVTE_CHECK(!is_fp4_dtype(output->dtype()), "Output must not be FP4.");
-  NVTE_CHECK(is_fp8_dtype(input->dtype()), "Input must have FP8 type.");
-
-  NVTE_CHECK(input->num_tensors == output->num_tensors,
-             "Number of input and output tensors must be same.");
-
-  ShapeRepresentation shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;
-  if (input->all_same_shape()) {
-    shape_rep = ShapeRepresentation::SAME_BOTH_DIMS;
-  } else if (input->all_same_first_dim()) {
-    shape_rep = ShapeRepresentation::VARYING_LAST_DIM;
-  } else if (input->all_same_last_dim()) {
-    shape_rep = ShapeRepresentation::VARYING_FIRST_DIM;
-  } else if (input->varying_both_dims()) {
-    shape_rep = ShapeRepresentation::VARYING_BOTH_DIMS;
-  }
-
-  const bool is_single_tensor = (shape_rep == ShapeRepresentation::SAME_BOTH_DIMS ||
-                                 shape_rep == ShapeRepresentation::VARYING_FIRST_DIM);
-
-  const size_t first_logical_dim = input->logical_shape.data[0];
-  const size_t last_logical_dim = input->logical_shape.data[1];
-  const size_t elts_total = first_logical_dim * last_logical_dim;
-
-  const size_t num_tensors = input->num_tensors;
-
   constexpr size_t CHUNK_DIM_Y = DequantizeConfig::CHUNK_DIM_Y;
   constexpr size_t CHUNK_DIM_X = DequantizeConfig::CHUNK_DIM_X;
   constexpr size_t THREADS_PER_CHUNK = DequantizeConfig::THREADS_PER_CHUNK;
@@ -602,16 +568,6 @@ inline void group_dequantize(const GroupedTensor *input, GroupedTensor *output,
   const dim3 grid(blocks);
   const dim3 block(THREADS_PER_CHUNK);
 
-  const int64_t *const offsets_ptr = reinterpret_cast<const int64_t *>(input->tensor_offsets.dptr);
-  const int64_t *const first_dims_ptr = reinterpret_cast<const int64_t *>(input->first_dims.dptr);
-  const int64_t *const last_dims_ptr = reinterpret_cast<const int64_t *>(input->last_dims.dptr);
-
-  const e8m0_t *const scales_ptr =
-      use_rowwise_scaling ? reinterpret_cast<e8m0_t *>(input->scale_inv.dptr)
-                          : reinterpret_cast<e8m0_t *>(input->columnwise_scale_inv.dptr);
-
-  const SimpleTensor &input_data = use_rowwise_scaling ? input->data : input->columnwise_data;
-
   TRANSFORMER_ENGINE_TYPE_SWITCH_FP8ONLY(
       input->dtype(), IType,
       TRANSFORMER_ENGINE_TYPE_SWITCH_NON_FP8ONLY(
diff --git a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
index ccc9d4a82d..a85bc1f10f 100644
--- a/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
+++ b/transformer_engine/common/transpose/quantize_transpose_vector_blockwise_fp4.cu
@@ -1298,8 +1298,9 @@ void quantize_transpose_vector_blockwise_fp4(
                                               scale_t_stride_x, scale_t_stride_y, kScaleBlockDim,
                                               epsilon, rng_state, noop_ptr, e4m3_max_for_arch,
                                               nvfp4_4over6_mode ==
-                                                  kNVTENVFP44Over6MinMSE);))  // kUse4Over6 / kRowScaledNVFP4
-                                  )                                           // kIs2DBlockScaling
+                                                  kNVTENVFP44Over6MinMSE);)  // kUse4Over6
+                                      )                                      // kRowScaledNVFP4
+                                  )                                          // kIs2DBlockScaling
                               )  // kApplyStochasticRounding
                           )      // kSwizzledScale
                       )          // kAligned

From 8b5318bb70b75bbb30af1043c67283e0cbc9bcda Mon Sep 17 00:00:00 2001
From: AllenFarcas <alfarcas@amd.com>
Date: Wed, 12 Aug 2026 14:19:37 +0000
Subject: [PATCH 5/5] Mark the grouped dequantize mismatch locator as
 fork-local

Restores the upstream comment above the search loop, whose code survives the
rewrite, and records why the loop scans bytes.

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>
---
 tests/cpp/operator/test_dequantize_mxfp8_grouped.cu | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
index a716569539..c5e420cfbe 100644
--- a/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
+++ b/tests/cpp/operator/test_dequantize_mxfp8_grouped.cu
@@ -347,6 +347,8 @@ void performTest(const ShapeRepresentation shape_rep, const size_t num_tensors,
     int result = memcmp(out_grouped_h.data() + data_offset, out_ref_h.data() + data_offset,
                         tensor_elts * sizeof(OutputType));
     if (result != 0) {
+      // Find first mismatch for error reporting
+      // ROCm scans bytes, not elements: memcmp flags differences element comparison misses.
       const auto *grouped_bytes =
           reinterpret_cast<const unsigned char *>(out_grouped_h.data() + data_offset);
       const auto *ref_bytes =