submission 488180

import os
import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_source = r"""
#include <cuda_fp16.h>
#include <cudaTypedefs.h>

#include <torch/extension.h>
#include <torch/library.h>


#define WARP_SIZE 32


void check_cu_error(CUresult error) {
  if (error == CUDA_SUCCESS) return;
  const char *error_msg_ptr;
  if (cuGetErrorString(error, &error_msg_ptr) != CUDA_SUCCESS)
    error_msg_ptr = "unable to get error string";
  TORCH_CHECK(false, "cuTensorMapEncodeTiled error: ", error_msg_ptr);
}


template <const int NUM_ELEMENTS>
inline void create_tmap_descriptor(
    CUtensorMap *tmap,
    const char *ptr,
    uint64_t global_height, uint64_t global_width,
    uint32_t shared_height, uint32_t shared_width,
    CUtensorMapSwizzle swizzle_type
) {
    /*
    The goal is to transfer multiple of [shared_height, NUM_ELEMENTS] spanning
    [shared_height, shared_width] --> [shared_width / NUM_ELEMENTS, shared_height, NUM_ELEMENTS].

    Code taken and modified from:
    - https://docs.nvidia.com/cuda/cuda-programming-guide/04-special-topics/async-copies.html#using-tma-to-transfer-multi-dimensional-arrays.
    - https://gau-nernst.github.io/tcgen05/ 
    */
    constexpr int rank{3};
    uint64_t global_dim[rank] = {NUM_ELEMENTS, global_height, global_width / (uint64_t) NUM_ELEMENTS};
    // 4 bits would be 1/2 bytes.
    uint64_t global_strides[rank - 1] = {global_width / 2, NUM_ELEMENTS / 2}; 
    uint32_t box_dim[rank] = {NUM_ELEMENTS, shared_height, shared_width / NUM_ELEMENTS};
    uint32_t element_strides[rank] = {1, 1, 1};

    auto error = cuTensorMapEncodeTiled(
        tmap,
        CUtensorMapDataType::CU_TENSOR_MAP_DATA_TYPE_16U4_ALIGN8B,
        rank,
        (void *)ptr,
        global_dim,
        global_strides,
        box_dim,
        element_strides,
        // Interleave patterns can be used to accelerate loading of values that
        // are less than 4 bytes long.
        CUtensorMapInterleave::CU_TENSOR_MAP_INTERLEAVE_NONE,
        swizzle_type,
        // L2 Promotion can be used to widen the effect of a cache-policy to a wider
        // set of L2 cache lines.
        CUtensorMapL2promotion::CU_TENSOR_MAP_L2_PROMOTION_NONE,
        // Any element that is outside of bounds will be set to zero by the TMA transfer.
        CUtensorMapFloatOOBfill::CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
    check_cu_error(error);
}


// https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cute/arch/cluster_sm90.hpp#L180
__device__ inline uint32_t elect_sync() {
    uint32_t pred = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred %%px;\n\t"
        "elect.sync _|%%px, %1;\n\t"
        "@%%px mov.s32 %0, 1;\n\t"
        "}"
        : "+r"(pred)
        : "r"(0xFFFFFFFF));
    return pred;
}


__device__ inline void mbarrier_init(int mbar_addr, int count) {
    asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(mbar_addr), "r"(count));
}


// https://github.com/NVIDIA/cutlass/blob/v4.2.1/include/cutlass/arch/barrier.h#L408
__device__ inline void mbarrier_wait(int mbar_addr, int phase) {
    uint32_t ticks = 0x989680; // arbitrarily large timer value.
    asm volatile(
        "{\n\t"
        ".reg .pred P1;\n\t"
        "LAB_WAIT:\n\t"
        "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 P1, [%0], %1, %2;\n\t"
        "@P1 bra.uni DONE;\n\t"
        "bra.uni LAB_WAIT;\n\t"
        "DONE:\n\t"
        "}" ::"r"(mbar_addr),
        "r"(phase), "r"(ticks));
}


__device__ inline
void tma_gmem2smem(int dst, const void *src, int size, int mbar_addr) {
    asm volatile(
        "cp.async.bulk.shared::cta.global.mbarrier::complete_tx::bytes [%0], [%1], %2, [%3];"
        :: "r"(dst), "l"(src), "r"(size), "r"(mbar_addr));
}


template <int CTA_GROUP = 1>
__device__ inline void tma_3d_gmem2smem(int dst, const void *tmap_ptr, int x, int y, int z, int mbar_addr)
{
    // when CTA_GROUP=1, we can use .shared::cta instead.
    // but .shared::cluster doesn't seem to be slower, so always use it unconditionally here.
    // .cta_group::2 allows mbar_addr and dst to be in different CTA's smem.
    asm volatile("cp.async.bulk.tensor.3d.shared::cluster.global.mbarrier::complete_tx::bytes.cta_group::%6 "
                 "[%0], [%1, {%2, %3, %4}], [%5];" ::"r"(dst),
                 "l"(tmap_ptr), "r"(x), "r"(y), "r"(z), "r"(mbar_addr), "n"(CTA_GROUP)
                 : "memory");
}


// Encodes the matrix descriptor and ensures 64 bits.
__device__ inline
constexpr uint64_t encode_descriptor(uint64_t x) { return (x & 0x3'FFFFULL) >> 4ULL; }


// Copy scale factors from shared memory to tensor memory.
// .32x128b = 32 rows x 16 bytes = one scale factor tile for one MMA.
// .warpx4 duplicates data across all 32-lane groups.
__device__ inline
void copy_sf_smem2tmem(int taddr, uint64_t s_desc) {
    asm volatile("tcgen05.cp.cta_group::1.32x128b.warpx4 [%0], %1;" :: "r"(taddr), "l"(s_desc));
}

// Issue FP4 MMA instruction with block scaling.
// d_tmem=0: accumulator always starts at TMEM column 0.
// enable_input_d: 0 = clear accumulator, nonzero = accumulate.
__device__ inline
void run_mma_nvfp4(
    uint64_t a_desc,
    uint64_t b_desc,
    uint32_t i_desc,
    int scale_A_tmem,
    int scale_B_tmem,
    int enable_input_d
) {
    const int d_tmem = 0;
    asm volatile(
        "{\n\t"
        ".reg .pred p;\n\t"
        "setp.ne.b32 p, %6, 0;\n\t"
        "tcgen05.mma.cta_group::1.kind::mxf4nvf4.block_scale.block16 [%0], %1, %2, %3, [%4], [%5], p;\n\t"
        "}"
        :: "r"(d_tmem), "l"(a_desc), "l"(b_desc), "r"(i_desc),
           "r"(scale_A_tmem), "r"(scale_B_tmem), "r"(enable_input_d)
    );
}


constexpr int MMA_K = 64;  // FP4 MMA K-dimension size.


template <const int NUM_THREADS, const int BLOCK_M, const int BLOCK_N, const int BLOCK_K>
__global__
__launch_bounds__(NUM_THREADS) void kernel_v01_naive(
    const __grid_constant__ CUtensorMap A_tmap,
    const __grid_constant__ CUtensorMap B_tmap,
    const char *SFA,
    const char *SFB,
    half *C,
    int M,
    int N,
    int K
) {
    const int thread_idx{static_cast<int>(threadIdx.x)};
    const int block_idx{static_cast<int>(blockIdx.x)};

    const int warp_idx{thread_idx / WARP_SIZE};

    const int grid_dim_n{N / BLOCK_N};

    const int block_idx_m{block_idx / grid_dim_n};
    const int block_idx_n{block_idx % grid_dim_n};

    const int offset_m{block_idx_m * BLOCK_M};
    const int offset_n{block_idx_n * BLOCK_N};

    // Set up shared memory.
    // Layout: [A tile | B tile | SFA tile | SFB tile]
    extern __shared__ __align__(1024) char smem[];
    const int A_smem{static_cast<int>(__cvta_generic_to_shared(smem))};
    const int B_smem{A_smem + BLOCK_M * BLOCK_K / 2};
    constexpr int SF_size = 128 * BLOCK_K / 16;
    const int SFA_smem{B_smem + BLOCK_N * BLOCK_K / 2};
    const int SFB_smem{SFA_smem + SF_size};

#pragma nv_diag_suppress static_var_with_dynamic_init
    __shared__ uint64_t mbars[1];
    const int mbar_addr{static_cast<int>(__cvta_generic_to_shared(mbars))};
    __shared__ int tmem_addr[1];

    // TMEM layout:
    // Columns [0, BLOCK_N)                          : accumulator D
    // Columns [BLOCK_N, BLOCK_N + 4*BLOCK_K/MMA_K)  : SFA
    // Columns [BLOCK_N + 4*BLOCK_K/MMA_K, ...)       : SFB
    constexpr int SFA_tmem = BLOCK_N;
    constexpr int SFB_tmem = SFA_tmem + 4 * (BLOCK_K / MMA_K);
    constexpr int TMEM_COLS = BLOCK_N + 8 * (BLOCK_K / MMA_K);  // = BLOCK_N + BLOCK_K/8

    if (warp_idx == 0 && elect_sync()) {
        mbarrier_init(mbar_addr, 1);
        asm volatile("fence.mbarrier_init.release.cluster;");
    } else if (warp_idx == 1) {
        // Allocate TMEM for accumulator + scale factors.
        const int addr{static_cast<int>(__cvta_generic_to_shared(tmem_addr))};
        asm volatile("tcgen05.alloc.cta_group::1.sync.aligned.shared::cta.b32 [%0], %1;"
            ::"r"(addr), "r"(TMEM_COLS));
    }
    __syncthreads();

    const int taddr{tmem_addr[0]};
    int phase{0};

    // Instruction descriptor for tcgen05.mma.kind::mxf4nvf4
    // atype=E2M1 (1), btype=E2M1 (1), MMA_N and MMA_M encoded in upper bits
    constexpr uint32_t i_desc = (1U << 7U)                         // atype=E2M1
                                | (1U << 10U)                      // btype=E2M1
                                | ((uint32_t)BLOCK_N >> 3U << 17U) // MMA_N
                                | ((uint32_t)BLOCK_M >> 7U << 27U) // MMA_M
                                ;

    for (int iter_k{0}; iter_k < (K / BLOCK_K); ++iter_k) {
        // === Phase 1: Load from global memory to shared memory via TMA ===
        if (warp_idx == 0 && elect_sync()) {
            const int off_k{iter_k * BLOCK_K};

            // Load A and B tiles via 3D TMA tensor map.
            // z-coordinate = off_k / 32 because NUM_ELEMENTS=32 (T for FP4) in the tensor map.
            tma_3d_gmem2smem(A_smem, &A_tmap, 0, offset_m, off_k / 32, mbar_addr);
            tma_3d_gmem2smem(B_smem, &B_tmap, 0, offset_n, off_k / 32, mbar_addr);

            // Load SFA/SFB via 1D bulk copy.
            // Underlying storage order is (L, M/128, rest_k, 32, 4, 4) — each atom is 512 bytes.
            const int rest_k = K / 16 / 4;  // number of K-atoms (each covers 64 K-elements)
            const char *SFA_src = SFA + ((offset_m / 128) * rest_k + off_k / (16 * 4)) * 512;
            const char *SFB_src = SFB + ((offset_n / 128) * rest_k + off_k / (16 * 4)) * 512;
            tma_gmem2smem(SFA_smem, SFA_src, SF_size, mbar_addr);
            tma_gmem2smem(SFB_smem, SFB_src, SF_size, mbar_addr);

            // Signal expected number of bytes for all TMA transfers.
            constexpr int cp_size = (BLOCK_M + BLOCK_N) * BLOCK_K / 2 + 2 * SF_size;
            asm volatile("mbarrier.arrive.expect_tx.release.cta.shared::cta.b64 _, [%0], %1;"
                ::"r"(mbar_addr), "r"(cp_size) : "memory");
        }

        // Wait for TMA to complete.
        mbarrier_wait(mbar_addr, phase);
        phase ^= 1;

        // === Phase 2: Copy scale factors to TMEM, then perform MMA ===
        if (warp_idx == 0 && elect_sync())
        {
            // Canonical no-swizzle shared memory descriptor for A/B operands.
            // Layout: ((8, m), (T, 2k)) where T = 128-bit / 4-bit = 32 elements.
            // SBO = 8 * T_bytes = 8 * 16 = 128 (stride between 8-row groups).
            // LBO = stride between T-groups in shared memory.
            //   TMA with NUM_ELEMENTS=32 lays out [BLOCK_K/32, BLOCK_M, 32],
            //   so T-group 1 starts at BLOCK_M * 16 bytes after T-group 0.
            auto make_desc_A = [](int addr) -> uint64_t
            {
                constexpr int SBO = 8 * 16;  // = 128 bytes
                constexpr int LBO = BLOCK_M * 16;
                return encode_descriptor(addr) | (encode_descriptor(LBO) << 16ULL) | (encode_descriptor(SBO) << 32ULL) | (1ULL << 46ULL);
            };
            auto make_desc_B = [](int addr) -> uint64_t
            {
                constexpr int SBO = 8 * 16;  // = 128 bytes
                constexpr int LBO = BLOCK_N * 16;
                return encode_descriptor(addr) | (encode_descriptor(LBO) << 16ULL) | (encode_descriptor(SBO) << 32ULL) | (1ULL << 46ULL);
            };

            // No-swizzle shared memory descriptor for scale factors.
            // SBO = stride between 8-row groups = 8 rows * 16 bytes/row.
            auto make_desc_SF = [](int addr) -> uint64_t
            {
                const int SBO = 8 * 16;  // = 128 bytes
                return encode_descriptor(addr) | (encode_descriptor(SBO) << 32ULL) | (1ULL << 46ULL);
            };

            // Copy scale factors from shared memory to tensor memory.
            // tcgen05.cp and tcgen05.mma are pipelined correctly per PTX docs.
            for (int k{0}; k < BLOCK_K / MMA_K; ++k) {
                uint64_t sfa_desc = make_desc_SF(SFA_smem + k * 512);  // 512 bytes per MMA_K atom
                uint64_t sfb_desc = make_desc_SF(SFB_smem + k * 512);
                copy_sf_smem2tmem(SFA_tmem + k * 4, sfa_desc);
                copy_sf_smem2tmem(SFB_tmem + k * 4, sfb_desc);
            }

            // Perform MMA for each MMA_K chunk within BLOCK_K.
            for (int k{0}; k < BLOCK_K / MMA_K; ++k)
            {
                // A tile offset: k * BLOCK_M * (MMA_K/2) bytes
                // B tile offset: k * BLOCK_N * (MMA_K/2) bytes
                uint64_t a_desc = make_desc_A(A_smem + k * BLOCK_M * (MMA_K / 2));
                uint64_t b_desc = make_desc_B(B_smem + k * BLOCK_N * (MMA_K / 2));

                const int scale_A_tmem = SFA_tmem + k * 4;
                const int scale_B_tmem = SFB_tmem + k * 4;

                // First MMA of first k-iter: clear accumulator (enable_input_d=0).
                // All subsequent: accumulate (enable_input_d=nonzero).
                const int enable_input_d = (k == 0) ? iter_k : 1;
                run_mma_nvfp4(a_desc, b_desc, i_desc, scale_A_tmem, scale_B_tmem, enable_input_d);
            }

            // Signal MMA completion on the mbarrier.
            asm volatile("tcgen05.commit.cta_group::1.mbarrier::arrive::one.shared::cluster.b64 [%0];"
                ::"r"(mbar_addr) : "memory");
        }

        // Wait for MMA to complete.
        mbarrier_wait(mbar_addr, phase);
        phase ^= 1;
    }

    // === Epilogue: Read accumulator from TMEM and store to global memory ===
    // PTX docs require this fence before tcgen05.ld, after tcgen05.mma.
    asm volatile("tcgen05.fence::after_thread_sync;");

    // Each thread handles one row. 4 warps * 32 threads = 128 rows = BLOCK_M.
    // Load 8 columns at a time from TMEM.
    for (int n{0}; n < BLOCK_N / 8; ++n) {
        float tmp[8];
        // TMEM address: 16 MSBs = row offset, 16 LSBs = column offset.
        const int addr = taddr + ((warp_idx * 32) << 16) + (n * 8);
        asm volatile("tcgen05.ld.sync.aligned.32x32b.x8.b32 {%0, %1, %2, %3, %4, %5, %6, %7}, [%8];"
                    : "=f"(tmp[0]), "=f"(tmp[1]), "=f"(tmp[2]), "=f"(tmp[3]),
                      "=f"(tmp[4]), "=f"(tmp[5]), "=f"(tmp[6]), "=f"(tmp[7])
                    : "r"(addr));
        asm volatile("tcgen05.wait::ld.sync.aligned;");

        // Convert f32 pairs to f16 pairs and write to global memory.
        half2 out[4];
        for (int i{0}; i < 4; ++i)
            out[i] = __float22half2_rn({tmp[i * 2], tmp[i * 2 + 1]});

        // Each thread writes 16 bytes (8 half values) to its row.
        half *out_ptr = C + (offset_m + thread_idx) * N + (offset_n + n * 8);
        reinterpret_cast<int4 *>(out_ptr)[0] = reinterpret_cast<int4 *>(out)[0];
    }
    __syncthreads();

    if (warp_idx == 0) {
        // Deallocate TMEM (accumulator + scale factor columns).
        asm volatile("tcgen05.dealloc.cta_group::1.sync.aligned.b32 %0, %1;" :: "r"(taddr), "r"(TMEM_COLS));
    }
}


torch::Tensor launch_kernel_naive_01(
    torch::Tensor A,
    torch::Tensor B,
    torch::Tensor sfa,
    torch::Tensor sfb,
    int M,
    int N,
    int K
) {
    auto C = torch::empty({M, N}, torch::dtype(torch::kFloat16).device(A.device()));

    // Tile sizes.
    // BLOCK_M=128; 1 CTA for .kind::mxf4nvf4.
    // BLOCK_N=128; ensure one SF atom covers all N rows in the tile.
    // BLOCK_K= 64; same as MMA_K for now.
    constexpr int BLOCK_M{128};
    constexpr int BLOCK_N{128};
    constexpr int BLOCK_K{64};
    constexpr int NUM_THREADS{4 * WARP_SIZE};  // 4 warps, 128 threads

    auto A_ptr{reinterpret_cast<const char *>(A.data_ptr())};
    auto B_ptr{reinterpret_cast<const char *>(B.data_ptr())};
    auto SFA_ptr{reinterpret_cast<const char *>(sfa.data_ptr())};
    auto SFB_ptr{reinterpret_cast<const char *>(sfb.data_ptr())};

    // Create 3D TMA tensor maps for A and B.
    // NUM_ELEMENTS=32: T = 128-bit / 4-bit = 32 (canonical atom width for FP4).
    CUtensorMap A_tmap{}, B_tmap{};
    create_tmap_descriptor<32>(&A_tmap, A_ptr, M, K, BLOCK_M, BLOCK_K,
        CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);
    create_tmap_descriptor<32>(&B_tmap, B_ptr, N, K, BLOCK_N, BLOCK_K,
        CUtensorMapSwizzle::CU_TENSOR_MAP_SWIZZLE_NONE);

    // Shared memory: A tile + B tile + SFA tile + SFB tile.
    constexpr int AB_SHARED_SIZE{(BLOCK_M + BLOCK_N) * BLOCK_K / 2};
    constexpr int SF_SHARED_SIZE{2 * 512}; // Each SF atom is 512 bytes.
    constexpr int SHARED_SIZE{AB_SHARED_SIZE + SF_SHARED_SIZE};

    dim3 num_threads(NUM_THREADS);
    dim3 num_blocks((M / BLOCK_M) * (N / BLOCK_N));

    auto kernel{kernel_v01_naive<NUM_THREADS, BLOCK_M, BLOCK_N, BLOCK_K>};

    if (SHARED_SIZE > 48'000)
        cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, SHARED_SIZE);

    kernel<<<num_blocks, num_threads, SHARED_SIZE>>>(
        A_tmap,
        B_tmap,
        SFA_ptr,
        SFB_ptr,
        reinterpret_cast<half *>(C.data_ptr<at::Half>()),
        M, N, K
    );
    return C;
}

"""

cpp_source = """
#include <torch/extension.h>

torch::Tensor launch_kernel_naive_01(
    torch::Tensor a_bytes,
    torch::Tensor b_bytes,
    torch::Tensor sfa,
    torch::Tensor sfb,
    int M,
    int N,
    int K);
"""

module = load_inline(
    name='kernel',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['launch_kernel_naive_01'],
    verbose=True,
    is_python_module=True,
    no_implicit_headers=True,
    extra_cuda_cflags=[
        "-O3",
        "-gencode=arch=compute_100a,code=sm_100a",
        "--use_fast_math",
        "--expt-relaxed-constexpr",
        "--relocatable-device-code=false",
        "-lineinfo",
        "-Xptxas=-v",
        # "--keep",
        # "--keep-dir",
        # f"{Path(__file__).parent}/tmp",
    ],
    extra_ldflags=["-lcuda"],
)


def custom_kernel(data: input_t) -> output_t:
    abc_tensors, _, sfasfb_reordered_tensors, problem_sizes = data

    BLOCK_M = 128
    BLOCK_N = 128

    results = []
    for (a, b, c), (sfa_reordered, sfb_reordered), (m, n, k, l) in zip(
        abc_tensors, sfasfb_reordered_tensors, problem_sizes
    ):
        for l_idx in range(l):
            a_ptr = a[:, :, l_idx].contiguous()  # [m, k//2]
            b_ptr = b[:, :, l_idx].contiguous()  # [n, k//2]

            # Pad M and N to multiples of BLOCK_M and BLOCK_N.
            padded_m = ((m + BLOCK_M - 1) // BLOCK_M) * BLOCK_M
            padded_n = ((n + BLOCK_N - 1) // BLOCK_N) * BLOCK_N

            if padded_m != m:
                a_padded = torch.zeros(padded_m, k // 2, dtype=torch.uint8, device=a_ptr.device).view(a_ptr.dtype)
                a_padded[:m] = a_ptr
                a_ptr = a_padded

            if padded_n != n:
                b_padded = torch.zeros(padded_n, k // 2, dtype=torch.uint8, device=b_ptr.device).view(b_ptr.dtype)
                b_padded[:n] = b_ptr
                b_ptr = b_padded

            c_out = module.launch_kernel_naive_01(
                a_ptr, b_ptr, sfa_reordered, sfb_reordered, padded_m, padded_n, k
            )

            # Trim back to original size.
            c[:, :, l_idx] = c_out[:m, :n]

        results.append(c)

    return results