// Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
// SPDX-License-Identifier: MIT
#pragma once
// Auto-generated — do not edit manually.
// Lightweight LN headers (~3K lines) replacing the 151K-line embedded version.
// Uses #include directives for CUDA TK headers (requires --include-path at NVRTC compile time).

#if defined(__clang__)
#pragma clang diagnostic push
#pragma clang diagnostic ignored "-Woverlength-strings"
#endif

#if defined(_MSC_VER)
#pragma warning(push)
#pragma warning(disable : 4068)
#endif

namespace cudnn_frontend {
namespace experimental {
namespace generated {

inline constexpr const char ln_headers_source[] =
    R"KERNEL(
#pragma once
// Auto-generated lightweight LN kernel header.
// Replaces the 151K-line persistent_ln_headers_13.0.h with:
//   - Standard CUDA TK #include directives (resolved via --include-path at NVRTC compile time)
//   - ~2900 lines of cuDNN-authored LN-specific code extracted from the original
//
// CGA/cluster support (USE_CLUSTER) is disabled:
//   - #include <cuda/ptx> and cooperative_groups are NOT included
//   - All if constexpr (USE_CLUSTER) branches changed to if constexpr(USE_CLUSTER) so dead code is not compiled
//
// Requires CUDA Toolkit headers available at the --include-path location.
// Generated by step3b_extract_ln_headers.py — do not edit manually.

// ============================================================
// Standard CUDA Toolkit headers (resolved via --include-path)
// ============================================================

// Minimal std stubs needed by NVRTC (no standard library available)
#if __cplusplus >= 201103L
namespace std {
template<bool B, class T = void> struct enable_if {};
template<class T>                struct enable_if<true, T> { typedef T type; };
#if __cplusplus >= 201402L
template< bool B, class T = void > using enable_if_t = typename enable_if<B,T>::type;
#endif

struct true_type  { enum { value = true };  operator bool() const { return true; } };
struct false_type { enum { value = false }; operator bool() const { return false; } };

template<typename T> struct is_floating_point    : false_type {};
template<> struct is_floating_point<float>       :  true_type {};
template<> struct is_floating_point<double>      :  true_type {};

template<typename T> struct is_integral           : false_type {};
template<> struct is_integral<bool>               :  true_type {};
template<> struct is_integral<char>               :  true_type {};
template<> struct is_integral<short>              :  true_type {};
template<> struct is_integral<int>                :  true_type {};
template<> struct is_integral<long>               :  true_type {};
template<> struct is_integral<long long>          :  true_type {};
template<> struct is_integral<unsigned char>      :  true_type {};
template<> struct is_integral<unsigned short>     :  true_type {};
template<> struct is_integral<unsigned int>       :  true_type {};
template<> struct is_integral<unsigned long>      :  true_type {};
template<> struct is_integral<unsigned long long> :  true_type {};

template<typename T, typename U> struct is_same       : false_type {};
template<typename T>             struct is_same<T, T> :  true_type {};
}
#endif

// CUDA numeric type headers
#include <cuda_fp16.h>
#include <cuda_bf16.h>
#include <cuda_fp8.h>
#include <cuda_fp4.h>

// NOTE: #include <cuda/ptx> and <cooperative_groups.h> are intentionally omitted.
// They are only needed for CGA/cluster support (USE_CLUSTER=true), which this
// OSS engine does not use. All USE_CLUSTER code paths are guarded with
// "if constexpr" so dead branches are not compiled.

// ============================================================
// Fixed-width integer typedefs for NVRTC (no <cstdint> available)
// ============================================================
constexpr int THREADS_PER_WARP = 32;

#ifndef __CUDACC_RTC__
#include <stdint.h>
#else
typedef signed char int8_t;
typedef unsigned char uint8_t;
typedef signed short int16_t;
typedef unsigned short uint16_t;
typedef signed int int32_t;
typedef unsigned int uint32_t;
typedef signed long long int64_t;
typedef unsigned long long uint64_t;
#endif

// ============================================================
// cuDNN-authored LN kernel utilities
// Extracted from persistent_ln_headers_13.0.h
// ============================================================


// cuDNN type aliases for FP8 types (from after inlined cuda_fp8.h)
typedef __nv_fp8_e4m3 nv_fp8_e4m3;
typedef __nv_fp8x2_e4m3 nv_fp8x2_e4m3;
typedef __nv_fp8_e5m2 nv_fp8_e5m2;
typedef __nv_fp8x2_e5m2 nv_fp8x2_e5m2;

// cuDNN type aliases for FP4 types (from after inlined cuda_fp4.h)
typedef __nv_fp4_e2m1 nv_fp4_e2m1;
typedef __nv_fp4x2_e2m1 nv_fp4x2_e2m1;
typedef __nv_fp4x4_e2m1 nv_fp4x4_e2m1;

// Pointer typedefs for NVRTC (from original Jitify stubs)
typedef int64_t intptr_t;
typedef uint64_t uintptr_t;
typedef int64_t intmax_t;
typedef uint64_t uintmax_t;

// Large vector types for 64-byte and 32-byte loads
struct uint16 {
    uint4 u;
    uint4 v;
    uint4 s;
    uint4 t;
};

struct uint8 {
    uint4 u;
    uint4 v;
};

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Datatype helpers

template <int BYTES>
struct BytesToType {};

template <>
struct BytesToType<64> {
    using Type = uint16;
};

template <>
struct BytesToType<32> {
    using Type = uint8;
};

template <>
struct BytesToType<16> {
    using Type = uint4;
};

template <>
struct BytesToType<8> {
    using Type = uint64_t;
};

template <>
struct BytesToType<4> {
    using Type = uint32_t;
};

template <>
struct BytesToType<2> {
    using Type = uint16_t;
};

template <>
struct BytesToType<1> {
    using Type = uint8_t;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename T>
struct TypeToVec2 {};

template <>
struct TypeToVec2<float> {
    using Type = float2;
};

template <>
struct TypeToVec2<half> {
    using Type = half2;
};

template <>
struct TypeToVec2<nv_bfloat16> {
    using Type = nv_bfloat162;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <int INDEX>
struct Get {
    template <typename T, typename R>
    static inline __device__ R
    of(const T& vec);
};

template <>
template <typename T, typename R>
inline __device__ R
Get<0>::of(const T& vec) {
    return vec.x;
}

template <>
template <typename T, typename R>
inline __device__ R
Get<1>::of(const T& vec) {
    return vec.y;
}

template <>
template <typename T, typename R>
inline __device__ R
Get<2>::of(const T& vec) {
    return vec.z;
}

template <>
template <typename T, typename R>
inline __device__ R
Get<3>::of(const T& vec) {
    return vec.w;
}

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename Src, typename Dst>
struct Converter {
    static inline __device__ Dst
    convert(const Src& from) {
        return Dst(from);
    }
};

template <>
struct Converter<float2, half2> {
    static inline __device__ half2
    convert(const float2& x) {
        return __float22half2_rn(x);
    }
};

template <>
struct Converter<float2, nv_bfloat162> {
    static inline __device__ nv_bfloat162
    convert(const float2& x) {
#if __CUDA_ARCH__ >= 800
        return __float22bfloat162_rn(x);
#else
        union {
            nv_bfloat162 raw;
            nv_bfloat16 x;
            nv_bfloat16 y;
        } tmp;
        tmp.x = __float2bfloat16_rn(x.x);
        tmp.y = __float2bfloat16_rn(x.y);
        return tmp.raw;
#endif
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename T>
struct Zeros {
    static inline __device__ T
    get() {
        return T(0.f);
    }
};

template <>
struct Zeros<float2> {
    static inline __device__ float2
    get() {
        return make_float2(0.f, 0.f);
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

__device__ __inline__ uint8_t
float_to_e8m0(float val) {
    // CL-15277: use rounding-up mode for float -> e8m0 conversion
#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000) &&                           \
    (defined(__CUDA_ARCH_FEAT_SM100_ALL) || defined(__CUDA_ARCH_FEAT_SM101_ALL) || defined(__CUDA_ARCH_FEAT_SM110_ALL) || \
     defined(__CUDA_ARCH_FEAT_SM120_ALL))
    uint16_t out;
    asm volatile(
        "{\n"
        "cvt.rp.satfinite.ue8m0x2.f32   %0, 0.0, %1;\n"
        "}"
        : "=h"(out)
        : "f"(val));
    return *reinterpret_cast<uint8_t*>(&out);
#else
#if CUDART_VERSION >= 12080
    // Use explicit round-up mode to match the SM_100+ hardware instruction
    return __nv_cvt_float_to_e8m0(val, __NV_SATFINITE, cudaRoundPosInf);
#else
    if (isnan(val) || isinf(val)) {
        return 0xFF;
    }
    uint32_t val_u32  = __float_as_uint(val);
    uint8_t exponent  = (val_u32 >> 23) & 0xFF;
    uint32_t mantissa = val_u32 & 0x7FFFFF;
    if ((mantissa > 0) && (exponent != 0xFE)) {  // exp can only be < 0xFE here
        exponent++;
    }
    return exponent;
#endif
#endif
}

__device__ __inline__ float
e8m0_to_float(uint8_t val) {
#if CUDART_VERSION >= 12080
    __nv_fp8_e8m0 e8m0;
    e8m0.__x = val;
    return static_cast<float>(e8m0);
#else
    return __uint_as_float(static_cast<uint32_t>(val) << 23);
#endif

}

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Vectorization data structures
template <typename ELT_TYPE, uint32_t NUM_ELT>
struct Vec {
    using Elt_type = ELT_TYPE;

    enum { BYTES = NUM_ELT * sizeof(Elt_type) };
    using Vec_type = typename BytesToType<BYTES>::Type;

    using Alias_type = union {
        Vec_type vec;
        Elt_type elt[NUM_ELT];
    };

    Alias_type data;

    template <typename S>
    inline __device__ void
    to(Vec<S, NUM_ELT>& other) {
#pragma unroll
        for (int it = 0; it < NUM_ELT; it++) {
            other.data.elt[it] = S(this->data.elt[it]);
        }
    }

    template <typename Op>
    inline __device__ void
    assign(const Op& op) {
#pragma unroll
        for (int it = 0; it < NUM_ELT; it++) {
            this->data.elt[it] = op(it);
        }
    }

    inline __device__ void
    load_from(const void* base_ptr, const size_t idx) {
        this->data.vec = static_cast<const Vec_type*>(base_ptr)[idx];
    }

    inline __device__ void
    store_to(void* base_ptr, const size_t idx) {
        static_cast<Vec_type*>(base_ptr)[idx] = this->data.vec;
    }

    inline __device__ void
    operator+=(const Vec<ELT_TYPE, NUM_ELT>& rhs) {
#pragma unroll
        for (int it = 0; it < NUM_ELT; it++) {
            this->data.elt[it] += rhs.data.elt[it];
        }
    }
};

inline __device__ float2
operator+(const float2& a, const float2& b) {
    return {a.x + b.x, a.y + b.y};
}

inline __device__ void
operator+=(float2& a, const float2& b) {
    a.x += b.x;
    a.y += b.y;
}

template <typename T>
struct Sum {
    inline __device__
    Sum() {}
    inline __device__ T
    operator()(const T& a, const T& b) {
        return a + b;
    }
};

////////////////////////////////////////////////////////////////////////////////////////////////////

inline __device__ void
__nv_ptx_builtin_ocg_fence_view_async_shared(void) {
    asm volatile("fence.proxy.async.shared::cta;\n");
}

extern "C" {
__device__ uint32_t
__nvvm_get_smem_pointer(void* ptr);
//__device__ void
//__nv_ptx_builtin_ocg_fence_view_async_shared(void);
__device__ void
__nv_ptx_builtin_ocg_write_async_shared_b32(uint32_t dstAddr, uint32_t mbarrierAddr, uint32_t b0);
__device__ void
__nv_ptx_builtin_ocg_write_async_shared_b64(uint32_t dstAddr, uint32_t mbarrierAddr, uint64_t b0);
__device__ void
__nv_ptx_builtin_ocg_write_async_shared_v2_b32(uint32_t dstAddr, uint32_t mbarrierAddr, uint32_t b0, uint32_t b1);
__device__ void
__nv_ptx_builtin_ocg_write_async_shared_v4_b32(uint32_t dstAddr,
                                               uint32_t mbarrierAddr,
                                               uint32_t b0,
                                               uint32_t b1,
                                               uint32_t b2,
                                               uint32_t b3);
}

namespace utils {

//////////////////////////////////////////////////////////////////////////////////////////////////
//// Shuffle syncronization helpers
template <typename T>
inline __device__ T
warp_shuffle_xor(const T& x, uint32_t idx) {
    return __shfl_xor_sync(uint32_t(-1), x, idx);
}

template <>
inline __device__ float2
warp_shuffle_xor<float2>(const float2& x, uint32_t idx) {
    return {warp_shuffle_xor(x.x, idx), warp_shuffle_xor(x.y, idx)};)KERNEL"
    R"KERNEL(}

template <typename T>
inline __device__ T
warp_shuffle_down(const T& x, uint32_t idx) {
    return __shfl_down_sync(uint32_t(-1), x, idx);
}

template <>
inline __device__ float2
warp_shuffle_down<float2>(const float2& x, uint32_t idx) {
    return {warp_shuffle_down(x.x, idx), warp_shuffle_down(x.y, idx)};
}

// This is a helper function to support __shfl_sync(float2)
template <typename T>
__device__ inline static T
shfl_sync_helper(unsigned mask, T var, int srcLane) {
    return __shfl_sync(mask, var, srcLane);
}

// specialize for float2
template <>
__device__ inline float2
shfl_sync_helper<float2>(unsigned mask, float2 var, int srcLane) {
    double ret = __shfl_sync(mask, *reinterpret_cast<double*>(&var), srcLane);  // cast to double
    return *reinterpret_cast<float2*>(&ret);
}

//////////////////////////////////////////////////////////////////////////////////////////////////
__device__ __forceinline__ void
namedBarrierSync(int name, int numThreads) {
    asm volatile("bar.sync %0, %1;" : : "r"(name), "r"(numThreads) : "memory");
}

}  // namespace utils

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Atomics
static inline __device__ void
atomicMaxFloat(float* addr, const float value) {
    atomicMax(reinterpret_cast<int*>(addr), __float_as_int(value));
}

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Warp reduction functions
template <int GROUP_SIZE>
static inline __device__ float
warp_reduce_max(const float m, unsigned int mask = 0xFFFFFFFF) {
    static constexpr int WARP_SIZE = 32;
    static_assert(GROUP_SIZE > 0 && GROUP_SIZE <= WARP_SIZE, "group size must be less than or equal to warp size");

    float max_val = m;
#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000) && defined(__CUDA_ARCH_FEAT_SM100_ALL)
    // Only enable the credux instruction for reduction on the full warp, as we saw perf regression for GROUP_SIZE < 32,
    // where we need to iterate over #(WARP_SIZE/GROUP_SIZE) groups in the warp
    if constexpr (GROUP_SIZE == WARP_SIZE) {
        // For sm_100a arch, we can use redux.sync.op{.abs.}{.NaN}.f32 instruction to reduce the min/max value
        // https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-redux-sync
        asm volatile("redux.sync.max.f32 %0, %1, %2;" : "=f"(max_val) : "f"(m), "r"(mask));
        return max_val;
    }
#endif
#pragma unroll
    for (int delta = GROUP_SIZE / 2; delta > 0; delta /= 2) {
        const float other_m = __shfl_down_sync(mask, max_val, delta);
        __builtin_assume(max_val >= 0);
        __builtin_assume(other_m >= 0);
        max_val = fmaxf(max_val, other_m);
    }
    return max_val;
}

template <int num_warps, typename compute_t>
static inline __device__ compute_t
reduce_max(const compute_t m, const int warpid) {
    __shared__ float staging[num_warps];
    constexpr int warp_size = 32;
    const float my_max      = m;
    const float my_warp_max = warp_reduce_max<warp_size>(my_max);
    if (threadIdx.x % 32 == 0) {
        staging[warpid] = my_warp_max;
    }
    __syncthreads();
    compute_t result = 0;
    if (warpid == 0) {
        const float my_max = threadIdx.x < num_warps ? staging[threadIdx.x] : 0;
        result             = warp_reduce_max<warp_size>(my_max);
    }
    return result;
}

template <int num_warps, typename compute_t>
static inline __device__ compute_t
reduce_max(const compute_t m, const int warpid, int tidx) {
    __shared__ float staging[num_warps];
    constexpr int warp_size = 32;
    const float my_max      = m;
    const float my_warp_max = warp_reduce_max<warp_size>(my_max);
    if (tidx % 32 == 0) {
        staging[warpid] = my_warp_max;
    }
    utils::namedBarrierSync(1, num_warps * warp_size);
    compute_t result = 0;
    if (warpid == 0) {
        const float my_max = tidx < num_warps ? staging[tidx] : 0;
        result             = warp_reduce_max<warp_size>(my_max);
    }
    return result;
}

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Misc
constexpr __device__ int32_t
find_largest_divisor(int32_t n, int32_t divisor_upper_bound) {
    if (n <= divisor_upper_bound) {
        return n;
    }
    for (int32_t i = divisor_upper_bound; i >= 1; --i) {
        if (n % i == 0) {
            return i;  // found the largest divisor
        }
    }
    return divisor_upper_bound;  // fallback
}

constexpr __device__ bool
is_power_of_2(unsigned int n) {
    return (n > 0) && ((n & (n - 1)) == 0);
}

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Cooperative groups related
//// !!!! Keep this block at the end of the file - start !!!!

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Runtime parameters

struct PersistentLnParamsBase {
    // For Multi-CTA, number of different CTA groups; otherwise same as gridDim.x.
    int ctas_per_col = 0;

    // Input is interpreted as matrix and we normalize across columns.
    int rows      = 0;
    int cols      = 0;
    int batchSize = 1;  // For AdaLN
    int seqLen    = 1;  // For AdaLN

    // Common data pointers for forward and backward passes.
    void *x     = nullptr;
    void *mu    = nullptr;
    void *rs    = nullptr;
    void *gamma = nullptr;

    // Multi-CTA workspace in gmem.
    void *workspace = nullptr;

    // Multi-CTA sync barriers in gmem.
    int *barrier = nullptr;

    bool isRMSNorm    = false;
    bool noScale      = false;
    bool noBias       = false;
    bool isAdaLN      = false;
    bool isBatchFirst = true;
};

struct PersistentLnFwdParams : public PersistentLnParamsBase {
    // Output of LN FWD.
    void *z = nullptr;

    void *beta    = nullptr;
    float epsilon = 0.f;

    // FP8 support
    bool fp8_out     = false;
    float *scale     = nullptr;
    float *scale_inv = nullptr;
    float *amax      = nullptr;

    // Reuse z for row scaled output (or block scaled output if we use 2d1x1x)
    void *scale_row = nullptr;
    void *scale_col = nullptr;  // only used in 1d2x2x
    void *z_col     = nullptr;  // only used in 1d2x2x
    void *z_math    = nullptr;  // only used when enabling the 1d2x2x colwise kernel
};

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Inter-CTA synchronization
template <uint32_t CTAS_PER_ROW>
struct InterCTASync {
    template <typename Params>
    inline __device__
    InterCTASync(Params &params, uint32_t bidm, uint32_t bidn, uint32_t tidx)
        : phase_counter_(0),
          b0_(params.barrier + bidm)  // The barrier for this group of CTAs.
          ,
          b1_(params.barrier + bidm + params.ctas_per_col)  // The barrier for this group of CTAs.
          ,
          tidx_(tidx) {
        // BARRIERS ARE ASSUMED TO BE INITIALIZED TO 0!
    }

    inline __device__ void
    spin_wait_(int *barrier, int step, int expected) {
        asm volatile("red.release.gpu.global.add.s32 [%0], %1;" ::"l"(barrier), "r"(step));
        for (int found = -1; found != expected;) {
            asm volatile("ld.global.acquire.gpu.b32 %0, [%1];" : "=r"(found) : "l"(barrier));
        }
    }

    inline __device__ void
    sync(uint32_t threads_per_cta) {
        // ALL THREADS MUST ENTER!

        // We switch barrier every iteration.
        int *barrier = phase_counter_ & 0x1 ? b1_ : b0_;
        // We decrement every other iteration.
        bool dec     = phase_counter_ & 0x2;
        int step     = dec ? -1 : 1;
        int expected = dec ? 0 : CTAS_PER_ROW;
        // There are only 4 phases: up/down for b0/b1.
        phase_counter_ = (phase_counter_ + 1) & 0x3;

        if (0 == tidx_) {
            spin_wait_(barrier, step, expected);
        }

        // CTA waits for thread 0
        utils::namedBarrierSync(0, threads_per_cta);
    }

    int phase_counter_;
    int *b0_;
    int *b1_;
    uint32_t tidx_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////
//// LayerNorm Reduction and stats computation
// NOTE: WHOLE_CTA means if we'll use one empty mbarrier for inter-CTA sync using cluster.
// In the TMA LN engines it should be "true" since we're syncing the whole CTA with one empty mbarrier, while in the 
// APEX LN engines it should be "false" as there's one empty mbarrier per warp m (and therefore WARPS_M empty mbarriers
// in total)
template <typename T, uint32_t CTAS_PER_ROW, uint32_t WARPS_M, uint32_t WARPS_N, bool USE_CLUSTER, bool WHOLE_CTA>
struct Reducer : public Reducer<T, 1, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA> {
    using InterCTASync = InterCTASync<CTAS_PER_ROW>;
    using Base         = Reducer<T, 1, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA>;
    using Type         = typename Base::Type;

    enum { SMEM_BYTES = Base::SMEM_BYTES };

    enum { WS_BARRIER_BYTES = 2 * sizeof(int) };
    enum { WS_DATA_BYTES = WARPS_M * CTAS_PER_ROW * sizeof(T) };

    // size of the barriers + temporary result per CTA (multiply with CTAS_PER_ROW to get total)
    enum { WORKSPACE_BYTES_PER_GROUP = Base::WORKSPACE_BYTES_PER_GROUP + WS_BARRIER_BYTES + WS_DATA_BYTES };

    template <typename Params>
    inline __device__
    Reducer(Params &params,
            uint32_t bidm,
            uint32_t bidn,
            uint32_t warp_m,
            uint32_t warp_n,
            uint32_t tidx,
            uint32_t lane,
            void *smem,
            uint64_t *smem_bar)
        : Base(params, bidm, bidn, warp_m, warp_n, tidx, lane, smem, smem_bar),
          inter_cta_(params, bidm, bidn, tidx),
          bidn_(bidn),  // CTA id within the group.
          warp_m_(warp_m),
          w0_(static_cast<T *>(params.workspace) + (bidm * WARPS_M + warp_m) * CTAS_PER_ROW),
          w1_(w0_ + params.ctas_per_col * WARPS_M * CTAS_PER_ROW),
          smem_bar_(smem_bar),
          parity_(0) {
        // the first several elements are the barriers
        smem_cga_ = reinterpret_cast<T *>(&smem_bar[WHOLE_CTA? 2: (WARPS_M + 1)]);
    }

    template <typename Op>
    inline __device__ T
    allreduce(T data, Op &op, int32_t warps_m) {
        data = Base::reduce(data, op, warps_m);

        if constexpr (USE_CLUSTER) {
#if (__CUDA_ARCH__ >= 900) && (CUDART_VERSION >= 12080)
            auto cluster = cooperative_groups::this_cluster();

            // NOTE: right now only lane 0 has the valid value so that a intra-warp broadcast is required here.
            // This issue is not in fwd (struct Stats) which does it earlier inside warp_chan_upd_dynamic().
            data = utils::shfl_sync_helper(uint32_t(-1), data, 0);

            // Broadcast local results to other CTAs inside the CGA
            // Size of smem_cga_: [WARPS_M][CTAS_PER_ROW]
            // Thread 0 sends to block 0, Thread 1 sends to block 1, etc...
            if ((this->warp_n_ == 0) && (this->lane_ < CTAS_PER_ROW)) {
                st_async_remote(&smem_cga_[warp_m_ * CTAS_PER_ROW + bidn_], data, this->lane_, &smem_bar_[WHOLE_CTA? 0: WARPS_M]);
            }

            // Leader thread arrives on local barrier to indicate expected tx count
            // We can't use (threadIdx.x == 0) directly here because it won't work for TMA
            if ((this->warp_n_ == 0) && (warp_m_ == 0) && (this->lane_ == 0)) {
                // It's possible that not all warp_m are active here, therefore, we can't use WARPS_M directly
                uint32_t expected_tx_count = sizeof(data) * CTAS_PER_ROW * warps_m;
                utils::mbarrier_arrive_expect_tx_relaxed_cluster(&smem_bar_[WHOLE_CTA? 0: WARPS_M], expected_tx_count);
            }

            // Wait on local barrier
            while (!utils::mbarrier_try_wait_parity_relaxed_cluster(&smem_bar_[WHOLE_CTA? 0: WARPS_M], parity_)) {}
            utils::fence_acquire_smem();  // Ensure we can read local smem values that are released by barrier flip

            T total = Zeros<T>::get();
            if (this->lane_ < CTAS_PER_ROW) {
                total = smem_cga_[warp_m_ * CTAS_PER_ROW + this->lane_];
            }
            utils::fence_release_smem();  // Ensure read of buffer remains ordered before smem_bar.arrive.relaxed
)KERNEL"
    R"KERNEL(            // Signal barrier is empty.
            // Each thread must arrive on all the barriers (like an arrive broadcast).
            // NOTE: here we only let each warp m arrive on the remote barriers of the same warp m (i.e. the same row),
            // because if we have one empty barrier for the entire CTA, there's a difficulty how to init the count of
            // the barrier (since some warps could be disabled here). The downside is we'll need a __syncthreads() here,
            // which doesn't have obvious perf impact though.
            // TODO: we could probably have one full bar per row (same as the empty bar) to get rid of the CTA sync,
            // while the downside is more smem (mbarriers) will be needed.
            if (!WHOLE_CTA) {
                utils::namedBarrierSync(0, warps_m * WARPS_N * 32);
            }
            for (int other_block_rank = 0; other_block_rank < CTAS_PER_ROW; ++other_block_rank) {
                uint64_t *remote_bar_empty = cluster.map_shared_rank(&smem_bar_[WHOLE_CTA? 1: warp_m_], other_block_rank);
                utils::mbarrier_arrive_relaxed_cluster(remote_bar_empty);
            }

            total = Reducer<T, 1, 1, 1, USE_CLUSTER, WHOLE_CTA>::allreduce_(total, op);

            // Wait for remote buffer to be writable again by waiting on local barrier:
            while (!utils::mbarrier_try_wait_parity_relaxed_cluster(&smem_bar_[WHOLE_CTA? 1: warp_m_], parity_)) {}
            utils::fence_acquire_smem();  // Order wrt st_async_remote

            // Flip parity bit for next iteration
            parity_ ^= 1;

            return total;
#else
            static_assert(true, "Cluster enabled on host side but not available on device");
#endif
        } else {
            // We switch workspace every iteration.
            T *workspace = inter_cta_.phase_counter_ & 0x1 ? w1_ : w0_;

            // Warp leaders 0 hold the CTA-local results.
            if (this->warp_n_ == 0 && this->lane_ == 0) {
                workspace[bidn_] = data;
            }

            // We need to check for the case where we are using CTA > 1 and WARP_M > 1 and N == 1
            // Because in this case we cannot rely on the __syncThreads in template<typename T, uint32_t WARPS_M,
            // uint32_t WARPS_N>::compute (because we will use the specialized template <typename T, uint32_t
            // WARPS_M>::compute which does not have a sync)
            if (CTAS_PER_ROW > 1 && WARPS_M > 1 && WARPS_N == 1) {
                utils::namedBarrierSync(0, warps_m * WARPS_N * 32);
            }

            // Wait for all CTAS_PER_ROW CTAS in the group to have written their result.
            inter_cta_.sync(warps_m * WARPS_N * 32);

            T total = Zeros<T>::get();

            // Every warp does the final reduction locally.
            if constexpr (CTAS_PER_ROW <= THREADS_PER_WARP) {
                if (this->lane_ < CTAS_PER_ROW) {
                    total = workspace[this->lane_];
                }
            } else {
                static constexpr int32_t LOOP_PER_THREAD = (CTAS_PER_ROW + THREADS_PER_WARP - 1) / THREADS_PER_WARP;

                // Collect stats for the current lane from ctas (#bidn % THREADS_PER_WARP == #lane) and then do the
                // intra-warp reduction
                total = workspace[this->lane_];
#pragma unroll
                for (int32_t loop_in_thread = 1; loop_in_thread < LOOP_PER_THREAD; ++loop_in_thread) {
                    const int32_t lane_in_loop = this->lane_ + loop_in_thread * THREADS_PER_WARP;
                    if (lane_in_loop < CTAS_PER_ROW) {
                        total = op(total, workspace[lane_in_loop]);
                    }
                }
            }

            total = Reducer<T, 1, 1, 1, false, false>::allreduce_(total, op);

            return total;
        }
    }

    InterCTASync inter_cta_;
    T *w0_;
    T *w1_;
    int bidn_;
    int warp_m_;

   private:
    T *smem_cga_;
    uint64_t *smem_bar_;
    int parity_;
};

template <typename T, uint32_t WARPS_M, bool USE_CLUSTER, bool WHOLE_CTA>
struct Reducer<T, 1, WARPS_M, 1, USE_CLUSTER, WHOLE_CTA> {
    using Type = T;
    enum { SMEM_BYTES = 0 };
    enum { WORKSPACE_BYTES_PER_GROUP = 0 };
    enum { THREADS_PER_WARP = 32 };

    template <typename Params>
    inline __device__
    Reducer(Params &params,
            uint32_t bidm,
            uint32_t bidn,
            uint32_t warp_m,
            uint32_t warp_n,
            uint32_t /*tidx*/,
            uint32_t lane,
            void *smem,
            uint64_t *smem_bar)
        : warp_n_(warp_n), lane_(lane) {}

    template <typename Op>
    static inline __device__ T
    allreduce_(T data, Op &op) {
#pragma unroll
        for (int it = 1; it < THREADS_PER_WARP; it *= 2) {
            data = op(data, utils::warp_shuffle_xor(data, it));
        }
        return data;
    }

    template <typename Op>
    inline __device__ T
    allreduce(T data, Op &op, int32_t warps_m) {
        return allreduce_(data, op);
    }

   protected:
    template <typename Op>
    inline __device__ T
    reduce(T data, Op &op, int32_t warps_m) {
// only lane 0 holds the result!
#pragma unroll
        for (int it = THREADS_PER_WARP / 2; it > 0; it /= 2) {
            data = op(data, utils::warp_shuffle_down(data, it));
        }
        return data;
    }

   public:
    int warp_n_;
    int lane_;
};

template <typename T, uint32_t WARPS_M, uint32_t WARPS_N, bool USE_CLUSTER, bool WHOLE_CTA>
struct Reducer<T, 1, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA> : public Reducer<T, 1, WARPS_M, 1, USE_CLUSTER, WHOLE_CTA> {
    using Base = Reducer<T, 1, WARPS_M, 1, USE_CLUSTER, WHOLE_CTA>;
    using Type = T;

    enum { SMEM_BYTES = Base::SMEM_BYTES + WARPS_M * WARPS_N * sizeof(T) * 2 };
    enum { WORKSPACE_BYTES_PER_GROUP = 0 };
    enum { THREADS_PER_WARP = 32 };

    template <typename Params>
    inline __device__
    Reducer(Params &params,
            uint32_t bidm,
            uint32_t bidn,
            uint32_t warp_m,
            uint32_t warp_n,
            uint32_t tidx,
            uint32_t lane,
            void *smem,
            uint64_t *smem_bar)
        : Base(params, bidm, bidn, warp_m, warp_n, tidx, lane, smem, smem_bar), use0_(true) {
        smem0_ = &static_cast<T *>(smem)[warp_m * WARPS_N];
        smem1_ = smem0_ + WARPS_M * WARPS_N;
    }

    template <typename Op>
    inline __device__ T
    allreduce(T data, Op &op, int32_t warps_m) {
        T *smem = use0_ ? smem0_ : smem1_;
        use0_   = !use0_;
        data    = Base::reduce(data, op, warps_m);
        if (this->lane_ == 0) {
            smem[this->warp_n_] = data;
        }
        // The number matters here. 0 will only use one barrier per block. "1" will use 2 according to ncu.
        // We should use 0 unless it is absolutely necessary to use seperate barriers. 
        utils::namedBarrierSync(0, warps_m * WARPS_N * 32);
        T out = Zeros<T>::get();
#pragma unroll
        for (int it = 0; it < WARPS_N; it++) {
            out = op(out, smem[it]);
        }
        return out;
    }

   protected:
    template <typename Op>
    inline __device__ T
    reduce(T data, Op &op, int32_t warps_m) {
        T *smem = use0_ ? smem0_ : smem1_;
        use0_   = !use0_;
        // only intra-CTA group leader holds the result!
        data = Base::reduce(data, op, warps_m);
        if (this->lane_ == 0) {
            smem[this->warp_n_] = data;
        }
        utils::namedBarrierSync(0, warps_m * WARPS_N * 32);
        T out = Zeros<T>::get();
        if (this->warp_n_ == 0 && this->lane_ == 0) {
#pragma unroll
            for (int it = 0; it < WARPS_N; it++) {
                out = op(out, smem[it]);
            }
        }
        return out;
    }

   private:
    T *smem0_;
    T *smem1_;
    bool use0_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////
template <bool isRMSNorm, typename T, bool handleNan = false>
inline __device__ void
update_norm_stats(T &m_a, T &m2_a, T &n_a, T m_b, T m2_b, T n_b) {
    const T n_ab = n_a + n_b;

    if constexpr (!isRMSNorm) {
        const T rn_ab     = (handleNan && (n_ab == Zeros<T>::get())) ? Zeros<T>::get() : (1.f / n_ab);
        const T delta     = m_a - m_b;
        const float m2_ab = m2_a + m2_b + delta * delta * n_a * n_b * rn_ab;
        const float m_ab  = (n_a * m_a + n_b * m_b) * rn_ab;
        n_a               = n_ab;
        m_a               = m_ab;
        m2_a              = m2_ab;
    } else {
        const float m2_ab = m2_a + m2_b;
        n_a               = n_ab;
        m_a               = Zeros<T>::get();
        m2_a              = m2_ab;
    }
}

template <bool isRMSNorm, typename T, bool handleNan = false>
inline __device__ void
warp_chan_upd_dynamic(T &m_a, T &m2_a, T &n_a, int num_active) {
    // Assume at least leftmost is valid and init: step = next_pow2(num_active) / 2 (might get NaN otherwise)
    int highest_bit_set = (8 * sizeof(num_active)) - __clz(num_active - 1);

#pragma unroll
    for (int step = (1 << (highest_bit_set - 1)); step > 0; step /= 2) {
        // Exchange
        T n_b  = utils::warp_shuffle_down(n_a, step);
        T m_b  = (!isRMSNorm ? utils::warp_shuffle_down(m_a, step) : Zeros<T>::get());
        T m2_b = utils::warp_shuffle_down(m2_a, step);
        update_norm_stats<isRMSNorm, T, handleNan>(m_a, m2_a, n_a, m_b, m2_b, n_b);
    }
    // Intra-warp broadcast (only lane 0 has valid stats).
    m_a  = __shfl_sync(uint32_t(-1), m_a, 0);
    m2_a = __shfl_sync(uint32_t(-1), m2_a, 0);
}

////////////////////////////////////////////////////////////////////////////////////////////////////
// NOTE: WHOLE_CTA means if we'll use one empty mbarrier for inter-CTA sync using cluster.
// In the TMA LN engines it should be "true" since we're syncing the whole CTA with one empty mbarrier, while in the 
// APEX LN engines it should be "false" as there's one empty mbarrier per warp m (and therefore WARPS_M empty mbarriers
// in total)
template <typename T, uint32_t CTAS_PER_ROW, uint32_t WARPS_M, uint32_t WARPS_N, bool USE_CLUSTER, bool WHOLE_CTA>
struct Stats {
    // This could be done generically with the Reducer. But then we would have to exchange 3 instead of 2 fields.

    using InterCTASync = InterCTASync<CTAS_PER_ROW>;
    using BlockStats   = Stats<T, 1, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA>;
    using stats_t      = typename BlockStats::stats_t;

    enum { SMEM_BYTES = BlockStats::SMEM_BYTES };

    template <typename Params>
    inline __device__
    Stats(Params &params,
          uint32_t bidm,
          uint32_t bidn,
          uint32_t warp_m,
          uint32_t warp_n,
          uint32_t tidx,
          uint32_t lane,
          void *smem,
          uint64_t *smem_bar)
        : inter_cta_(params, bidm, bidn, tidx),
          block_stats_(params, bidm, bidn, warp_m, warp_n, tidx, lane, smem, smem_bar),
          bidn_(bidn),  // CTA id within the group.
          w0_(static_cast<stats_t *>(params.workspace) + (bidm * WARPS_M + warp_m) * CTAS_PER_ROW),
          w1_(w0_ + params.ctas_per_col * WARPS_M * CTAS_PER_ROW),
          warp_n_(warp_n),
          warp_m_(warp_m),
          lane_(lane),
          smem_bar_(smem_bar),
          parity_(0) {
        // the first several elements are the barriers
        smem_cga_ = reinterpret_cast<stats_t *>(&smem_bar[WHOLE_CTA? 2: (WARPS_M + 1)]);
    }

    template <bool isRMSNorm, uint32_t LDGS, uint32_t NUM_ELTS>
    inline __device__ stats_t
    compute(const T (&elts)[LDGS * NUM_ELTS], const T rn, int32_t warps_m) {
        static constexpr uint32_t N      = LDGS * NUM_ELTS;
        constexpr T ELTS_PER_ROW_PER_CTA = N * WARPS_N * THREADS_PER_WARP;
        // TODO rn is not really needed here..
        constexpr T block_rn = 1.f / T(ELTS_PER_ROW_PER_CTA);
        stats_t block_stats  = block_stats_.template compute<isRMSNorm, LDGS, NUM_ELTS>(elts, block_rn, warps_m);

        if constexpr (USE_CLUSTER) {
#if (__CUDA_ARCH__ >= 900) && (CUDART_VERSION >= 12080)
            auto cluster = cooperative_groups::this_cluster();

            // Broadcast local results to other CTAs inside the CGA
            // Size of smem_cga_: [WARPS_M][CTAS_PER_ROW]
            // Thread 0 sends to block 0, Thread 1 sends to block 1, etc...
            if ((warp_n_ == 0) && (lane_ < CTAS_PER_ROW)) {
                st_async_remote(&smem_cga_[warp_m_ * CTAS_PER_ROW + bidn_], block_stats, lane_, &smem_bar_[WHOLE_CTA? 0: WARPS_M]);
            }

            // Leader thread arrives on local barrier to indicate expected tx count
            // We can't use (threadIdx.x == 0) directly here because it won't work for TMA
            if ((warp_n_ == 0) && (warp_m_ == 0) && (lane_ == 0)) {
                // It's possible that not all warp_m are active here, therefore, we can't use WARPS_M directly
                uint32_t expected_tx_count = sizeof(block_stats) * CTAS_PER_ROW * warps_m;
                utils::mbarrier_arrive_expect_tx_relaxed_cluster(&smem_bar_[WHOLE_CTA? 0: WARPS_M], expected_tx_count);
            }

            // Wait on local barrier
            while (!utils::mbarrier_try_wait_parity_relaxed_cluster(&smem_bar_[WHOLE_CTA? 0: WARPS_M], parity_)) {}
            utils::fence_acquire_smem();  // Ensure we can read local smem values that are released by barrier flip

            T n  = Zeros<T>::get();
            T m  = Zeros<T>::get();
            T m2 = Zeros<T>::get();
            if (lane_ < CTAS_PER_ROW) {
                stats_t result = smem_cga_[warp_m_ * CTAS_PER_ROW + lane_];
                n              = ELTS_PER_ROW_PER_CTA;
                m              = Get<0>::of<stats_t, T>(result);  // should be 0 for RMSNorm
                m2             = Get<1>::of<stats_t, T>(result);
            }
            utils::fence_release_smem();  // Ensure read of buffer remains ordered before smem_bar.arrive.relaxed
)KERNEL"
    R"KERNEL(            // Signal barrier is empty.
            // Each thread must arrive on all the barriers (like an arrive broadcast).
            // NOTE: here we only let each warp m arrive on the remote barriers of the same warp m (i.e. the same row),
            // because if we have one empty barrier for the entire CTA, there's a difficulty how to init the count of
            // the barrier (since some warps could be disabled here). The downside is we'll need a __syncthreads() here,
            // which doesn't have obvious perf impact though.
            // TODO: we could probably have one full bar per row (same as the empty bar) to get rid of the CTA sync,
            // while the downside is more smem (mbarriers) will be needed.
            if (!WHOLE_CTA) {
                utils::namedBarrierSync(0, warps_m * WARPS_N * 32);
            }
            for (int other_block_rank = 0; other_block_rank < CTAS_PER_ROW; ++other_block_rank) {
                uint64_t *remote_bar_empty = cluster.map_shared_rank(&smem_bar_[WHOLE_CTA? 1: warp_m_], other_block_rank);
                utils::mbarrier_arrive_relaxed_cluster(remote_bar_empty);
            }

            warp_chan_upd_dynamic<isRMSNorm, T, !is_power_of_2(CTAS_PER_ROW)>(m, m2, n, CTAS_PER_ROW);

            // Wait for remote buffer to be writable again by waiting on local barrier:
            while (!utils::mbarrier_try_wait_parity_relaxed_cluster(&smem_bar_[WHOLE_CTA? 1: warp_m_], parity_)) {}
            utils::fence_acquire_smem();  // Order wrt st_async_remote

            // Flip parity bit for next iteration
            parity_ ^= 1;

            return {m, m2};
#else
            static_assert(true, "Cluster enabled on host side but not available on device");
#endif
        } else {
            // We switch workspace every iteration.
            stats_t *workspace = inter_cta_.phase_counter_ & 0x1 ? w1_ : w0_;

            // Warp leaders 0 hold the CTA-local results.
            if (warp_n_ == 0 && lane_ == 0) {
                workspace[bidn_] = block_stats;
            }

            // We need to check for the case where we are using CTA > 1 and WARP_M > 1 and N == 1
            // Because in this case we cannot rely on the __syncThreads in template<typename T, uint32_t WARPS_M,
            // uint32_t WARPS_N>::compute (because we will use the specialized template <typename T, uint32_t
            // WARPS_M>::compute which does not have a sync)
            // Bug 5221388: For WARPS_M > 1 and WARPS_N > 1, we also need the sync here, otherwise the stats received in
            // `workspace[lane_]` may be NaN (as if some stats from the other CTAs are not initialized).
            // Reason unknown yet and tracked in CL-16775.
            if (CTAS_PER_ROW > 1 && WARPS_M > 1) {
                utils::namedBarrierSync(0, warps_m * WARPS_N * 32);
            }

            // Wait for all CTAS_PER_ROW CTAS in the group to have written their result.
            inter_cta_.sync(warps_m * WARPS_N * 32);

            T n  = Zeros<T>::get();
            T m  = Zeros<T>::get();
            T m2 = Zeros<T>::get();

            // Every warp does the final reduction locally.
            if constexpr (CTAS_PER_ROW <= THREADS_PER_WARP) {
                if (lane_ < CTAS_PER_ROW) {
                    stats_t result = workspace[lane_];
                    n              = ELTS_PER_ROW_PER_CTA;
                    m              = Get<0>::of<stats_t, T>(result);  // should be 0 for RMSNorm
                    m2             = Get<1>::of<stats_t, T>(result);
                }
                warp_chan_upd_dynamic<isRMSNorm, T, !is_power_of_2(CTAS_PER_ROW)>(m, m2, n, CTAS_PER_ROW);
            } else {
                static constexpr int32_t LOOP_PER_THREAD = (CTAS_PER_ROW + THREADS_PER_WARP - 1) / THREADS_PER_WARP;

                stats_t curr = workspace[lane_];
                n            = ELTS_PER_ROW_PER_CTA;
                m            = Get<0>::of<stats_t, T>(curr);  // should be 0 for RMSNorm
                m2           = Get<1>::of<stats_t, T>(curr);

                // Collect stats for the current lane from ctas (#bidn % THREADS_PER_WARP == #lane) and then do the
                // intra-warp reduction
#pragma unroll
                for (int32_t loop_in_thread = 1; loop_in_thread < LOOP_PER_THREAD; ++loop_in_thread) {
                    const int32_t lane_in_loop = lane_ + loop_in_thread * THREADS_PER_WARP;
                    if (lane_in_loop < CTAS_PER_ROW) {
                        curr   = workspace[lane_in_loop];
                        T m_b  = Get<0>::of<stats_t, T>(curr);
                        T m2_b = Get<1>::of<stats_t, T>(curr);
                        update_norm_stats<isRMSNorm>(m, m2, n, m_b, m2_b, ELTS_PER_ROW_PER_CTA);
                    }
                }
                warp_chan_upd_dynamic<isRMSNorm>(m, m2, n, THREADS_PER_WARP);
            }

            return {m, m2};
        }
    }

    InterCTASync inter_cta_;
    BlockStats block_stats_;

    stats_t *w0_;
    stats_t *w1_;
    int bidn_;
    int warp_n_;
    int warp_m_;
    int lane_;

   private:
    stats_t *smem_cga_;
    uint64_t *smem_bar_;
    int parity_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename T, uint32_t WARPS_M, uint32_t WARPS_N, bool USE_CLUSTER, bool WHOLE_CTA>
struct Stats<T, 1, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA> {
    using WarpStats = Stats<T, 1, WARPS_M, 1, USE_CLUSTER, WHOLE_CTA>;
    using stats_t   = typename WarpStats::stats_t;

    enum { SMEM_BYTES = WARPS_M * WARPS_N * sizeof(stats_t) * 2 };

    template <typename Params>
    inline __device__
    Stats(Params &params,
          uint32_t bidm,
          uint32_t bidn,
          uint32_t warp_m,
          uint32_t warp_n,
          uint32_t tidx,
          uint32_t lane,
          void *smem,
          uint64_t *smem_bar)
        : warp_stats_(params, bidm, bidn, warp_m, warp_n, tidx, lane, smem, smem_bar), use0_(true) {
        smem0_ = static_cast<stats_t *>(smem) + warp_m * WARPS_N;
        smem1_ = smem0_ + WARPS_M * WARPS_N;
    }

    template <bool isRMSNorm, uint32_t LDGS, uint32_t NUM_ELTS>
    inline __device__ stats_t
    compute(const T (&elts)[LDGS * NUM_ELTS], const T rn, int32_t warps_m) {
        stats_t *smem = use0_ ? smem0_ : smem1_;
        use0_         = !use0_;
        // Compute warp local for all WARPS_N
        static constexpr uint32_t N = LDGS * NUM_ELTS;
        constexpr T warp_rn         = 1.f / T(N * THREADS_PER_WARP);
        stats_t warp_stats          = warp_stats_.template compute<isRMSNorm, LDGS, NUM_ELTS>(elts, warp_rn, warps_m);

        // Each warp warp leader stores its stats
        const auto warp_n = warp_stats_.reducer_.warp_n_;
        const auto lane   = warp_stats_.reducer_.lane_;
        if (lane == 0) {
            smem[warp_n] = warp_stats;
        }
        utils::namedBarrierSync(0, warps_m * WARPS_N * 32);

        T n  = Zeros<T>::get();
        T m  = Zeros<T>::get();
        T m2 = Zeros<T>::get();

        if (lane < WARPS_N) {
            stats_t result = smem[lane];
            n              = N * THREADS_PER_WARP;
            m              = Get<0>::of<stats_t, T>(result);  // should be 0 for RMSNorm
            m2             = Get<1>::of<stats_t, T>(result);
        }

        warp_chan_upd_dynamic<isRMSNorm>(m, m2, n, WARPS_N);

        return {m, m2};
    }

    WarpStats warp_stats_;
    stats_t *smem0_;
    stats_t *smem1_;
    bool use0_;
};

template <typename T, uint32_t WARPS_M, bool USE_CLUSTER, bool WHOLE_CTA>
struct Stats<T, 1, WARPS_M, 1, USE_CLUSTER, WHOLE_CTA> {
    using stats_t = typename TypeToVec2<T>::Type;
    // The simple Warp reducer.
    using Reducer = Reducer<T, 1, WARPS_M, 1, USE_CLUSTER, WHOLE_CTA>;

    enum { SMEM_BYTES = 0 };

    template <typename Params>
    inline __device__
    Stats(Params &params,
          uint32_t bidm,
          uint32_t bidn,
          uint32_t warp_m,
          uint32_t warp_n,
          uint32_t tidx,
          uint32_t lane,
          void *smem,
          uint64_t *smem_bar)
        : reducer_(params, bidm, bidn, warp_m, warp_n, tidx, lane, smem, smem_bar) {}

    template <bool isRMSNorm, uint32_t LDGS, uint32_t NUM_ELTS>
    inline __device__ stats_t
    compute(const T (&elts)[LDGS * NUM_ELTS], const T rn, int32_t warps_m) {
        static constexpr uint32_t N = LDGS * NUM_ELTS;
        auto sum                    = Sum<T>();

        T m = Zeros<T>::get();
        if (!isRMSNorm) {
// CL-14115: The unroll factor 128 for LDGS was chosen based on the compilation/perf results for APEX LN_fwd engines
#pragma unroll(128 * NUM_ELTS)
            for (int it = 0; it < N; it++) {
                m += elts[it];
            }
            m = reducer_.allreduce(m, sum, warps_m) * rn;
        }

        T m2 = Zeros<T>::get();
#pragma unroll(128 * NUM_ELTS)
        for (int it = 0; it < N; it++) {
            T diff = (elts[it] - m);
            m2 = __fmaf_ieee_rn(diff, diff, m2);
        }
        m2 = reducer_.allreduce(m2, sum, warps_m);

        return {m, m2};
    }

    Reducer reducer_;
};

////////////////////////////////////////////////////////////////////////////////////////////////////
//// Fast division helpers

// Count leading zeros - start from most significant bit.
static int
clz(int32_t x) {
    for (int32_t i = 31; i >= 0; --i)
        if (((1 << i) & x) != 0) {
            return 31 - i;
        }
    return 32;
}

static int32_t
find_log_2(int32_t x, bool round_up = false) {
    int32_t a = 31 - clz(x);
    if (round_up) {
        a += !(0 == (x & (x - 1)));
    }
    return a;
}

static void
find_divisor(int32_t denom, uint32_t &mul_coeff, uint32_t &shift_coeff) {
    if (denom == 0) {
        return;
    }
    if (denom == 1) {
        // if dividing by 1, reduced math doesn't work because mul_coeff would
        // need to be 2^32, which doesn't fit into unsigned int.  the div()
        // routine handles this special case separately.
        mul_coeff   = 0;
        shift_coeff = 0;
        return;
    }
    // To express the division N/D in terms of a multiplication, what we first
    // imagine is simply N*(1/D).  However, 1/D will always evaluate to 0 (for D>1),
    // so we need another way.  There's nothing that says we have to use exactly
    // the fraction 1/D; instead it could be any X/Y that reduces to 1/D (i.e.,
    // Y=X*D), or at least to "close enough" to it.  If we pick Y that is a power
    // of two, then the N*(X/Y) can be N*X followed by a right-shift by some amount.
    // The power of two we should pick should be at least 2^32, because in the
    // div() routine we'll use umulhi(), which returns only the upper 32 bits --
    // this being equivalent to a right-shift by 32.  But we might want a higher
    // power of two for better accuracy depending on the magnitude of the denominator.
    // Once we've picked Y, then X [our mul_coeff value] is simply Y/D, rounding up,
    // and we save shift_coeff as whatever further shift we have to do beyond
    // what the umulhi() implies.
    uint32_t p  = 31 + find_log_2(denom, true);
    uint32_t m  = ((1ull << p) + static_cast<uint32_t>(denom - 1)) / static_cast<uint32_t>(denom);
    mul_coeff   = m;
    shift_coeff = p - 32;
}

__device__ __forceinline__ uint32_t
umulhi(uint32_t x, uint32_t y) {
#if defined(__CUDA_ARCH__)
    return __umulhi(x, y);
#else
    uint64_t z = static_cast<uint64_t>(x) * y;
    return static_cast<uint32_t>(z >> 32);
#endif
}

class reduced_divisor {
   public:
    reduced_divisor() {}
    __forceinline__
    reduced_divisor(int32_t _y)
        : y(_y) {
        mul_coeff   = 0U;
        shift_coeff = 0U;
        find_divisor(y, mul_coeff, shift_coeff);
    }
    __device__ __forceinline__
    reduced_divisor(uint32_t _mul_coeff, uint32_t _shift_coeff, int32_t _y)
        : mul_coeff(_mul_coeff), shift_coeff(_shift_coeff), y(_y) {}
    __device__ __forceinline__ int32_t
    div(int32_t x) const {
        // if dividing by 1, then find_divisor wouldn't have worked because
        // mul_coeff would have had to be 2^32, which can't be represented,
        // so we have to special case that one.
        return (y != 1) ? umulhi(static_cast<uint32_t>(x), mul_coeff) >> shift_coeff : x;
    }
    __device__ __forceinline__ int32_t
    mod(int32_t x) const {
        return x - (div(x) * y);
    }
    __device__ __forceinline__ void
    divmod(int32_t x, int32_t &q, int32_t &mod) const {
        q   = div(x);
        mod = x - (q * y);
    }
    __device__ __forceinline__ int32_t
    get() const {
        return y;
    }

   protected:
    uint32_t mul_coeff;
    uint32_t shift_coeff;
    int32_t y;
};


////////////////////////////////////////////////////////////////////////////////////////////////////

template <uint32_t HIDDEN_SIZE_,
          typename weight_t_,
          typename input_t_,
          typename output_t_,
          typename compute_t_,
          typename index_t_,
          uint32_t THREADS_PER_CTA_>
struct Kernel_traits_base {
    using weight_t  = weight_t_;
    using input_t   = input_t_;
    using output_t  = output_t_;
    using compute_t = compute_t_;
    using index_t   = index_t_;

    enum { HIDDEN_SIZE = HIDDEN_SIZE_ };
    enum { THREADS_PER_CTA = THREADS_PER_CTA_ };
    enum { THREADS_PER_WARP = 32 };
};)KERNEL"
    R"KERNEL(

////////////////////////////////////////////////////////////////////////////////////////////////////

template <typename weight_t_,
          typename input_t_,
          typename output_t_,
          typename compute_t_,
          typename norm_output_t_,
          typename index_t_,
          uint32_t HIDDEN_SIZE_,
          uint32_t BATCH_SIZE_,
          uint32_t CTAS_PER_ROW_,
          uint32_t WARPS_M_,
          uint32_t WARPS_N_,
          uint32_t BYTES_PER_LDG_ = 16,
          bool isRMSNorm_         = false,
          bool isAdaLN_           = false,
          bool isBatchFirst_      = true,
          bool hasGamma_          = true,
          bool hasBeta_           = true,
          bool useGammaSmem_      = true,  // for FWD kernels it denotes the use of shmem for both gamma and beta
          bool useCluster_        = false,
          bool wholeCTA_          = false,
          uint32_t ROWS_PER_WARP_ = 1,
          typename Base           = Kernel_traits_base<HIDDEN_SIZE_,
                                             weight_t_,
                                             input_t_,
                                             output_t_,
                                             compute_t_,
                                             index_t_,
                                             WARPS_M_ * WARPS_N_ * THREADS_PER_WARP> >
struct Kernel_traits : public Base {
    using input_t       = typename Base::input_t;
    using weight_t      = typename Base::weight_t;
    using compute_t     = typename Base::compute_t;
    using output_t      = typename Base::output_t;
    using index_t       = typename Base::index_t;
    using norm_output_t = norm_output_t_;
    enum { isRMSNorm = isRMSNorm_ };
    enum { isAdaLN = isAdaLN_ };
    enum { isBatchFirst = isBatchFirst_ };
    enum { hasGamma = hasGamma_ };
    enum { hasBeta = hasBeta_ };
    enum { CTAS_PER_ROW = CTAS_PER_ROW_ };
    enum { WARPS_M = WARPS_M_ };
    enum { WARPS_N = WARPS_N_ };
    enum { COLS = HIDDEN_SIZE_ };
    enum { HIDDEN_SIZE = HIDDEN_SIZE_ };
    enum { BATCH_SIZE = BATCH_SIZE_ };
    enum { BYTES_PER_LDG = BYTES_PER_LDG_ };
    enum { NUM_ELTS = BYTES_PER_LDG / sizeof(input_t) };
    enum { USE_GAMMA_SMEM = useGammaSmem_ };
    enum { USE_CLUSTER = useCluster_ };
    enum { WHOLE_CTA = wholeCTA_ };
    enum { THREADS_PER_ROW = WARPS_N * THREADS_PER_WARP };
    enum { THREADS_PER_CTA = WARPS_M * THREADS_PER_ROW };
    enum { ROWS_PER_WARP = ROWS_PER_WARP_ };
    enum { ROWS_PER_CTA = WARPS_M * ROWS_PER_WARP };

    enum { BYTES_PER_ROW = COLS * sizeof(input_t) };
    enum { BYTES_PER_ROW_PER_CTA = THREADS_PER_ROW * BYTES_PER_LDG };
    // Multi-row per CTA not supported for multi-CTA => no smem for WGRAD needed
    enum { SMEM_BYTES_WGRAD = ROWS_PER_CTA == 1 ? 0 : BATCH_SIZE * ROWS_PER_CTA * COLS * sizeof(compute_t) };

    using reduce_t   = typename TypeToVec2<compute_t>::Type;
    using RMSReducer = Reducer<compute_t, CTAS_PER_ROW, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA>;
    using Reducer    = Reducer<reduce_t, CTAS_PER_ROW, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA>;

    enum { SMEM_BYTES_DGRAD = Reducer::SMEM_BYTES };
    enum { SMEM_BYTES = SMEM_BYTES_DGRAD + SMEM_BYTES_WGRAD };

    using Ivec     = Vec<input_t, NUM_ELTS>;
    using Ovec     = Vec<output_t, NUM_ELTS>;
    using Wvec     = Vec<weight_t, NUM_ELTS>;
    using Cvec     = Vec<compute_t, NUM_ELTS>;
    using NormOvec = Vec<norm_output_t, NUM_ELTS>;
    enum { ELTS_PER_LDG = BYTES_PER_LDG / sizeof(input_t) };

    // Assume that each thread can handle the same number of elements in the output and weights as in the input.
    // The number of columns fetched per load from input: one per thread.
    enum { VEC_COLS_PER_LDG = CTAS_PER_ROW * THREADS_PER_ROW };
    // The total number of vectorized loads/stores per hidden vector.
    enum { VEC_COLS = COLS / ELTS_PER_LDG };
    // The number of loads per thread for the input.
    enum { LDGS = VEC_COLS / VEC_COLS_PER_LDG };

    using Stats = Stats<compute_t, CTAS_PER_ROW, WARPS_M, WARPS_N, USE_CLUSTER, WHOLE_CTA>;
    enum { SMEM_BYTES_FWD = Stats::SMEM_BYTES };
};

////////////////////////////////////////////////////////////////////////////////////////////////////


constexpr float CUDNN_FLT_MIN = 1.17549435082228750796873653722224568e-38F;  // Minimum positive normalized float value
/**
 * Base class to provide common block scaling utilities, private to `BlockScaleHelper`
 */
template <typename Cvec, typename Ovec, int32_t BlockSize>
class BlockScaleHelperBase;

/**
 * Base class to provide common block scaling utilities, for nvfp4
 */
template <typename Cvec, typename Ovec>
class BlockScaleHelperBase<Cvec, Ovec, 16> {
   public:
    using scale_t = nv_fp8_e4m3;

   protected:
    using compute_t = typename Cvec::Elt_type;
    using output_t  = typename Ovec::Elt_type;

    using index_t = uint32_t;

    static constexpr int32_t THREADS_PER_WARP = 32;
    static constexpr int32_t BLOCK_SIZE       = 16;
    static constexpr int32_t NUM_ELTS         = sizeof(Cvec) / sizeof(compute_t);
    static constexpr int32_t GROUP_SIZE       = BLOCK_SIZE / NUM_ELTS;
    static constexpr int32_t SF_COUNT         = THREADS_PER_WARP / GROUP_SIZE;

    static __device__ compute_t
    scale(const compute_t value) {
        constexpr compute_t nvfp4_max = 6.0f;
        compute_t result = value / nvfp4_max;
        // Clamp to CUDNN_FLT_MIN to match reference implementation and avoid 0 scaling factors
        return fmaxf(result, CUDNN_FLT_MIN);
    }

    static __device__ scale_t
    fp32ToE4M3(float value) {
        return scale_t(value);
    }

    static __device__ Vec<uint8_t, NUM_ELTS>
    vecDivideAndCast(const Cvec &input, compute_t scale) {
        // Construct the nv_fp4x2_e2m1 outputs.
        // If NUM_ELTS == 1, nv_fp4(ele1,0)
        // If NUM_ELTS % 2 == 0; f2_i=nv_fp4(ele_i, ele_i+1) --> out[i/2]
        // Per cuda_fp4.h and cuda_fp8.h
        // The data type of  __nv_fp4x2_storage_t is __nv_fp8_storage_t
        // The data type of __nv_fp8_storage_t is unsigned char:
        // typedef unsigned char __nv_fp8_storage_t;
        // The reason why we use  unsigned char is that the __shfl_xor_sync don't support nv_fp4 data type
        Vec<uint8_t, NUM_ELTS> output;
#pragma unroll
        for (int i = 0; i < NUM_ELTS; i += 2) {
            float2 f2 =
                make_float2((input.data.elt[i] / scale), (i + 1 < NUM_ELTS ? (input.data.elt[i + 1] / scale) : 0));
            output.data.elt[i / 2] = static_cast<uint8_t>(nv_fp4x2_e2m1(f2).__x);
        }
        return output;
    }
};

/**
 * Base class to provide common block scaling utilities, for mxfp8
 */
template <typename Cvec, typename Ovec>
class BlockScaleHelperBase<Cvec, Ovec, 32> {
   public:
    using lne8m0_t = uint8_t;
    using scale_t  = lne8m0_t;

   protected:
    using compute_t = typename Cvec::Elt_type;
    using output_t  = typename Ovec::Elt_type;
    using index_t   = uint32_t;

    static constexpr int32_t THREADS_PER_WARP = 32;
    static constexpr int32_t BLOCK_SIZE       = 32;
    static constexpr int32_t NUM_ELTS         = sizeof(Cvec) / sizeof(compute_t);
    static constexpr int32_t GROUP_SIZE       = BLOCK_SIZE / NUM_ELTS;

    static __device__ compute_t
    scale(const compute_t value) {
        static constexpr float FP8_MAX = std::is_same<output_t, nv_fp8_e5m2>::value ? 57344.f : 448.f;
        compute_t result = value / FP8_MAX;
        // Clamp to CUDNN_FLT_MIN to match reference implementation and avoid 0 scaling factors
        return fmaxf(result, CUDNN_FLT_MIN);
    }

    template <typename IOvec>
    static __device__ Ovec
    vecDivideAndCast(const IOvec &input, compute_t scale) {
        Ovec output;
#pragma unroll
        for (int32_t i = 0; i < NUM_ELTS; ++i) {
            output.data.elt[i] = output_t((compute_t(input.data.elt[i]) / scale));
        }
        return output;
    };
};

/**
 * Helper class to define the block scale operations
 */
template <typename Cvec, typename Ovec, int32_t BlockSize>
class BlockScaleRowHelper;

template <typename Cvec, typename Ovec, int32_t BlockSize>
class BlockScaleColHelper;

/**
 * Specialization for MXFP8 rowwise scaling
 */
template <typename Cvec, typename Ovec>
class BlockScaleRowHelper<Cvec, Ovec, 32> : private BlockScaleHelperBase<Cvec, Ovec, 32> {
   private:
    using Base = BlockScaleHelperBase<Cvec, Ovec, 32>;
    using Base::BLOCK_SIZE;
    using Base::GROUP_SIZE;
    using Base::NUM_ELTS;
    using Base::THREADS_PER_WARP;
    using typename Base::compute_t;
    using typename Base::index_t;
    using typename Base::scale_t;

   public:
    __device__
    BlockScaleRowHelper() = default;

    template <typename IOvec>
    __device__ void
    blockQuantizeStore(const IOvec &z_math, void *sf_ptr, index_t sf_idx, void *z_ptr, index_t z_idx) {
        const index_t group_id = lane / GROUP_SIZE;

        // 1. compute scaling factor
        compute_t sf = 0;
#pragma unroll
        // compute the amax across NUM_ELTS elements
        for (int jt = 0; jt < NUM_ELTS; jt++) {
            sf = fmaxf(sf, fabsf(compute_t(z_math.data.elt[jt])));
        }

// Currently disabled for rowwise block scaling as perf regression observed
#if 0 && (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000) && defined(__CUDA_ARCH_FEAT_SM100_ALL)
#pragma unroll
        for (int jt = 0; jt < NUM_ELTS; jt++) {
            if (group_id == jt) {
                asm volatile("redux.sync.max.f32 %0, %1, %2;" : "=f"(sf) : "f"(sf), "r"(__activemask()));
            }
        }
#else
        // warp reduce (binary reduction) starting from GROUP_SIZE / 2 to aggregate
        // a total of NUM_ELTS * GROUP_SIZE = BLOCK_SIZE number of z values in sf
#pragma unroll
        for (int offset = GROUP_SIZE / 2; offset > 0; offset /= 2) {
            sf = fmaxf(sf, __shfl_down_sync(0xffffffff, sf, offset));
        }

        // broadcast from the lane with the valid values to the other threads in the group
        sf = __shfl_sync(0xffffffff, sf, group_id * GROUP_SIZE);
#endif
        // scale the row amax to get the scaling factor
        sf = Base::scale(sf);

        // 2. store scaling factor
        // Warp orchestration to compute scale_row:
        // e.g. for NUM_ELTS == 2:
        //   |  W0SF0 W0SF1  |  W0SF0 W0SF1  | ... |        W0SF0 W0SF1       |  W0SF0 W0SF1  |...
        //   |  W1SF0 W1SF1  |  W1SF0 W1SF1  | ... |        W1SF0 W1SF1       |  W1SF0 W1SF1  |...
        //   |      ...      |      ...      | ... |            ...           |      ...      |...
        //   | W31SF0 W31SF1 | W31SF0 W31SF1 | ... |        W31SF0 W31SF1     | W31SF0 W31SF1 |...
        //   |---------------|---------------|-----|--------------------------|---------------|...
        //   |----- CTA0 ----|----- CTA1 ----| ... |--- CTA<ctas_per_row-1> --|----- CTA0 ----|...
        //   |                                                                |
        //   |-------------------------------LDG0-----------------------------|-------LDG1-----...
        Vec<scale_t, NUM_ELTS> sf_out;
        // only correct for group_id == 0, but it's ok because only lane == 0 is storing sf_out
        sf_out.data.elt[0] = float_to_e8m0(sf);
#pragma unroll
        for (int jt = 1; jt < NUM_ELTS; jt++) {
            sf_out.data.elt[jt] = float_to_e8m0(__shfl_sync(0xffffffff, sf, jt * GROUP_SIZE));
        }
        if (lane == 0) {
            sf_out.store_to(sf_ptr, sf_idx);
        }
        // No need to check for 0x00 - scale() now clamps to FLT_MIN which converts to 0x01
        float scale = e8m0_to_float(sf_out.data.elt[0]);
        Ovec z_row_scaled = Base::vecDivideAndCast(z_math, scale);
        z_row_scaled.store_to(z_ptr, z_idx);
    }

   private:
    const index_t lane = threadIdx.x % THREADS_PER_WARP;
};

/**
 * Specialization for MXFP8 colwise scaling in 1D2X2X
 */
template <typename Cvec, typename Ovec>
class BlockScaleColHelper<Cvec, Ovec, 32> : private BlockScaleHelperBase<Cvec, Ovec, 32> {
   private:
    using Base = BlockScaleHelperBase<Cvec, Ovec, 32>;
    using Base::BLOCK_SIZE;
    using Base::GROUP_SIZE;
    using Base::NUM_ELTS;
    using Base::THREADS_PER_WARP;
    using typename Base::compute_t;
    using typename Base::index_t;
    using typename Base::output_t;
    using typename Base::scale_t;

   public:
    __device__
    BlockScaleColHelper(float *mxfp8_tile)
        : mxfp8_tile_(mxfp8_tile) {}

    template <typename IOvec>
    __device__ void
    initTile(const IOvec &z_math, int num_threads) {)KERNEL"
    R"KERNEL(        return initTile(z_math, warp, num_threads);
    }

    template <typename IOvec>
    __device__ void
    initTile(const IOvec &z_math, index_t row_in_tile, int num_threads) {
        constexpr int entry_barrier_arbitrary = 5;
        utils::namedBarrierSync(entry_barrier_arbitrary, num_threads);
#pragma unroll
        for (index_t jt = 0; jt < NUM_ELTS; jt++) {
            index_t idx_in_mxfp8_smem =
                (jt * (BLOCK_SIZE + /*padding row*/ 1) + row_in_tile) * MXFP8_1D2X2X_SMEM_COLS + lane;
            mxfp8_tile_[idx_in_mxfp8_smem] = float(z_math.data.elt[jt]);
        }
    }

    template <bool Transpose>
    __device__ void
    blockQuantizeStore(void *sf_ptr,
                       index_t sf_row_idx,
                       index_t sf_col_idx,
                       index_t sf_row_width,
                       void *z_ptr,
                       index_t z_idx,
                       int num_threads) {
        return blockQuantizeStore<Transpose>(
            sf_ptr, sf_row_idx, sf_col_idx, sf_row_width, z_ptr, z_idx, num_threads, warp);
    }

    template <bool Transpose>
    __device__ void
    blockQuantizeStore(void *sf_ptr,
                       index_t sf_row_idx,
                       index_t sf_col_idx,
                       index_t sf_row_width,
                       void *z_ptr,
                       index_t z_idx,
                       int num_threads,
                       index_t row_in_tile) {
        // Ensure all temp outputs have been written to mxfp8 smem for the current ldg
        constexpr int entry_barrier_arbitrary = 4;
        utils::namedBarrierSync(entry_barrier_arbitrary, num_threads);
        const index_t group_id = lane / GROUP_SIZE;

#if (defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 1000) && defined(__CUDA_ARCH_FEAT_SM100_ALL)
        // For sm_100a arch, we can use redux.sync.op{.abs.}{.NaN}.f32 instruction to reduce the min/max value
        // https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-redux-sync
        compute_t amax_vals[NUM_ELTS];
#pragma unroll
        for (int jt = 0; jt < NUM_ELTS; jt++) {
            const index_t idx_in_mxfp8_smem = jt * NEXT_COL_SMEM_OFFSET + lane * MXFP8_1D2X2X_SMEM_COLS + row_in_tile;
            asm volatile("redux.sync.max.abs.f32 %0, %1, 0xffffffff;"
                         : "=f"(amax_vals[jt])
                         : "f"(mxfp8_tile_[idx_in_mxfp8_smem]));
        }
        compute_t sf = amax_vals[group_id];
#else
        // 1. compute scaling factor
        compute_t sf = 0;
#pragma unroll
        // compute the amax across NUM_ELTS elements in the same col (in shared memory, adjacent rows)
        for (int jt = 0; jt < NUM_ELTS * MXFP8_1D2X2X_SMEM_COLS; jt += MXFP8_1D2X2X_SMEM_COLS) {
            const index_t idx_in_mxfp8_smem = ROW_SMEM_OFFSET + jt + row_in_tile;
            sf                              = fmaxf(sf, fabsf(mxfp8_tile_[idx_in_mxfp8_smem]));
        }

        // warp reduce (binary reduction) starting from GROUP_SIZE / 2 to aggregate
        // a total of NUM_ELTS * GROUP_SIZE = BLOCK_SIZE number of z values in sf
#pragma unroll
        for (int offset = GROUP_SIZE / 2; offset > 0; offset /= 2) {
            sf = fmaxf(sf, __shfl_down_sync(0xffffffff, sf, offset));
        }

        // broadcast from the lane with the valid values to the other threads in the group
        sf = __shfl_sync(0xffffffff, sf, group_id * GROUP_SIZE);
#endif

        // scale the col amax to get the scaling factor
        sf = Base::scale(sf);

        // Warp orchestration to compute scale_col:
        // e.g. for NUM_ELTS == 2:
        //   | W0SF0 W0SF1 W1SF0 W1SF1 ... W31SF0 W31SF1 | W0SF0 W0SF1 ... W31SF0 W31SF1 |...
        //   |----warp0---|---warp1--| ... |----warp31---|----warp0---|... |----warp31---|...
        //   |--------------------CTA0-------------------|---------------CTA1------------|...
        //   |------------------------------------------LDG0------------------------------...
        Ovec z_col_scaled;
        if constexpr (!Transpose) {
            // 2. store scaling factor
            // different columns require different sf, so all threads must exchange sf
            Vec<scale_t, NUM_ELTS> sf_out;
#pragma unroll
            for (int jt = 0; jt < NUM_ELTS; jt++) {
                sf_out.data.elt[jt] = float_to_e8m0(__shfl_sync(0xffffffff, sf, jt * GROUP_SIZE));
            }
            if (lane == 0) {
                sf_out.store_to(sf_ptr, sf_row_idx * sf_row_width + sf_col_idx);
            }

            // 3. scale and store output
#pragma unroll
            for (int jt = 0; jt < NUM_ELTS; jt++) {
                // iterate over NUM_ELTS elements in a row (in shared memory, every other NUM_ELTS in a column)
                index_t z_idx_in_mxfp8_smem = jt * NEXT_COL_SMEM_OFFSET + lane * MXFP8_1D2X2X_SMEM_COLS + row_in_tile;
                // No need to check for 0x00 - scale() now clamps to FLT_MIN which converts to 0x01
                float scale = e8m0_to_float(sf_out.data.elt[jt]);
                z_col_scaled.data.elt[jt] =
                    output_t(mxfp8_tile_[z_idx_in_mxfp8_smem] / scale);
            }
        } else {
            // 2. store scaling factor
            uint8_t sf_e8m0 = float_to_e8m0(sf);
            // No need to check for 0x00 - scale() now clamps to FLT_MIN which converts to 0x01
            if (lane % GROUP_SIZE == 0) {
                static_cast<scale_t *>(sf_ptr)[(sf_row_idx + group_id) * sf_row_width + sf_col_idx] = sf_e8m0;
            }

            // 3. scale and store output
            float scale = e8m0_to_float(sf_e8m0);
            // No need to check for 0.0f - scale() now clamps to FLT_MIN
#pragma unroll
            for (int jt = 0; jt < NUM_ELTS; jt++) {
                // iterate over NUM_ELTS elements in a column (in shared memory, adjacent NUM_ELTS rows)
                index_t z_idx_in_mxfp8_smem = ROW_SMEM_OFFSET + jt * MXFP8_1D2X2X_SMEM_COLS + row_in_tile;
                z_col_scaled.data.elt[jt]   = output_t(mxfp8_tile_[z_idx_in_mxfp8_smem] / scale);
            }
        }
        z_col_scaled.store_to(z_ptr, z_idx);

    }

   private:
    static constexpr index_t MXFP8_1D2X2X_SMEM_ROWS = BLOCK_SIZE * NUM_ELTS + /*padding rows between*/ (NUM_ELTS - 1);
    static constexpr index_t MXFP8_1D2X2X_SMEM_COLS = BLOCK_SIZE + /*padding col*/ 1;
    static constexpr index_t NEXT_COL_SMEM_OFFSET   = (BLOCK_SIZE + /*padding row*/ 1) * MXFP8_1D2X2X_SMEM_COLS;
    // For MXFP8 1d2x2x, we store the 32x32xNUM_ELTS tile in shared memory to compute the col amax for scaling within
    // the current CTA
    // e.g. for NUM_ELTS == 2, the shared memory layout for 1d2x2x is as follows:
    //
    //     col0   col2   col4   col6  ...  col62
    //      ↓      ↓      ↓      ↓           ↓
    //   |-------------------------------------------|
    //   | T0V0 | T1V0 | T2V0 | T3V0 |...| T31V0 | 0 |  <-- warp0  to store at jt=0
    //   | T0V0 | T1V0 | T2V0 | T3V0 |...| T31V0 | 0 |  <-- warp1  to store at jt=0
    //   |  ..  |  ..  |  ..  |  ..  |...|  ...  | 0 |  <-- ...
    //   | T0V0 | T1V0 | T2V0 | T3V0 |...| T31V0 | 0 |  <-- warp31 to store at jt=0
    //   |--------------------|------|---|-------|---|
    //   |  0   |  0   |  0   |  0   |...|   0   | 0 |  <-- padding row between to avoid bank conflicts when col-scaling
    //   |--------------------|------|---|-------|---|
    //   | T0V1 | T1V1 | T2V1 | T3V1 |...| T31V1 | 0 |  <-- warp0  to store at jt=1
    //   | T0V1 | T1V1 | T2V1 | T3V1 |...| T31V1 | 0 |  <-- warp1  to store at jt=1
    //   |  ..  |  ..  |  ..  |  ..  |...|  ...  | 0 |  <-- ...
    //   | T0V1 | T1V1 | T2V1 | T3V1 |...| T31V1 | 0 |  <-- warp31 to store at jt=1
    //   |-------------------------------------------|
    //      ↑      ↑      ↑      ↑           ↑     ↑
    //     col1   col3   col5   col7  ...  col63   padding col to avoid bank conflicts when computing col-scale factors
    //
    //      ↑      ↑      ↑      ↑           ↑
    //    warp0  warp1  warp2  warp3  ...  warp31
    //  to compute and apply the col-scale factors
    float *mxfp8_tile_ = nullptr;

    const index_t lane = threadIdx.x % THREADS_PER_WARP;
    const index_t warp = threadIdx.x / THREADS_PER_WARP;

    const index_t ROW_SMEM_OFFSET = ((lane * NUM_ELTS) + /*padding rows*/ (lane / GROUP_SIZE)) * MXFP8_1D2X2X_SMEM_COLS;
};

/**
 * Specialization for NVFP4 rowwise scaling
 */
template <typename Cvec, typename Ovec>
class BlockScaleRowHelper<Cvec, Ovec, 16> : private BlockScaleHelperBase<Cvec, Ovec, 16> {
    using Base = BlockScaleHelperBase<Cvec, Ovec, 16>;
    using Base::BLOCK_SIZE;
    using Base::GROUP_SIZE;
    using Base::NUM_ELTS;
    using Base::SF_COUNT;
    using Base::THREADS_PER_WARP;
    using typename Base::compute_t;
    using typename Base::index_t;
    using typename Base::output_t;
    using typename Base::scale_t;

   public:
    __device__
    BlockScaleRowHelper() = default;

    __device__ void
    blockQuantizeStore(const Cvec &z_math, void *sf_ptr, index_t sf_idx, void *z_ptr, index_t z_idx) {
        // 1. compute scaling factor
        compute_t sf = 0;
#pragma unroll
        // compute the amax across NUM_ELTS elements
        for (int jt = 0; jt < NUM_ELTS; jt++) {
            sf = fmaxf(sf, fabsf(z_math.data.elt[jt]));
        }

        // warp reduce (binary reduction) starting from GROUP_SIZE / 2 to aggregate
        // a total of NUM_ELTS * GROUP_SIZE = BLOCK_SIZE number of z values in sf
#pragma unroll
        for (int offset = GROUP_SIZE / 2; offset > 0; offset /= 2) {
            sf = fmaxf(sf, __shfl_down_sync(0xffffffff, sf, offset));
        }

        // broadcast from the lane with the valid values to the other threads in the group
        const index_t group_id = lane / GROUP_SIZE;
        sf                     = __shfl_sync(0xffffffff, sf, group_id * GROUP_SIZE);

        // scale the row amax to get the scaling factor
        sf = Base::scale(sf);

        // 2. store scaling factor
        Vec<scale_t, SF_COUNT> sf_out;
        // only correct for group_id == 0, but it's ok because only lane == 0 is storing sf_out
        sf_out.data.elt[0] = Base::fp32ToE4M3(sf);
#pragma unroll
        for (int jt = 1; jt < SF_COUNT; jt++) {
            sf_out.data.elt[jt] = Base::fp32ToE4M3(__shfl_sync(0xffffffff, sf, jt * GROUP_SIZE));
        }
        if (lane == 0) {
            sf_out.store_to(sf_ptr, sf_idx);
        }

        // TODO: can consider doing something like the following to remove the need for a shuffle loop
        // template <int32_t N>
        // struct UintContainer {
        //     using type = std::conditional_t<N == 1, uint8_t,
        //                  std::conditional_t<N == 2, uint16_t,
        //                  std::conditional_t<N == 4, uint32_t,
        //                  std::conditional_t<N == 8, uint64_t,
        //                  void>>>>;  // void for unsupported sizes
        // };
        // Base::vecDivideAndCast returns Vec<UintContainer<NUM_ELTS / 2>, 2>
        // 3. scale and store output
        auto z_row_scaled = Base::vecDivideAndCast(z_math, sf);
#pragma unroll
        // Consturct z_row_scaled by combining the adjacent thread's z_row_scaled
        // Vec<uint8_t, NUM_ELTS> z_row_scaled
        for (int i = 0; i < NUM_ELTS / 2; i++) {
            uint8_t adjacent_val                    = __shfl_xor_sync(0xffffffff, z_row_scaled.data.elt[i], 1);
            z_row_scaled.data.elt[i + NUM_ELTS / 2] = adjacent_val;
        }
        if (NUM_ELTS == 1) {
            uint8_t adjacent_val = __shfl_xor_sync(0xffffffff, z_row_scaled.data.elt[0], 1);
            z_row_scaled.data.elt[0] |= adjacent_val << 4U;
        }
        if (lane % 2 == 0) {
            z_row_scaled.store_to(z_ptr, z_idx / 2);
        }
    }

   private:
    const index_t lane = threadIdx.x % THREADS_PER_WARP;
};

)KERNEL";
inline constexpr size_t ln_headers_source_len = sizeof(ln_headers_source) - 1;

}  // namespace generated
}  // namespace experimental
}  // namespace cudnn_frontend

#if defined(_MSC_VER)
#pragma warning(pop)
#endif

#if defined(__clang__)
#pragma clang diagnostic pop
#endif
