leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity
Files
291300ffffa3533a78ee104f08a8490a29ce9ccb
cutlass/examples/94_ada_fp8_blockwise/ada_fp8_blockwise.cu
Junkai-Wu 9fba3195f9 v4.4 update. (#2979)
2026-01-24 11:46:17 -05:00

490 lines
18 KiB
Plaintext
Raw Blame History

/***************************************************************************************************
* Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: BSD-3-Clause
*
* Redistribution and use in source and binary forms, with or without
* modification, are permitted provided that the following conditions are met:
*
* 1. Redistributions of source code must retain the above copyright notice, this
* list of conditions and the following disclaimer.
*
* 2. Redistributions in binary form must reproduce the above copyright notice,
* this list of conditions and the following disclaimer in the documentation
* and/or other materials provided with the distribution.
*
* 3. Neither the name of the copyright holder nor the names of its
* contributors may be used to endorse or promote products derived from
* this software without specific prior written permission.
*
* THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
* AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
* IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
* DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
* FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
* DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
* SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
* CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
*
**************************************************************************************************/
/*! \file
\brief An FP8 blockwise scaled GEMM example for the NVIDIA Ada SM89 architecture using CUTLASS.
*/
#include "cutlass/arch/arch.h"
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/thread/linear_combination.h"
#include "cutlass/gemm/device/gemm_blockwise.h"
#include "cutlass/gemm/threadblock/mma_multistage_blockwise.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/numeric_types.h"
#include "cutlass/tensor_ref.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/device_memory.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/reference/host/gemm.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/util/reference/host/tensor_copy.h"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/tensor_norm.h"
#include "helper.h"
#include <algorithm>
#include <cmath>
#include <cstdio>
#include <cuda_runtime.h>
#include <iostream>
#include <random>
#include <vector>
#include "cutlass/util/command_line.h"
using cutlass::ceil_div;
using ElementA = cutlass::float_e4m3_t;
using ElementB = cutlass::float_e4m3_t;
using ElementOutput = cutlass::bfloat16_t;
using ElementAccumulator = float;
using ElementScale = float;
using LayoutA = cutlass::layout::RowMajor;
using LayoutB = cutlass::layout::ColumnMajor;
using LayoutC = cutlass::layout::RowMajor; // Currently only RowMajor is supported
using LayoutScale = cutlass::layout::RowMajor; // Currently only RowMajor is supported
static int const AlignmentA = 16;
static int const AlignmentB = 16;
static int const Stages = 3;
static int const BlockSize = 128;
constexpr float epsilon = 0.51f;
constexpr float floor_val = 1.0f;
// -----------------------------------------------------------------------------
// Command line options structure
// -----------------------------------------------------------------------------
struct Options {
bool help = false;
int m = 1024;
int n = 1024;
int k = 1024;
float alpha = 1.f;
float beta = 0.f;
bool verify = true;
int iterations = 1000;
int warmup = 1000;
// Parse command line arguments using CUTLASS helper
void parse(int argc, char const **argv) {
cutlass::CommandLine cmd(argc, argv);
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
cmd.get_cmd_line_argument("m", m);
cmd.get_cmd_line_argument("n", n);
cmd.get_cmd_line_argument("k", k);
cmd.get_cmd_line_argument("alpha", alpha, alpha);
cmd.get_cmd_line_argument("beta", beta, beta);
cmd.get_cmd_line_argument("verify", verify);
cmd.get_cmd_line_argument("iterations", iterations);
cmd.get_cmd_line_argument("warmup", warmup);
}
std::ostream &print_usage(std::ostream &out) const {
out << "94_ada_fp8_blockwise\n\n"
<< " FP8 GEMM with blockwise scaling (Ada, Sm89).\n\n"
<< "Options:\n\n"
<< " --help Display this help string\n"
<< " --m=<int> GEMM M dimension (default 1024)\n"
<< " --n=<int> GEMM N dimension (default 1024)\n"
<< " --k=<int> GEMM K dimension (default 1024)\n"
<< " --alpha=<f32> Epilogue alpha (default 1.0)\n"
<< " --beta=<f32> Epilogue beta (default 0.0)\n"
<< " --verify=<bool> Verify the results (default true)\n"
<< " --iterations=<int> Number of timing iterations (default 1000)\n"
<< " --warmup=<int> Number of warmup iterations (default 1000)\n";
return out;
}
double gflops(double runtime_s) const {
uint64_t flop = uint64_t(2) * m * n * k;
double gflop = double(flop) / double(1.0e9);
return gflop / runtime_s;
}
};
using EpilogueOutputOp = cutlass::epilogue::thread::LinearCombination<
ElementOutput, 8, ElementAccumulator, ElementAccumulator>;
using Gemm = cutlass::gemm::device::GemmBlockwise<
ElementA, LayoutA, ElementB, LayoutB, ElementOutput, LayoutC,
ElementAccumulator, cutlass::arch::OpClassTensorOp, cutlass::arch::Sm89,
cutlass::gemm::GemmShape<64, 128, 128>,
cutlass::gemm::GemmShape<64, 64, 128>, cutlass::gemm::GemmShape<16, 8, 32>,
ElementScale, LayoutScale, BlockSize, EpilogueOutputOp,
cutlass::gemm::threadblock::GemmIdentityThreadblockSwizzle<>, Stages,
AlignmentA, AlignmentB, false, cutlass::arch::OpMultiplyAdd>;
// Host-side verification
static bool verify_gemm(int M, int N, int K, ElementA const *A,
ElementB const *B, ElementOutput const *C,
ElementOutput const *D, ElementScale const *scale_A,
ElementScale const *scale_B, float alpha, float beta) {
std::vector<ElementA> A_fp8_host(M * K);
std::vector<ElementB> B_fp8_host(K * N);
std::vector<ElementOutput> C_bf16_host(M * N);
std::vector<ElementOutput> D_bf16_host(M * N);
int kBlocks = ceil_div(K, BlockSize);
int mBlocks = ceil_div(M, BlockSize);
int nBlocks = ceil_div(N, BlockSize);
std::vector<float> ScaleA(mBlocks * kBlocks);
std::vector<float> ScaleB(nBlocks * kBlocks);
CUDA_CHECK(
cudaMemcpy(A_fp8_host.data(), A,
sizeof(ElementA) * A_fp8_host.size(),
cudaMemcpyDeviceToHost));
CUDA_CHECK(
cudaMemcpy(B_fp8_host.data(), B,
sizeof(ElementB) * B_fp8_host.size(),
cudaMemcpyDeviceToHost));
CUDA_CHECK(
cudaMemcpy(C_bf16_host.data(), C,
sizeof(ElementOutput) * C_bf16_host.size(),
cudaMemcpyDeviceToHost));
CUDA_CHECK(
cudaMemcpy(D_bf16_host.data(), D,
sizeof(ElementOutput) * D_bf16_host.size(),
cudaMemcpyDeviceToHost));
CUDA_CHECK(
cudaMemcpy(ScaleA.data(), scale_A,
sizeof(ElementScale) * ScaleA.size(),
cudaMemcpyDeviceToHost));
CUDA_CHECK(
cudaMemcpy(ScaleB.data(), scale_B,
sizeof(ElementScale) * ScaleB.size(),
cudaMemcpyDeviceToHost));
// --------------------------------------------------------------------
// De-quantize A and B into FP32 using blockwise scales
// --------------------------------------------------------------------
std::vector<float> A_f32(M * K);
for (int m = 0; m < M; ++m) {
int blk_m = m / BlockSize;
for (int k = 0; k < K; ++k) {
int blk_k = k / BlockSize;
float sA = ScaleA[blk_m * kBlocks + blk_k];
A_f32[m * K + k] = static_cast<float>(A_fp8_host[m * K + k]) * sA;
}
}
std::vector<float> B_f32(K * N);
for (int n = 0; n < N; ++n) {
int blk_n = n / BlockSize;
for (int k = 0; k < K; ++k) {
int blk_k = k / BlockSize;
float sB = ScaleB[blk_n * kBlocks + blk_k];
B_f32[k + n * K] = static_cast<float>(B_fp8_host[k + n * K]) * sB;
}
}
// --------------------------------------------------------------------
// Prepare tensor refs and run reference GEMM
// --------------------------------------------------------------------
cutlass::TensorRef<float, LayoutA> A_ref(A_f32.data(), K);
cutlass::TensorRef<float, LayoutB> B_ref(B_f32.data(), K);
// Convert C (BF16) to FP32 once
std::vector<float> C_f32(M * N);
cutlass::NumericConverter<float, ElementOutput> bf16_to_f32;
std::transform(C_bf16_host.begin(), C_bf16_host.end(), C_f32.begin(),
[&](ElementOutput x) { return bf16_to_f32(x); });
cutlass::TensorRef<float, LayoutC> C_ref(C_f32.data(), N);
// Output buffer in FP32 (will later hold ReLU result)
std::vector<float> D_ref_f32(M * N, 0.0f);
cutlass::TensorRef<float, LayoutC> D_ref_tensor(D_ref_f32.data(), N);
cutlass::gemm::GemmCoord problem{M, N, K};
cutlass::reference::host::Gemm<float, LayoutA, float, LayoutB, float,
LayoutC, float, float,
cutlass::arch::OpMultiplyAdd>
gemm_ref;
gemm_ref(problem, alpha, A_ref, B_ref,
beta, C_ref, D_ref_tensor, 0.0f);
// Build tensor views for relative comparison
cutlass::TensorView<float, LayoutC> ref_view(D_ref_f32.data(), LayoutC(N), {M, N});
// Convert device output (BF16) to float for comparison
std::vector<float> kernel_output(M * N);
cutlass::TensorView<float, LayoutC> kernel_output_view(kernel_output.data(), LayoutC(N), {M, N});
cutlass::TensorView<ElementOutput, LayoutC> d_bf16_view(D_bf16_host.data(), LayoutC(N), {M, N});
cutlass::reference::host::TensorCopy(kernel_output_view, d_bf16_view);
bool result = cutlass::reference::host::TensorRelativelyEquals(
ref_view, kernel_output_view, epsilon, floor_val);
// std::cout << "ref:\n" << ref_view << "\n\n\n";
// std::cout << "compute:\n" << kernel_output_view << "\n\n\n";
// Compute error metrics
double mse = cutlass::reference::host::TensorMSE(kernel_output_view, ref_view);
double mre = cutlass::reference::host::TensorMRE(kernel_output_view, ref_view);
double max_error = cutlass::reference::host::TensorGreatestError(kernel_output_view, ref_view);
std::cout << " Result MSE: " << mse << ", MRE: " << mre << ", greatest error: " << max_error << std::endl;
std::cout << "GEMM Verification result is " << (result ? "PASS" : "FAIL")
<< std::endl;
return result;
}
static void initialize_tensors(
int M, int N, int K, cutlass::HostTensor<ElementA, LayoutA> &tensor_A,
cutlass::HostTensor<ElementB, LayoutB> &tensor_B,
cutlass::HostTensor<ElementOutput, LayoutC> &tensor_C,
cutlass::HostTensor<ElementOutput, LayoutC> &tensor_D,
cutlass::HostTensor<float, Gemm::LayoutScale> &tensor_ScaleA,
cutlass::HostTensor<float, Gemm::LayoutScale> &tensor_ScaleB) {
int mBlocks = ceil_div(M, BlockSize);
int nBlocks = ceil_div(N, BlockSize);
int kBlocks = ceil_div(K, BlockSize);
uint64_t seed = 2024;
// Resize tensors ------------------------------------------------------
tensor_A.resize({M, K});
tensor_B.resize({K, N});
tensor_C.resize({M, N});
tensor_D.resize({M, N});
tensor_ScaleA.resize({mBlocks, kBlocks});
tensor_ScaleB.resize({nBlocks, kBlocks});
// Fill A and B with random uniform values in [-2, 2]
cutlass::reference::host::TensorFillRandomUniform(
tensor_A.host_view(), seed + 1, 2.0, -2.0);
cutlass::reference::host::TensorFillRandomUniform(
tensor_B.host_view(), seed + 2, 2.0, -2.0);
// Fill C and D with random uniform values in [-2, 2]
cutlass::reference::host::TensorFillRandomUniform(
tensor_C.host_view(), seed + 3, 2.0, -2.0);
cutlass::reference::host::TensorFillRandomUniform(
tensor_D.host_view(), seed + 4, 2.0, -2.0);
// Fill scale tensors with random uniform values in [-1, 1]
cutlass::reference::host::TensorFillRandomUniform(
tensor_ScaleA.host_view(), seed + 5, 1.0, -1.0);
cutlass::reference::host::TensorFillRandomUniform(
tensor_ScaleB.host_view(), seed + 6, 1.0, -1.0);
tensor_A.sync_device();
tensor_B.sync_device();
tensor_C.sync_device();
tensor_D.sync_device();
tensor_ScaleA.sync_device();
tensor_ScaleB.sync_device();
}
// This example requires CUDA 12.4 or greater and sm89 or higher
bool sufficient() {
if (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 4)) {
std::cerr << "This example requires CUDA 12.4 or greater." << std::endl;
return false;
}
size_t smem_size = sizeof(typename Gemm::GemmKernel::SharedStorage);
cudaDeviceProp properties;
int device_idx;
cudaError_t result = cudaGetDevice(&device_idx);
if (result != cudaSuccess) {
std::cerr << "cudaGetDevice() failed with error: " << cudaGetErrorString(result) << std::endl;
return false;
}
result = cudaGetDeviceProperties(&properties, device_idx);
if (result != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() failed with error: " << cudaGetErrorString(result) << std::endl;
return false;
}
if (properties.major < 8 || (properties.major == 8 && properties.minor < 9)) {
std::cerr << "CUTLASS's Ada FP8 BlockwiseGEMM example requires a device of compute capability 89 or higher.\n" << std::endl;
return false;
}
if (properties.sharedMemPerBlockOptin < smem_size) {
std::cerr << "Insufficient shared memory. Need " << smem_size
<< ", but device only has " << properties.sharedMemPerBlockOptin << std::endl;
return false;
}
return true;
}
int main(int argc, char const **args) {
// ---------------------------------------------------------------------------
// Parse command-line options
// ---------------------------------------------------------------------------
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
// Problem dimensions
const int M = options.m;
const int N = options.n;
const int K = options.k;
// Waive test if insufficient CUDA device
if (!sufficient()) {
std::cerr << "Insufficient resources to run the kernel." << std::endl;
return 0;
}
cutlass::HostTensor<ElementA, LayoutA> tensor_A;
cutlass::HostTensor<ElementB, LayoutB> tensor_B;
cutlass::HostTensor<ElementOutput, LayoutC> tensor_C;
cutlass::HostTensor<ElementOutput, LayoutC> tensor_D;
cutlass::HostTensor<float, Gemm::LayoutScale> tensor_ScaleA;
cutlass::HostTensor<float, Gemm::LayoutScale> tensor_ScaleB;
initialize_tensors(M, N, K, tensor_A, tensor_B, tensor_C, tensor_D, tensor_ScaleA, tensor_ScaleB);
ElementA const *ptr_A = tensor_A.device_data();
ElementB const *ptr_B = tensor_B.device_data();
ElementOutput const *ptr_C = tensor_C.device_data();
ElementOutput *ptr_D = tensor_D.device_data();
float *a_ptr = tensor_ScaleA.device_data();
float *b_ptr = tensor_ScaleB.device_data();
typename Gemm::EpilogueOutputOp::Params epilogue_params(options.alpha, options.beta);
int kBlocks = ceil_div(K, BlockSize);
int ldA = kBlocks;
int ldB = kBlocks;
using LayoutScale = Gemm::LayoutScale;
using TensorRefScale = Gemm::TensorRefScale;
LayoutScale layout_A(ldA);
LayoutScale layout_B(ldB);
TensorRefScale ref_scale_A(a_ptr, layout_A);
TensorRefScale ref_scale_B(b_ptr, layout_B);
// Construct the argument list ----------------------------------------
typename Gemm::Arguments arguments(
/* problem_size */ {M, N, K},
/* A */ {ptr_A, K},
/* B */ {ptr_B, K},
/* C */ {ptr_C, N},
/* D */ {ptr_D, N},
/* scale_A tensor ref */ ref_scale_A,
/* scale_B tensor ref */ ref_scale_B,
/* epilogue params */ epilogue_params,
/* split-k slices */ 1,
/* gather A indices */ nullptr,
/* gather B indices */ nullptr,
/* scatter D indices */ nullptr);
Gemm gemm_op;
cutlass::Status status = gemm_op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "GEMM cannot implement the given problem." << std::endl;
return -1;
}
size_t workspace_bytes = Gemm::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_bytes);
status = gemm_op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "GEMM initialization failed." << std::endl;
return -1;
}
status = gemm_op();
if (status != cutlass::Status::kSuccess) {
std::cerr << "GEMM execution failed." << std::endl;
return -1;
}
// Verification
if (options.verify) {
std::cout << "Verifying GEMM" << std::endl;
bool ok = verify_gemm(M, N, K, ptr_A, ptr_B, ptr_C, ptr_D, a_ptr, b_ptr,
options.alpha, options.beta);
if (!ok) {
std::cerr << "Verification failed." << std::endl;
return -1;
}
}
// Profiling loop
if (options.iterations > 0) {
GpuTimer timer;
for (int iter = 0; iter < options.warmup + options.iterations; ++iter) {
if (iter == options.warmup)
timer.start();
CUTLASS_CHECK(gemm_op.run());
}
timer.stop();
float elapsed_ms = timer.elapsed_millis();
double avg_runtime_ms = double(elapsed_ms) / double(options.iterations);
double gflops = options.gflops(avg_runtime_ms / 1000.0);
std::cout << "Avg runtime: " << avg_runtime_ms << " ms" << std::endl;
std::cout << "GFLOPS: " << gflops << std::endl;
fflush(stdout);
}
return 0;
}
Reference in New Issue View Git Blame Copy Permalink