Synthetic benchmark from long gemini session for >170 TFLOPS FP8

benchmark_fp8_ptx.cu

#include <iostream>
#include <cstdint>
#include <cuda_runtime.h>

#define CUDA_CHECK(status) \
    do { \
        cudaError_t err = (status); \
        if (err != cudaSuccess) { \
            std::cerr << "[CUDA Error] " << cudaGetErrorString(err) \
                      << " (" << err << ") at line " << __LINE__ << std::endl; \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

#define UNROLL_FACTOR 32

__global__ void __launch_bounds__(256, 2) fp8_ptx_peak_kernel_optimized(int64_t iterations) {
    uint32_t a[4] = {0x3c003c00, 0x3c003c00, 0x3c003c00, 0x3c003c00};
    uint32_t b[2] = {0x3c003c00, 0x3c003c00};
    
    // 4 independent sets of accumulators to break latency dependency chain
    float c0[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    float c1[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    float c2[4] = {0.0f, 0.0f, 0.0f, 0.0f};
    float c3[4] = {0.0f, 0.0f, 0.0f, 0.0f};

    for (int64_t i = 0; i < iterations; ++i) {
        #pragma unroll
        for (int u = 0; u < UNROLL_FACTOR; ++u) {
            asm volatile (
                "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%0,%1,%2,%3};\n"
                "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%10,%11,%12,%13}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
                "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%14,%15,%16,%17}, {%4,%5,%6,%7}, {%8,%9}, {%14,%15,%16,%17};\n"
                "mma.sync.aligned.kind::f8f6f4.m16n8k32.row.col.f32.e4m3.e4m3.f32 {%18,%19,%20,%21}, {%4,%5,%6,%7}, {%8,%9}, {%18,%19,%20,%21};\n"
                : "+f"(c0[0]), "+f"(c0[1]), "+f"(c0[2]), "+f"(c0[3]),
                  "+f"(c1[0]), "+f"(c1[1]), "+f"(c1[2]), "+f"(c1[3]),
                  "+f"(c2[0]), "+f"(c2[1]), "+f"(c2[2]), "+f"(c2[3]),
                  "+f"(c3[0]), "+f"(c3[1]), "+f"(c3[2]), "+f"(c3[3])
                : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]),
                  "r"(b[0]), "r"(b[1])
            );
        }
    }

    if (threadIdx.x == 0 && c0[0] == 12345.0f) {
        printf("Sink: %f\n", c0[0] + c1[0] + c2[0] + c3[0]);
    }
}

int main() {
    CUDA_CHECK(cudaFree(0));

    int device = 0;
    cudaDeviceProp prop;
    CUDA_CHECK(cudaGetDeviceProperties(&prop, device));

    std::cout << "Benchmarking Device: " << prop.name << " (" << prop.multiProcessorCount << " SMs)" << std::endl;

    int threads_per_block = 256;
    int num_blocks = prop.multiProcessorCount * 8; // Higher occupancy
    int64_t iterations = 50000;

    // Warmup
    fp8_ptx_peak_kernel_optimized<<<num_blocks, threads_per_block>>>(iterations / 10);
    CUDA_CHECK(cudaDeviceSynchronize());

    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop));

    CUDA_CHECK(cudaEventRecord(start));
    fp8_ptx_peak_kernel_optimized<<<num_blocks, threads_per_block>>>(iterations);
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));

    float ms = 0.0f;
    CUDA_CHECK(cudaEventElapsedTime(&ms, start, stop));

    int total_warps = (num_blocks * threads_per_block) / 32;
    // 4 MMA instructions per unroll loop step
    double total_mma_ops = (double)total_warps * iterations * UNROLL_FACTOR * 4.0;
    double total_flops = total_mma_ops * 8192.0;

    double seconds = ms / 1000.0;
    double tflops = (total_flops / seconds) / 1e12;

    std::cout << "Execution Time: " << ms << " ms" << std::endl;
    std::cout << "Raw Peak FP8 Tensor Core Performance: " << tflops << " TFLOPS" << std::endl;

    CUDA_CHECK(cudaEventDestroy(start));
    CUDA_CHECK(cudaEventDestroy(stop));

    return 0;
}