leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 24 Packages Projects Releases Wiki Activity

v4.4 update. (#2979)

Browse Source
This commit is contained in:
Junkai-Wu
2026-01-24 11:46:17 -05:00
committed by GitHub
parent 2fafefb7b9
commit 9fba3195f9
293 changed files with 46343 additions and 2995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/111_hopper_ssd/111_hopper_ssd.cu
+849
View File
@@ -0,0 +1,849 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#include <iostream>
#include "cutlass/util/command_line.h"
#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cute/layout.hpp"
#include "cutlass/kernel_hardware_info.hpp"
#include "thrust/universal_vector.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/tensor_copy.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
#include "reference/reference_ssd_cumsum.hpp"
#include "reference/reference_ssd.hpp"
#include "cutlass/transform/device/transform_universal_adapter.hpp"
#include "device/ssd.hpp"
#include "kernel/sm90_ssd_kernel_builder.hpp"
using namespace cute;
// Command line options parsing
struct Options {
using Element = cutlass::bfloat16_t;
using ElementAcc = float;
using ElementDA = float;
static constexpr bool D_HAS_HDIM = true;
static constexpr bool HAS_D = true;
static constexpr bool HAS_Z = true;
bool help;
bool error;
// All static number now
int G = 2;
int B = 3;
int E = 2;
int H = 2;
// Reference kernel doesn't support dynamic C now.
static constexpr auto C = Int<8>{};
static constexpr auto D = Int<64>{};
static constexpr auto L = Int<128>{};
static constexpr auto N = Int<128>{};
int EH = E * H;
int iterations;
bool verify;
bool verbose;
int warmups;
bool measure;
Options():
help(false),
error(false),
iterations(1), verify(true),
measure(false), warmups(3)
{}
// Parses the command line
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
Options defaults;
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
cmd.get_cmd_line_argument("iterations", iterations, defaults.iterations);
cmd.get_cmd_line_argument("G", G, defaults.G);
cmd.get_cmd_line_argument("B", B, defaults.B);
cmd.get_cmd_line_argument("E", E, defaults.E);
cmd.get_cmd_line_argument("H", H, defaults.H);
verbose = cmd.check_cmd_line_flag("verbose");
verify = !(cmd.check_cmd_line_flag("without_verify"));
EH = E*H;
if (iterations > 1) {
measure = true;
verbose = true;
}
auto problem_shape = cute::make_tuple(G, B, EH, C, L, D, N);
cute::print("problem_shape : "); cute::print(problem_shape); cute::print("\n");
}
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "111_hopper_ssd\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --iterations=<int> Benchmarking iterations.\n"
<< " --without_verify Don't verify the results.\n"
<< " --verbose Print execution time per kernel\n"
<< " --G=<int> Group\n"
<< " --B=<int> Batch\n"
<< " --E=<int> Expanded factor\n"
<< " --H=<int> Number of heads\n"
<< "\n";
return out;
}
auto get_problem_shape() const {
return cute::make_tuple(G, B, EH, C, L, D, N);
}
// acceptable layout by cuDNN
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
auto layoutX() const {
auto layout = make_layout(make_shape(L, C, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutDelta() const {
auto layout = make_layout(make_shape(L, C, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutDeltaA() const {
auto layout = make_layout(make_shape(L, C, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutB() const {
auto layout = make_layout(make_shape(L, C, N, G, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutC() const {
auto layout = make_layout(make_shape(L, C, N, G, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutY() const {
auto layout = make_layout(make_shape(L, C, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutF() const {
auto layout = make_layout(make_shape(N, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutD() const {
if constexpr (D_HAS_HDIM) {
auto layout = make_layout(make_shape(D, EH));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
else {
auto layout = make_layout(make_shape(Int<1>{}, EH));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
}
auto layoutZ() const {
auto layout = make_layout(make_shape(L, C, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
// transformed layout for kernel parameters
auto layoutX_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B));
return make_layout(
make_shape(D,L,int32_t(C),EH*B),
make_stride(
stride<2>(layout),
stride<0>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutB_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),N,G*B));
return make_layout(
make_shape(L,N,int32_t(C),G*B),
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutC_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),N,G*B));
return make_layout(
make_shape(L,N,int32_t(C),G*B),
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutDelta_transformed() const {
return make_layout(make_shape(L,int32_t(C),EH*B));
}
auto layoutY_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B));
return make_layout(
make_shape(L,D,int32_t(C),EH*B), // (M,K,L,...)
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutF_transformed() const {
auto layout = make_layout(make_shape(N,D,EH*B));
return make_layout(
make_shape(D,N,EH*B),
make_stride(
stride<1>(layout),
stride<0>(layout),
stride<2>(layout)
)
);
}
auto layoutD_transformed() const {
if constexpr (D_HAS_HDIM) {
return make_layout(make_shape(D, EH));
}
else {
return make_layout(make_shape(Int<1>{}, EH));
}
}
auto layoutZ_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B));
return make_layout(
make_shape(L,D,int32_t(C),EH*B),
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
};
template <typename Element>
static void
initialize_values(
thrust::universal_vector<Element>& dst_ptr,
cutlass::Distribution::Kind dist_kind,
uint64_t seed,
Element var = Element(1.f)) {
if (cutlass::Distribution::Uniform == dist_kind) {
int scope = 2;
cutlass::reference::host::BlockFillRandomUniform(
dst_ptr.data().get(), dst_ptr.size(), seed, scope, -scope, 0);
}
else if (cutlass::Distribution::AllZeros == dist_kind) {
cutlass::reference::host::BlockFillRandomUniform(
dst_ptr.data().get(), dst_ptr.size(), seed, 0, 0, 0);
}
else if (cutlass::Distribution::AllOnes == dist_kind) {
cutlass::reference::host::BlockFillRandomUniform(
dst_ptr.data().get(), dst_ptr.size(), seed, 1, 1, 0);
}
else if (cutlass::Distribution::Gaussian == dist_kind) {
cutlass::reference::device::BlockFillRandomGaussian(
dst_ptr.data().get(), dst_ptr.size(), seed, (Element) 0, var);
}
else if (cutlass::Distribution::Sequential == dist_kind) {
cutlass::reference::host::BlockFillSequential(dst_ptr.data().get(), dst_ptr.size());
}
else {
std::cerr << "Invalid distribution kind!\n.";
exit(1);
}
}
template <
class Options_
>
struct TestBed {
using Option = Options_;
using Element = typename Option::Element;
using ElementDA = typename Option::ElementDA;
using ElementAcc = typename Option::ElementAcc;
thrust::universal_vector<Element> tensor_X;
thrust::universal_vector<Element> tensor_DeltaA;
thrust::universal_vector<ElementDA> tensor_DeltaA_cumsum;
thrust::universal_vector<Element> tensor_Delta;
thrust::universal_vector<Element> tensor_B;
thrust::universal_vector<Element> tensor_C;
thrust::universal_vector<Element> tensor_D;
thrust::universal_vector<Element> tensor_Y;
thrust::universal_vector<Element> tensor_Z;
thrust::universal_vector<Element> tensor_Y_ref_0;
thrust::universal_vector<Element> tensor_Y_ref_1;
thrust::universal_vector<Element> tensor_F;
thrust::universal_vector<Element> tensor_F_ref_0;
thrust::universal_vector<Element> tensor_F_ref_1;
cutlass::Distribution::Kind init_X = cutlass::Distribution::Uniform;
cutlass::Distribution::Kind init_DeltaA = cutlass::Distribution::Gaussian;
cutlass::Distribution::Kind init_Delta = cutlass::Distribution::Gaussian;
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform;
cutlass::Distribution::Kind init_C = cutlass::Distribution::Uniform;
using TileShape = decltype(make_shape(Options::L, Options::D, Options::N)); // (L, D, N)
using SsdOperation = cutlass::ssd::device::SSD<
typename cutlass::ssd::kernel::Sm90SsdBuilder<
Element, ElementDA, ElementAcc, Element,
TileShape,
Option::HAS_D, Option::D_HAS_HDIM, Option::HAS_Z
>::Kernel>;
using CumsumKenrel = cutlass::ssd::kernel::CumsumKernel<Element, ElementDA, TileShape>;
using CumsumOperation = cutlass::transform::device::TransformUniversalAdapter<CumsumKenrel>;
bool initialize(Options const& options, const cutlass::KernelHardwareInfo& hw_info, uint64_t seed = 2023) {
auto [g, b, eh, c, l, d, n] = options.get_problem_shape();
assert(g == 1 && "Only group size == 1 is supported") ;
auto size_X = b * eh * c * l * d;
auto size_DeltaA = b * eh * c * l;
auto size_Delta = b * eh * c * l;
auto size_B = g * b * c * n * l;
auto size_C = g * b * c * n * l;
auto size_Y = b * eh * c * l * d;
auto size_F = b * eh * d * n;
tensor_X .resize(sizeof(Element) * size(options.layoutX()));
tensor_DeltaA .resize(sizeof(Element) * size(options.layoutDeltaA()));
tensor_Delta .resize(sizeof(Element) * size(options.layoutDelta()));
tensor_B .resize(sizeof(Element) * size(options.layoutB()));
tensor_C .resize(sizeof(Element) * size(options.layoutC()));
tensor_D .resize(sizeof(Element) * size(options.layoutD()));
tensor_Z .resize(sizeof(Element) * size(options.layoutZ()));
tensor_Y .resize(sizeof(Element) * size(options.layoutY()));
tensor_Y_ref_0.resize(sizeof(Element) * size(options.layoutY()));
tensor_Y_ref_1.resize(sizeof(Element) * size(options.layoutY()));
tensor_F .resize(sizeof(Element) * size(options.layoutF()));
tensor_F_ref_0.resize(sizeof(Element) * size(options.layoutF()));
tensor_F_ref_1.resize(sizeof(Element) * size(options.layoutF()));
tensor_DeltaA_cumsum.resize(sizeof(ElementDA) * size(options.layoutDeltaA()));
// Limit distribution to reduce skew between hosts and devices
initialize_values(tensor_X, init_X, seed);
initialize_values(tensor_DeltaA, init_DeltaA, seed + 1, Element(0.05f));
initialize_values(tensor_Delta, init_Delta, seed + 3, Element(0.05f));
initialize_values(tensor_B, init_B, seed + 5);
initialize_values(tensor_C, init_C, seed + 7);
initialize_values(tensor_D, init_C, seed + 9);
initialize_values(tensor_Z, init_X, seed);
cudaError_t result;
result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the Initialization kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
}
// apply cumsum(device) before kernel launch
typename CumsumOperation::Arguments arguments{
make_shape(int(b), int(eh), int(c), int(l)),
{
tensor_DeltaA.data().get(),
tensor_DeltaA_cumsum.data().get(),
},
hw_info
};
CumsumOperation op;
size_t workspace_size = CumsumOperation::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
// may be used uninitialized
cudaEvent_t start;
cudaEvent_t end;
cudaEventCreate(&start);
cudaEventCreate(&end);
// warm up
if (options.measure) {
for (int i = 0; i < options.warmups; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
}
result = cudaEventRecord(start);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
// Run
for (int i = 0; i < options.iterations; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
result = cudaEventRecord(end);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
float runtime_ms = 0;
result = cudaEventElapsedTime(&runtime_ms, start, end);
if (result != cudaSuccess) {
std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
runtime_ms /= static_cast<float>(options.iterations);
if (options.verbose) {
printf("[iters = %d, warmups = %d] cumsum kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms);
}
return true;
}
bool sufficient() const {
int device_idx;
cudaError_t result = cudaGetDevice(&device_idx);
if (result != cudaSuccess) {
throw std::runtime_error("cudaGetDevice() API call failed.");
}
int max_smem_size;
result = cudaDeviceGetAttribute(&max_smem_size, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_idx);
if (result != cudaSuccess) {
throw std::runtime_error("cudaDeviceGetAttribute() failed");
}
return true;
}
bool run(Options const& options, const cutlass::KernelHardwareInfo& hw_info) {
if (!sufficient()) {
std::cerr << "Test waived due to insufficient CUDA device.\n";
return true;
}
if (!initialize(options, hw_info)) {
std::cerr << "Failed to initialize the test.\n";
return true;
};
auto [g, b, eh, c, l, d, n] = options.get_problem_shape();
typename SsdOperation::Arguments arguments{
make_shape(int(g), int(b), int(eh), int(c), int(l), int(d), int(n)),
{
tensor_X.data().get(),
tensor_DeltaA_cumsum.data().get(),
tensor_Delta.data().get(),
tensor_B.data().get(),
tensor_C.data().get(),
options.layoutX_transformed(),
options.layoutB_transformed(),
options.layoutC_transformed(),
options.layoutDelta_transformed()
},
{
tensor_Y.data().get(),
tensor_F.data().get(),
tensor_D.data().get(),
tensor_Z.data().get(),
options.layoutY_transformed(),
options.layoutF_transformed(),
options.layoutD_transformed(),
options.layoutZ_transformed()
},
hw_info
};
SsdOperation op;
size_t workspace_size = SsdOperation::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
cudaError_t result;
// may be used uninitialized
cudaEvent_t start;
cudaEvent_t end;
cudaEventCreate(&start);
cudaEventCreate(&end);
// warm up
if (options.measure) {
for (int i = 0; i < options.warmups; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
}
result = cudaEventRecord(start);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
// Run
for (int i = 0; i < options.iterations; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
result = cudaEventRecord(end);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
float runtime_ms = 0;
result = cudaEventElapsedTime(&runtime_ms, start, end);
if (result != cudaSuccess) {
std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
runtime_ms /= static_cast<float>(options.iterations);
if (options.verbose) {
printf("[iters = %d, warmups = %d] ssd kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms);
printf("smem size = %d\n", SsdOperation::Kernel::SharedStorageSize);
}
// Matrix
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
auto mY_ref_0 = cute::make_tensor(tensor_Y_ref_0.data().get(), options.layoutY());
auto mY_ref_1 = cute::make_tensor(tensor_Y_ref_1.data().get(), options.layoutY());
auto mY_res = cute::make_tensor(tensor_Y.data().get(), options.layoutY());
auto mF_ref_0 = cute::make_tensor(tensor_F_ref_0.data().get(), options.layoutF());
auto mF_ref_1 = cute::make_tensor(tensor_F_ref_1.data().get(), options.layoutF());
auto mF_res = cute::make_tensor(tensor_F.data().get(), options.layoutF());
auto mX = cute::make_tensor(tensor_X.data().get(), options.layoutX());
auto mB = cute::make_tensor(tensor_B.data().get(), options.layoutB());
auto mC = cute::make_tensor(tensor_C.data().get(), options.layoutC());
auto mD = cute::make_tensor(tensor_D.data().get(), options.layoutD());
auto mZ = cute::make_tensor(tensor_Z.data().get(), options.layoutZ());
auto mDelta = cute::make_tensor(tensor_Delta.data().get(), options.layoutDelta());
auto mDeltaA = cute::make_tensor(tensor_DeltaA.data().get(), options.layoutDeltaA());
// Reference Device kernel
if (options.verify) {
ssd_reference<Option::HAS_D, Option::D_HAS_HDIM, Option::HAS_Z>(
mY_ref_1,
mF_ref_1,
mX,
mDelta,
mDeltaA,
mB,
mC,
mD,
mZ,
options
);
}
bool passed = true;
if (options.verify) {
printf("[TensorY]verifying...\n");
passed &= compare_reference<5>(mY_ref_1, mY_res);
printf("[TensorF]verifying...\n");
passed &= compare_reference<4>(mF_ref_1, mF_res);
}
return passed;
}
template<
int TensorDim,
class Engine, class Layout
>
static constexpr bool
compare_reference(
cute::Tensor<Engine, Layout> const& reference,
cute::Tensor<Engine, Layout> const& computed,
float epsilon = 0.05f) {
if (size(reference) != size(computed)) {
return false;
}
bool passed = true;
if (epsilon == 0.0f) {
// fast refcheck w/o epsilon
for (size_t i = 0; i < size_t(size(reference)); ++i) {
if (reference(i) != computed(i)) {
passed = false;
printf("[%llu] %f, %f\n", static_cast<unsigned long long>(i),
float(reference(i)), float(computed(i)));
break;
}
}
}
else {
// refcheck with epsilon
for (size_t i = 0; i < size_t(size(reference)); ++i) {
auto ref = static_cast<float>(reference(i));
auto act = static_cast<float>(computed(i));
auto abs_error = std::abs(act - ref);
auto rel_error = abs_error / (std::max(std::abs(act), std::abs(ref)) + 0.00001f);
if (std::isnan(abs_error) || std::isnan(rel_error) ||
std::min(rel_error, abs_error) > epsilon) {
passed = false;
printf("[%llu] %f, %f\n", static_cast<unsigned long long>(i),
float(reference(i)), float(computed(i)));
break;
}
}
}
if (not passed) {
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
auto m = cute::shape<2>(reference);
auto n = cute::shape<TensorDim-1>(reference);
printf("reference:\n");
for (int mi = 0; mi < m; ++mi) {
for (int ni = 0; ni < n; ++ni) {
if constexpr (TensorDim == 5) {
printf("%.4f ", static_cast<float>(reference(0,0,mi,2,ni)));
}
else {
printf("%.4f ", static_cast<float>(reference(0,0,mi,ni)));
}
}
printf("\n");
}
printf("\n");
printf("computed:\n");
for (int mi = 0; mi < m; ++mi) {
for (int ni = 0; ni < n; ++ni) {
if constexpr (TensorDim == 5) {
printf("%.4f ", static_cast<float>(computed(0,0,mi,2,ni)));
}
else {
printf("%.4f ", static_cast<float>(computed(0,0,mi,ni)));
}
}
printf("\n");
}
printf("\n");
}
return passed;
}
};
#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
int main(int argc, char const **args) {
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (__CUDACC_VER_MAJOR__ < 12 || props.major < 9) {
std::cout
<< "This example requires a GPU of NVIDIA's Hopper Architecture or "
<< "later (compute capability 90 or greater) and CUDA 12.0 or greater.\n";
return 0;
}
else if (__CUDACC_VER_MAJOR__ < 12 || (props.major != 9 || props.minor != 0)) {
std::cout
<< "This example requires a GPU of NVIDIA's Hopper Architecture "
<< "(compute capability 90) and CUDA 12.0 or greater.\n";
return 0;
}
#if defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
//
// Parse options
//
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
if (options.error) {
std::cerr << "Aborting execution." << std::endl;
return -1;
}
// Execute kernel
printf("start testing....\n");
// The KernelHardwareInfo struct holds the number of SMs on the GPU with a given device ID. This
// information is used by the underlying kernel.
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
hw_info.device_id = 0;
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
// Check Device/Host ref kernel
TestBed<Options> testbed{};
bool passed = testbed.run(options, hw_info);
if (passed) {
printf("everything is ok.\n");
}
else {
printf("something is wrong!!!!!\n");
}
#endif // defined(CUTLASS_ARCH_MMA_SM90_SUPPORTED)
return 0;
}
examples/111_hopper_ssd/CMakeLists.txt
+41
View File
@@ -0,0 +1,41 @@
# 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.
set_property(
SOURCE 111_hopper_ssd.cu
PROPERTY COMPILE_FLAGS "--use_fast_math"
)
cutlass_example_add_executable(
111_hopper_ssd
111_hopper_ssd.cu
)
if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")))
endif()
examples/111_hopper_ssd/README.md
+67
View File
@@ -0,0 +1,67 @@
# NVIDIA Hopper SSD (State Space Decomposition) CUDA Example
## Overview
This example demonstrates the implementation of State Space Decomposition (SSD) operations on NVIDIA's Hopper GPU architecture. It showcases the use of CUTLASS library components for high-performance tensor computations that efficiently leverage Hopper's advanced hardware capabilities, including TMA (Tensor Memory Accelerator) and warp specialization.
## System Requirements
+ NVIDIA GPU with Hopper Architecture (compute capability 9.0)
+ CUDA Toolkit 12.0 or newer
+ C++17 compatible compiler
## Build the example
Follow the cutlass example building.
## Command Line Options
The example supports the following command line options:
--help: Display the usage statement
--iterations=<int>: Number of iterations for benchmarking (default: 1)
--without_verify: Skip result verification
--verbose: Print execution time per kernel
--G=<int>: Group size (default: 2)
--B=<int>: Batch size (default: 3)
--E=<int>: Expanded factor (default: 2)
--H=<int>: Number of heads (default: 2)
## Limitation
+ Only support LxDxN = 128x64x128
+ Uses bfloat16 precision for input/output tensors
## Performance
+ Utilizes TMA (Tensor Memory Accelerator) for efficient memory access
+ Implements warp-specialized kernels for optimal resource utilization
+ Limited by the SEGSUM (segment sum) computation part
+ ALU bound operation
# Copyright
Copyright (c) 2024 - 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.
```
examples/111_hopper_ssd/collective/common.hpp
+174
View File
@@ -0,0 +1,174 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/kernel_hardware_info.h"
#include "cute/tensor.hpp"
namespace cutlass::ssd::collective {
using namespace cute;
template<typename Atom, typename TA, typename TB, typename TC>
CUTE_DEVICE void gemm_reset_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) {
constexpr int rA = decltype(rank(tA))::value;
constexpr int rB = decltype(rank(tB))::value;
constexpr int rC = decltype(rank(tC))::value;
if constexpr (rA == 2 && rB == 2 && rC == 1) {
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<1>(tA); k_block++) {
cute::gemm(atom, tA(_,k_block), tB(_,k_block), tC);
atom.accumulate_ = GMMA::ScaleOut::One;
}
}
else {
static_assert(rA == 3 && rB == 3 && rC == 3);
CUTLASS_PRAGMA_UNROLL
for (int k_block = 0; k_block < size<2>(tA); k_block++) {
cute::gemm(atom, tA(_,_,k_block), tB(_,_,k_block), tC);
atom.accumulate_ = GMMA::ScaleOut::One;
}
}
}
template<typename Atom, typename TA, typename TB, typename TC>
CUTE_DEVICE void gemm_zero_acc(Atom& atom, TA const& tA, TB const& tB, TC&& tC) {
atom.accumulate_ = GMMA::ScaleOut::Zero;
gemm_reset_zero_acc(atom, tA, tB, tC);
}
template<template<cute::GMMA::Major, cute::GMMA::Major, cute::GMMA::ScaleIn, cute::GMMA::ScaleIn> class Primitive, cute::GMMA::Major tA, cute::GMMA::Major tB, cute::GMMA::ScaleIn sA, cute::GMMA::ScaleIn sB>
inline auto __device__ constexpr convert_to_gmma_rs(cute::MMA_Atom<Primitive<tA, tB, sA, sB>> const& tiled_mma) {
using Atom = cute::MMA_Atom<Primitive<tA, tB, sA, sB>>;
using ElementA = typename Atom::ValTypeA;
using ElementB = typename Atom::ValTypeB;
using ElementC = typename Atom::ValTypeC;
using Shape_MNK = typename Atom::Shape_MNK;
using RS = decltype(cute::GMMA::rs_op_selector<ElementA, ElementB, ElementC, Shape_MNK, tA, tB, sA, sB>());
return cute::MMA_Atom<RS>{};
}
template<template<cute::GMMA::ScaleIn, cute::GMMA::ScaleIn> class Primitive, cute::GMMA::ScaleIn sA, cute::GMMA::ScaleIn sB>
inline auto __device__ constexpr convert_to_gmma_rs(cute::MMA_Atom<Primitive<sA, sB>> const& tiled_mma) {
using Atom = cute::MMA_Atom<Primitive<sA, sB>>;
using ElementA = typename Atom::ValTypeA;
using ElementB = typename Atom::ValTypeB;
using ElementC = typename Atom::ValTypeC;
using Shape_MNK = typename Atom::Shape_MNK;
constexpr auto tA = cute::GMMA::Major::K;
constexpr auto tB = cute::GMMA::Major::K;
using RS = decltype(cute::GMMA::rs_op_selector<ElementA, ElementB, ElementC, Shape_MNK, tA, tB, sA, sB>());
return cute::MMA_Atom<RS>{};
}
template<class Atom, class... Args>
CUTE_DEVICE auto constexpr convert_to_gmma_rs(cute::TiledMMA<Atom, Args...> const& tiled_mma) {
return cute::TiledMMA<decltype(convert_to_gmma_rs(Atom{})), Args...>{};
}
template<typename CLayout, typename AValueShape>
CUTE_DEVICE auto constexpr convert_c_layout_to_a_layout(CLayout const& c, AValueShape const& a) {
return make_layout(
make_shape(a, shape<1>(c), make_shape(shape<2>(c), size<0>(c) / size(a))),
make_stride(stride<0>(c), stride<1>(c), make_stride(stride<2>(c), size<2>(a) * stride<0,2>(c))));
}
template<class Layout, class Stages = _1>
CUTE_DEVICE constexpr auto unstageSmemLayout(Layout const& layout, Stages stages = {}) {
return composition(layout, make_tuple(_, _, make_layout(stages)));
}
template<class Element, class Accumulator, class OperandLayout_TV>
CUTE_DEVICE auto make_acc_into_op(Accumulator const& acc, OperandLayout_TV const& operand_layout_tv) {
Tensor operand = make_fragment_like<Element>(convert_c_layout_to_a_layout(acc.layout(), shape<1>(operand_layout_tv)));
Tensor operand_as_acc = make_tensor(operand.data(), acc.layout());
cute::copy(acc, operand_as_acc);
if constexpr (sizeof(Element) == 1) {
// 00 11 22 33 00 11 22 33 acc layout
// 00 00 11 11 22 22 33 33 operand layout
// BB AA AA BB AA BB BB AA conflict-free exchange pattern
// 16-bit exchange; so process two at a time potentially
int tid = threadIdx.x % 4;
auto values_u32 = recast<uint32_t>(operand);
CUTE_UNROLL
for (int n = 0; n < size<1>(values_u32); n++) {
CUTE_UNROLL
for (int k = 0; k < size<2>(values_u32); k++) {
CUTE_UNROLL
for (int ii = 0; ii < 8; ii += 4) {
uint32_t values_tmp_0 = values_u32(ii / 2 + 0, n, k);
uint32_t values_tmp_1 = values_u32(ii / 2 + 1, n, k);
// step A:
// t 1 v 0 -> t 0 v 1
// t 2 v 0 -> t 1 v 0
// t 0 v 1 -> t 2 v 0
// t 3 v 1 -> t 3 v 1
int v_to_send = tid == 1 || tid == 2 ? 0 : 1;
int v_to_recv = v_to_send;
int t_to_recv_from = (0x3021 >> (tid * 4)) & 0xF;
uint32_t values_tmp_a = v_to_send == 0 ? values_tmp_0 : values_tmp_1;
values_tmp_a = __shfl_sync(0xFFFFFFFF, values_tmp_a, t_to_recv_from, 4);
// step B:
// t 0 v 0 -> t 0 v 0
// t 3 v 0 -> t 1 v 1
// t 1 v 1 -> t 2 v 1
// t 2 v 1 -> t 3 v 0
v_to_send = 1 - v_to_send;
v_to_recv = 1 - v_to_recv;
t_to_recv_from = (0x2130 >> (tid * 4)) & 0xF;
uint32_t values_tmp_b = v_to_send == 0 ? values_tmp_0 : values_tmp_1;
values_tmp_b = __shfl_sync(0xFFFFFFFF, values_tmp_b, t_to_recv_from, 4);
values_u32(ii / 2 + 0, n, k) = __byte_perm(values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x1054 : 0x5410);
values_u32(ii / 2 + 1, n, k) = __byte_perm(values_tmp_a, values_tmp_b, v_to_send == 0 ? 0x3276 : 0x7632);
}
}
}
}
return operand;
}
} // namespace cutlass::fmha::collective
examples/111_hopper_ssd/collective/sm90_ssd_epilogue.hpp
+838
View File
@@ -0,0 +1,838 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/collective/collective_builder.hpp"
#include "cutlass/fast_math.h"
namespace cutlass::ssd::collective {
using namespace cute;
template<
class ElementAcc_,
class Element_,
class TileShape_,
class EpilogueTile_,
class SmemLayoutX_,
class SmemLayoutY_,
class SmemLayoutPartialY_,
class SmemLayoutP_,
class SmemLayoutZ_,
int StagesD_,
int StagesY_,
int StagesZ_,
bool HAS_D_,
bool D_HAS_HDIM_,
bool HAS_Z_>
struct SsdEpilogue {
using TileShape = TileShape_;
using ElementAcc = ElementAcc_;
using Element = Element_;
using ElementY = Element_;
using ElementP = Element_;
using ElementD = Element_;
using ElementX = Element_;
using ElementZ = Element_;
using EpilogueTile = EpilogueTile_;
using SmemLayoutX = SmemLayoutX_;
using SmemLayoutY = SmemLayoutY_;
using SmemLayoutPartialY = SmemLayoutPartialY_;
using SmemLayoutP = SmemLayoutP_;
using SmemLayoutZ = SmemLayoutZ_;
static constexpr int StagesD = StagesD_;
static constexpr int StagesY = StagesY_;
static constexpr int StagesZ = StagesZ_;
static constexpr bool HAS_D = HAS_D_;
static constexpr bool D_HAS_HDIM = D_HAS_HDIM_;
static constexpr bool HAS_Z = HAS_Z_;
// avoid "warning #3357-D: capturing structured bindings is a C++20 feature"
static constexpr auto L = get<0>(TileShape{});
static constexpr auto D = get<1>(TileShape{});
static constexpr auto N = get<2>(TileShape{});
constexpr static size_t SmemAlignmentY = cutlass::detail::alignment_for_swizzle(SmemLayoutY{});
struct CollectiveStorage {
alignas(SmemAlignmentY) ArrayEngine<ElementAcc, cosize_v<SmemLayoutPartialY>> smem_y_partial;
alignas(SmemAlignmentY) ArrayEngine<ElementY , cosize_v<SmemLayoutY>> smem_y;
alignas(SmemAlignmentY) ArrayEngine<ElementZ , cosize_v<SmemLayoutZ>> smem_z;
};
using EpiloadPipelineD = cutlass::PipelineTmaAsync<StagesD>;
using EpiloadPipelineZ = cutlass::PipelineTmaAsync<StagesZ>;
static constexpr int kEpiloadDBytes = D_HAS_HDIM ? D * sizeof(ElementD) : 0;
static constexpr int kEpiloadZBytes = HAS_Z ? cosize_v<SmemLayoutZ> * sizeof(ElementZ) : 0;
// TMA pipeline for storing D
using StorePipeline = cutlass::PipelineTmaStore<StagesY>;
using StorePipelineState = cutlass::PipelineState<StagesY>;
// TMA pipeline for storing P
using StorePPipeline = cutlass::PipelineTmaStore<1>;
using StorePPipelineState = cutlass::PipelineState<1>;
using CooperatePipeline = cutlass::PipelineAsync<StagesY>;
using CooperatePipelineState = cutlass::PipelineState<StagesY>;
struct SharedStorage {
using TensorStorage = CollectiveStorage;
TensorStorage tensors;
};
using TensorStorage = typename SharedStorage::TensorStorage;
using StrideY = cute::tuple<_1, int, int, int>; // (L,D,C,B)
using StrideP = cute::tuple<int, _1, int>; // (D,N,B)
using StrideZ = cute::tuple<_1, int, int, int>; // (L,D,C,B)
using LayoutY = decltype(make_layout(make_shape(L, D, int32_t(0), int32_t(0)),
make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,D,C,B)
using LayoutP = decltype(make_layout(make_shape(D, N, int32_t(0)), make_stride(N, _1{}, D*N))); // (D,N,B)
using LayoutD_2D = decltype(make_layout(make_shape(D, int32_t(0)), make_stride(_1{}, D))); // (D,EH)
using LayoutD_1D = decltype(make_layout(make_shape(_1{}, int32_t(0)), make_stride(_0{}, _1{}))); // (D,EH)
using LayoutD = cute::conditional_t<
D_HAS_HDIM,
LayoutD_2D,
LayoutD_1D
>;
using LayoutZ = decltype(make_layout(make_shape(L, D, int32_t(0), int32_t(0)),
make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,D,C,B)
struct Arguments {
ElementY* ptr_Y{nullptr};
ElementP* ptr_P{nullptr};
const ElementD* ptr_D{nullptr};
const ElementZ* ptr_Z{nullptr};
LayoutY layout_Y{};
LayoutP layout_P{};
LayoutD layout_D{};
LayoutZ layout_Z{};
};
using CopyOpS2G = SM90_TMA_STORE;
using CopyOpG2S = SM90_TMA_LOAD;
struct Params {
using TMA_Y = decltype(make_tma_atom(
CopyOpS2G{},
make_tensor(make_gmem_ptr(static_cast<ElementY const*>(nullptr)),LayoutY{}),
take<0,2>(SmemLayoutY{}),
EpilogueTile{}));
using TMA_P = decltype(make_tma_atom(
CopyOpS2G{},
make_tensor(make_gmem_ptr(static_cast<ElementP const*>(nullptr)),LayoutP{}),
take<0,2>(SmemLayoutP{}),
make_tile(shape<1>(TileShape{}), shape<2>(TileShape{}))));
using TMA_Z = decltype(make_tma_atom(
CopyOpG2S{},
make_tensor(make_gmem_ptr(static_cast<ElementZ const*>(nullptr)),LayoutZ{}),
take<0,2>(SmemLayoutZ{}),
make_tile(shape<0>(TileShape{}), shape<1>(TileShape{}))));
using TensorD = decltype(make_tensor(
make_gmem_ptr(static_cast<ElementD const*>(nullptr)),
LayoutD{}));
TMA_Y tma_store_y;
TMA_P tma_store_p;
TMA_Z tma_load_z;
TensorD tensor_d;
};
template<class ProblemShape>
static Params to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace = nullptr) {
using X = Underscore;
auto [G, B, EH, C, L, D, N] = problem_size;
auto tensor_y = make_tensor(make_gmem_ptr(args.ptr_Y), args.layout_Y);
auto tensor_p = make_tensor(make_gmem_ptr(args.ptr_P), args.layout_P);
auto tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), args.layout_D);
auto tensor_z = make_tensor(make_gmem_ptr(args.ptr_Z), args.layout_Z);
auto tma_store_y = make_tma_atom(
CopyOpS2G{},
tensor_y,
take<0,2>(SmemLayoutY{}),
EpilogueTile{});
auto tma_store_p = make_tma_atom(
CopyOpS2G{},
tensor_p,
take<0,2>(SmemLayoutP{}),
make_tile(shape<1>(TileShape{}), shape<2>(TileShape{})));
auto tma_load_z = make_tma_atom(
CopyOpG2S{},
tensor_z,
take<0,2>(SmemLayoutZ{}),
make_tile(shape<0>(TileShape{}), shape<1>(TileShape{})));
return Params{
tma_store_y,
tma_store_p,
tma_load_z,
tensor_d
};
}
template<
class Params, class ProblemShape,
class EpiloadPipeline, class PipelineState,
class TensorStorage
>
CUTLASS_DEVICE
void load_d(
int const& blk_coord, Params const& params, ProblemShape const& problem_size,
EpiloadPipeline& pipeline, PipelineState& pipeline_state,
TensorStorage& shared_tensors) {
if constexpr (D_HAS_HDIM) {
int lane_predicate = cute::elect_one_sync();
if (lane_predicate) {
auto& gD = params.tensor_d;
ElementD* ptr_d = shared_tensors.smem_d.data();
auto smem_layout = make_layout(make_shape(get<1>(TileShape{}), Int<StagesD>{})); // (D,)
auto sD = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_d), smem_layout));
auto bulk_atom = Copy_Atom<SM90_BULK_COPY_AUTO, ElementD>{};
int write_stage = pipeline_state.index();
// LOCK pipeline_state for _writing_
pipeline.producer_acquire(pipeline_state);
using BarrierType = typename EpiloadPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_state);
copy(bulk_atom.with(*tma_barrier), gD(_,blk_coord), sD(_,write_stage));
// Advance pipeline_state
++pipeline_state;
}
}
}
template<
class EpiloadPipeline, class PipelineState
>
CUTLASS_DEVICE void
load_d_tail(EpiloadPipeline pipeline, PipelineState pipeline_state) {
int lane_predicate = cute::elect_one_sync();
if (lane_predicate) {
pipeline.producer_tail(pipeline_state);
}
}
template <
class Params, class ProblemShape
>
CUTLASS_DEVICE
auto load_z_init(Params const& params, ProblemShape const& problem_size) {
using X = Underscore;
auto [G, B, EH, C, L, D, N] = problem_size;
Tensor mZ_mkl = params.tma_load_z.get_tma_tensor(make_shape(L,D,C,EH*B));
Tensor gZ_mkl = local_tile(mZ_mkl, TileShape{}, make_coord(_,_,_), Step<_1,_1,X>{});
return make_tuple(gZ_mkl);
}
template<
class Params, class ProblemShape,
class EpiLoadPipeline, class PipelineState,
class GTensor,
class TensorStorage
>
CUTLASS_DEVICE
void load_z(
int const& blk_coord, Params const& params, ProblemShape const& problem_size,
EpiLoadPipeline& pipeline, PipelineState& pipeline_state,
cute::tuple<GTensor> const& load_inputs,
TensorStorage& shared_tensors) {
if constexpr (HAS_Z) {
int lane_predicate = cute::elect_one_sync();
if (lane_predicate) {
Tensor sZ_ = make_tensor(make_smem_ptr(shared_tensors.smem_z.begin()), SmemLayoutZ{});
Tensor sZ = as_position_independent_swizzle_tensor(sZ_);
//
// Prepare the TMA loads for Z
//
Tensor gZ_mkl = get<0>(load_inputs);
// Partition the inputs based on the current block coordinates.
Tensor gZ = gZ_mkl(_,_,_0{},_0{},_,blk_coord);
auto oprands = tma_partition(params.tma_load_z, Int<0>{}, Layout<_1>{},
group_modes<0,2>(sZ), group_modes<0,2>(gZ)); // (TMA,k) and (TMA,PIPE)
// avoid "warning #3357-D: capturing structured bindings is a C++20 feature"
auto tZgZ = get<0>(oprands);
auto tZsZ = get<1>(oprands);
// Disable multicast
uint16_t mcast_mask_z = 0;
auto chunk = get<3>(problem_size);
// Mainloop
CUTLASS_PRAGMA_NO_UNROLL
for (int chunk_idx = 0; chunk_idx < chunk; ++chunk_idx) {
int write_stage = pipeline_state.index();
// LOCK pipeline_state for _writing_
pipeline.producer_acquire(pipeline_state);
using BarrierType = typename EpiLoadPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_state);
copy(params.tma_load_z.with(*tma_barrier, mcast_mask_z), tZgZ(_,chunk_idx), tZsZ(_,write_stage));
// Advance pipeline_state
++pipeline_state;
}
}
}
}
template<
class EpiLoadPipeline, class PipelineState
>
CUTLASS_DEVICE void
load_z_tail(EpiLoadPipeline pipeline, PipelineState pipeline_state) {
int lane_predicate = cute::elect_one_sync();
// Issue the epilogue waits
if (lane_predicate) {
/* This helps avoid early exit of blocks in Cluster
* Waits for all stages to either be released (all
* Consumer UNLOCKs), or if the stage was never used
* then would just be acquired since the phase was
* still inverted from make_producer_start_state
*/
pipeline.producer_tail(pipeline_state);
}
}
template<
class Params, class ProblemShape,
class CooperatePipeline, class CooperatePipelineState,
class TensorIntra,
class TiledMma,
class TensorStorage
>
CUTLASS_DEVICE
auto store_intra(
int& chunk, int const& blk_coord, Params const& params, ProblemShape const& problem_size,
CooperatePipeline& cooperate_pipeline, CooperatePipelineState& cooperate_pipe_producer_state,
TensorIntra& tIntra,
TiledMma tiled_mma,
TensorStorage& shared_tensors) {
int thread_idx = int(threadIdx.x % 128);
auto [G, B, EH, C, L, D, N] = problem_size;
Tensor mY_mn = params.tma_store_y.get_tma_tensor(make_shape(L,D,C,B*EH));
Tensor gY_mn = local_tile(mY_mn, take<0,2>(TileShape{}), make_coord(_,_,_))(_,_,_0{},_0{},chunk,blk_coord);
// Apply epilogue subtiling
Tensor gY_epi = flat_divide(gY_mn, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
auto epi_tile_m = size<0>(EpilogueTile{});
auto epi_tile_n = size<1>(EpilogueTile{});
auto partial_m = Int<128>{};
auto partial_n = Int<epi_tile_m * epi_tile_n / 128>{};
auto ptr_sY = shared_tensors.smem_y_partial.begin();
Tensor sY_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sY),
tile_to_shape(UMMA::Layout_K_SW64_Atom<ElementAcc>{}, make_shape(partial_m, partial_n, Int<2>{}))));
using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, Element>;
TiledCopy tiled_copy_C_atom = make_tiled_copy_C_atom(CopyAtomC{}, tiled_mma);
using CopyOpR2S = SM90_U16x8_STSM_T;
TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom<CopyOpR2S,ElementY>{}, tiled_copy_C_atom);
ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx);
Tensor tRS_rIntra = thread_r2s.retile_S(tIntra); // ((R2S,R2S_V),MMA_M,MMA_N)
// Hard code
static constexpr int FragmentSize = 4;
auto tRS_rY = make_tensor<ElementAcc>(shape(sY_epi(thread_idx,_,_0{})));
Tensor tRS_rIntra_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rIntra);
Tensor tRS_sY_frg = recast<Array<ElementAcc, FragmentSize>>(sY_epi(thread_idx,_,_));
Tensor tRS_rY_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rY);
auto mma_tile_m = size<0>(TileShape{}) / size<1>(tRS_rIntra);
auto mma_tile_n = size<1>(TileShape{}) / size<2>(tRS_rIntra);
// For each output tile
CUTLASS_PRAGMA_UNROLL
for (int epi_n = 0; epi_n < size<3>(gY_epi); ++epi_n) {
CUTLASS_PRAGMA_UNROLL
for (int epi_m = 0; epi_m < size<2>(gY_epi); ++epi_m) {
int mma_m = epi_m;
int mma_n = (epi_n * size<1>(EpilogueTile{})) / mma_tile_n;
Tensor tRS_rIntra_frg_mn = tRS_rIntra_frg(_,mma_m,mma_n);
// Vectorized fragment loop with visitor callback entry point
// Epilogue op
int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n);
int r2s_v = epi_n_in_mma * size(tRS_rY_frg);
CUTLASS_PRAGMA_UNROLL
for (int epi_v = 0; epi_v < size(tRS_rY_frg); ++epi_v) {
tRS_rY_frg(epi_v) = tRS_rIntra_frg_mn(r2s_v + epi_v);
}
cooperate_pipeline.producer_acquire(cooperate_pipe_producer_state);
copy(tRS_rY_frg, tRS_sY_frg(_,cooperate_pipe_producer_state.index()));
cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to Cooperate warps
cooperate_pipeline.producer_commit(cooperate_pipe_producer_state);
++cooperate_pipe_producer_state;
}
}
}
template<
class Params,
class EpiloadPipeline, class EpiloadPipelineState,
class TiledMma,
class TensorStorage
>
CUTLASS_DEVICE
auto update_d(
int const& blk_coord, Params params,
bool is_first_iteration,
EpiloadPipeline& epi_load_pipeline, EpiloadPipelineState& epi_load_pipe_consumer_state,
TiledMma tiled_mma,
TensorStorage& shared_tensors
) {
int thread_idx = int(threadIdx.x % 128);
int read_stage = epi_load_pipe_consumer_state.index();
// Load delta for epilogue
auto row_layout = make_layout(make_shape(get<0>(TileShape{}), get<1>(TileShape{}), Int<StagesD>{}), make_stride(_0{}, _1{}, get<1>(TileShape{})));
auto c_tv_layout = typename TiledMma::LayoutC_TV{};
auto tile_mn = make_shape(tile_size<0>(tiled_mma), tile_size<1>(tiled_mma));
auto d_layout = zipped_divide(row_layout, tile_mn);
auto d_tv_layout = composition(d_layout, make_tuple(c_tv_layout,_));
auto sD = make_tensor(make_smem_ptr(shared_tensors.smem_d.data()), d_tv_layout)(make_coord(thread_idx,_), make_coord(_,_,read_stage));
auto tSR_sD = as_position_independent_swizzle_tensor(sD);
auto tSR_rD = make_tensor<ElementD>(shape(sD));
auto tD = make_tensor<ElementAcc>(shape(sD));
if constexpr (D_HAS_HDIM) {
if (is_first_iteration) {
epi_load_pipeline.consumer_wait(epi_load_pipe_consumer_state);
}
copy(tSR_sD, tSR_rD);
type_convert<ElementD, ElementAcc>(tSR_rD, tD);
}
else if constexpr (HAS_D) {
auto& gD = params.tensor_d;
auto value = static_cast<ElementAcc>(gD(_0{}, blk_coord));
CUTLASS_PRAGMA_UNROLL
for (int ii = 0; ii < size(tSR_rD); ++ii) {
tD(ii) = value;
}
}
return make_tuple(tD);
}
template<
class Params, class ProblemShape,
class StorePipeline, class StorePipelineState,
class CooperatePipeline, class CooperatePipelineState,
class MainloopPipelineX, class PipelineStateX,
class EpiloadPipelineZ, class PipelineStateZ,
class TensorInter, class TensorDeltaA, class TensorD,
class TiledMma,
class TensorStorage, class TensorStorageX
>
CUTLASS_DEVICE
auto store(
int& chunk, int const& blk_coord, Params const& params, ProblemShape const& problem_size,
StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state,
CooperatePipeline& cooperate_pipeline, CooperatePipelineState& cooperate_pipe_consumer_state,
MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_state_x,
EpiloadPipelineZ& pipeline_z, PipelineStateZ& pipeline_state_z,
TensorInter& tInter, TensorDeltaA& tDeltaA, TensorD& tD,
TiledMma tiled_mma,
TensorStorage& shared_tensors, TensorStorageX& shared_tensors_x) {
int thread_idx = int(threadIdx.x % 128);
auto [G, B, EH, C, L, D, N] = problem_size;
Tensor mY_mn = params.tma_store_y.get_tma_tensor(make_shape(L,D,C,B*EH));
Tensor gY_mn = local_tile(mY_mn, take<0,2>(TileShape{}), make_coord(_,_,_))(_,_,_0{},_0{},chunk,blk_coord);
// Apply epilogue subtiling
Tensor gY_epi = flat_divide(gY_mn, EpilogueTile{}); // (EPI_TILE_M,EPI_TILE_N,EPI_M,EPI_N)
// Construct the corresponding pipelined smem tensors
auto ptr_sY = shared_tensors.smem_y.begin();
Tensor sY_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sY), SmemLayoutY{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
auto ptr_sX = shared_tensors_x.smem_x.data();
Tensor sX_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sX), SmemLayoutX{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
auto ptr_sZ = shared_tensors.smem_z.begin();
Tensor sZ_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sZ), SmemLayoutZ{})); // (EPI_TILE_M,EPI_TILE_N,PIPE_D)
auto epi_tile_m = size<0>(EpilogueTile{});
auto epi_tile_n = size<1>(EpilogueTile{});
auto partial_m = Int<128>{};
auto partial_n = Int<epi_tile_m * epi_tile_n / 128>{};
auto ptr_sY_partial = shared_tensors.smem_y_partial.begin();
Tensor sY_epi_partial = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sY_partial),
tile_to_shape(UMMA::Layout_K_SW64_Atom<ElementAcc>{}, make_shape(partial_m, partial_n, Int<2>{}))));
using CopyAtomC = Copy_Atom<SM90_U32x4_STSM_N, Element>;
TiledCopy tiled_copy_C_atom = make_tiled_copy_C_atom(CopyAtomC{}, tiled_mma);
using CopyOpR2S = SM90_U16x8_STSM_T;
TiledCopy tiled_r2s = make_tiled_copy_S(Copy_Atom<CopyOpR2S,ElementY>{}, tiled_copy_C_atom);
ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx);
Tensor tRS_rInter = thread_r2s.retile_S(tInter); // ((R2S,R2S_V),MMA_M,MMA_N)
Tensor tRS_rDeltaA = thread_r2s.retile_S(tDeltaA); // ((R2S,R2S_V),MMA_M,MMA_N)
Tensor tRS_rD = thread_r2s.retile_S(tD); // ((R2S,R2S_V),MMA_M,MMA_N)
Tensor tRS_sY = thread_r2s.partition_D(sY_epi); // (R2S,R2S_M,R2S_N,PIPE_D)
using CopyOpS2R = SM75_U16x8_LDSM_T;
TiledCopy tiled_s2r = make_tiled_copy_S(Copy_Atom<CopyOpS2R,ElementX>{}, tiled_copy_C_atom);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
Tensor tSR_sX = thread_s2r.partition_S(flat_divide(sX_epi, EpilogueTile{}));
Layout tSR_rX_layout = make_layout(take<0,3>(shape(thread_s2r.partition_D(flat_divide(sX_epi, EpilogueTile{})))));
Tensor tSR_rX = make_tensor<ElementX>(tSR_rX_layout);
Tensor tRS_rX = make_tensor<ElementAcc>(tSR_rX_layout);
using CopyOpS2R_Z = SM75_U16x8_LDSM_T;
TiledCopy tiled_s2r_z = make_tiled_copy_S(Copy_Atom<CopyOpS2R_Z,ElementZ>{}, tiled_copy_C_atom);
ThrCopy thread_s2r_z = tiled_s2r_z.get_slice(thread_idx);
Tensor tSR_sZ = thread_s2r_z.partition_S(flat_divide(sZ_epi, EpilogueTile{}));
Layout tSR_rZ_layout = make_layout(take<0,3>(shape(thread_s2r_z.partition_D(flat_divide(sZ_epi, EpilogueTile{})))));
Tensor tSR_rZ = make_tensor<ElementX>(tSR_rZ_layout);
Tensor tRS_rZ = make_tensor<ElementAcc>(tSR_rZ_layout);
// Hard code
static constexpr int FragmentSize = 4;
// Allocate Y registers
Layout tRS_rY_layout = make_layout(take<0,3>(shape(thread_r2s.partition_S(sY_epi))));
Tensor tRS_rY = make_tensor<ElementY>(tRS_rY_layout); // (R2S,R2S_M,R2S_N)
Tensor tRS_rY_frg = recast<Array<ElementY, FragmentSize>>(tRS_rY);
Tensor tRS_rX_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rX);
Tensor tRS_rZ_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rZ);
// Tensor tRS_rIntra_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rIntra);
Tensor tRS_rInter_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rInter);
Tensor tRS_rDeltaA_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rDeltaA);
Tensor tRS_rD_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rD);
Tensor tRS_rCompute = make_tensor<ElementAcc>(tRS_rY_layout); // (R2S,R2S_M,R2S_N)
Tensor tRS_rCompute_frg = recast<Array<ElementAcc, FragmentSize>>(tRS_rCompute);
auto tSR_rY = make_tensor<ElementAcc>(shape(sY_epi_partial(thread_idx,_,_0{})));
Tensor tSR_sY_frg = recast<Array<ElementAcc, FragmentSize>>(sY_epi_partial(thread_idx,_,_));
Tensor tSR_rY_frg = recast<Array<ElementAcc, FragmentSize>>(tSR_rY);
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
auto oprands = tma_partition(params.tma_store_y, Int<0>{}, Layout<_1>{},
group_modes<0,2>(sY_epi), group_modes<0,2>(gY_epi)); // (TMA,k) and (TMA,PIPE)
// avoid "warning #3357-D: capturing structured bindings is a C++20 feature"
auto bSG_gY = get<0>(oprands);
auto bSG_sY = get<1>(oprands);
auto mma_tile_m = size<0>(TileShape{}) / size<1>(tRS_rInter);
auto mma_tile_n = size<1>(TileShape{}) / size<2>(tRS_rInter);
#if 0
if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0) && chunk == 0) {
print("tRS_rInter : ");print(tRS_rInter);print("\n");
print("tRS_sY : ");print(tRS_sY);print("\n");
print("bSG_sY : ");print(bSG_sY);print("\n");
print("bSG_gY : ");print(bSG_gY);print("\n");
print("tRS_rY_frg : ");print(tRS_rY_frg);print("\n");
print("tSR_rY_frg : ");print(tSR_rY_frg);print("\n");
print("tRS_rInter_frg : ");print(tRS_rInter_frg);print("\n");
print("tRS_rCompute_frg : ");print(tRS_rCompute_frg);print("\n");
}
#endif
// Thread synchronizer for previously issued waits or fences
// to ensure visibility of smem reads/writes to threads or TMA unit
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
// Predication for TMA store (one warp issues TMA store)
bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0;
int epi_m_prev = 0, epi_n_prev = 0;
auto tma_store_fn = [&] (int epi_m, int epi_n) {
// Write the tile from smem to gmem with TMA
cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
synchronize(); // ensure all threads have issued their async fence
if (issue_tma_store) {
copy(params.tma_store_y, bSG_sY(_,store_pipe_producer_state.index()), bSG_gY(_,epi_m,epi_n));
}
// Commit the TMA stores for this stage
if (issue_tma_store) {
store_pipeline.producer_commit(store_pipe_producer_state);
}
++store_pipe_producer_state;
// Wait for the next smem buffer to be available
if (issue_tma_store) {
store_pipeline.producer_acquire(store_pipe_producer_state);
}
synchronize();
};
static constexpr bool DelayTmaStore = true;
int subtile_idx = -1;
if constexpr (HAS_Z) {
pipeline_z.consumer_wait(pipeline_state_z);
}
// For each output tile
CUTLASS_PRAGMA_UNROLL
for (int epi_n = 0; epi_n < size<3>(gY_epi); ++epi_n) {
CUTLASS_PRAGMA_UNROLL
for (int epi_m = 0; epi_m < size<2>(gY_epi); ++epi_m) {
bool is_first_iteration = epi_m == 0 && epi_n == 0;
bool is_last_iteration = epi_m == size<2>(gY_epi)-1 && epi_n == size<3>(gY_epi)-1;
if (subtile_idx != -1 && (epi_n * static_cast<int>(size<2>(gY_epi)) + epi_m) != subtile_idx) {
continue;
}
if constexpr (HAS_D) {
// load x
copy(tiled_s2r, tSR_sX(_,_,_,epi_m,epi_n,pipeline_state_x.index()), tSR_rX);
type_convert<ElementX, ElementAcc>(tSR_rX, tRS_rX);
}
if constexpr (HAS_Z) {
// load x
copy(tiled_s2r_z, tSR_sZ(_,_,_,epi_m,epi_n,pipeline_state_z.index()), tSR_rZ);
type_convert<ElementZ, ElementAcc>(tSR_rZ, tRS_rZ);
// SiLu
cutlass::epilogue::thread::SiLu<ElementAcc> op;
for (int ii = 0; ii < size(tRS_rZ); ++ii) {
tRS_rZ(ii) = op(tRS_rZ(ii));
}
}
cooperate_pipeline.consumer_wait(cooperate_pipe_consumer_state);
copy(tSR_sY_frg(_,cooperate_pipe_consumer_state.index()), tSR_rY_frg);
cooperate_pipeline.consumer_release(cooperate_pipe_consumer_state);
++cooperate_pipe_consumer_state;
cutlass::arch::fence_view_async_shared();
synchronize();
int mma_m = epi_m;
int mma_n = (epi_n * size<1>(EpilogueTile{})) / mma_tile_n;
// Tensor tRS_rIntra_frg_mn = tRS_rIntra_frg(_,mma_m,mma_n);
Tensor tRS_rInter_frg_mn = tRS_rInter_frg(_,mma_m,mma_n);
Tensor tRS_rDeltaA_frg_mn = tRS_rDeltaA_frg(_,mma_m,mma_n);
Tensor tRS_rD_frg_mn = tRS_rD_frg(_,mma_m,mma_n);
// Vectorized fragment loop with visitor callback entry point
// Epilogue op
int epi_n_in_mma = epi_n % (mma_tile_n / epi_tile_n);
int r2s_v = epi_n_in_mma * size(tRS_rCompute_frg);
CUTLASS_PRAGMA_UNROLL
for (int epi_v = 0; epi_v < size(tRS_rCompute_frg); ++epi_v) {
tRS_rCompute_frg(epi_v) = tRS_rDeltaA_frg_mn(r2s_v + epi_v) * tRS_rInter_frg_mn(r2s_v + epi_v) + tSR_rY_frg(epi_v);
if constexpr (HAS_D) {
tRS_rCompute_frg(epi_v) = tRS_rD_frg_mn(r2s_v + epi_v) * tRS_rX_frg(epi_v) + tRS_rCompute_frg(epi_v);
}
if constexpr (HAS_Z) {
tRS_rCompute_frg(epi_v) = tRS_rCompute_frg(epi_v) * tRS_rZ_frg(epi_v);
}
}
// The latest we can delay the TMA store is right before the smem store of the next iteration
// since the current TMA store needs to be committed before we can acquire the next smem buffer
if constexpr (DelayTmaStore) {
// Issue TMA stores for the previous subtile
if (not is_first_iteration and subtile_idx == -1) {
tma_store_fn(epi_m_prev, epi_n_prev);
}
epi_m_prev = epi_m;
epi_n_prev = epi_n;
}
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < size(tRS_rY_frg); ++i) {
tRS_rY_frg(i) = cutlass::NumericArrayConverter<ElementY, ElementAcc, FragmentSize>{}(tRS_rCompute_frg(i));
}
copy(tiled_r2s, tRS_rY, tRS_sY(_,_,_,store_pipe_producer_state.index()));
if constexpr (not DelayTmaStore) {
// Issue TMA stores for this subtile
tma_store_fn(epi_m, epi_n);
}
} // for epi_m
} // for epi_n
pipeline_x.consumer_release(pipeline_state_x);
++pipeline_state_x;
pipeline_z.consumer_release(pipeline_state_z);
++pipeline_state_z;
if constexpr (DelayTmaStore) {
// Issue TMA stores for the last subtile
tma_store_fn(epi_m_prev, epi_n_prev);
}
}
template<
class Params, class ProblemShape,
class StorePipeline, class StorePipelineState,
class TensorStorage
>
CUTLASS_DEVICE
auto store_p(
int const& blk_coord, Params const& params, ProblemShape const& problem_size,
StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state,
TensorStorage& shared_tensors) {
int thread_idx = int(threadIdx.x % 128);
auto [G, B, EH, C, L, D, N] = problem_size;
Tensor mP_mn = params.tma_store_p.get_tma_tensor(make_shape(D,N,B*EH));
Tensor gP_mn = local_tile(mP_mn, take<1,3>(TileShape{}), make_coord(_,_,blk_coord));
auto gP_epi = gP_mn(_,_,_0{},_0{});
// Construct the corresponding pipelined smem tensors
auto ptr_sP = shared_tensors.smem_p.begin();
Tensor sP_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sP), SmemLayoutP{})); // (EPI_TILE_M,EPI_TILE_N)
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
auto oprands = tma_partition(params.tma_store_p, Int<0>{}, Layout<_1>{},
group_modes<0,2>(sP_epi), group_modes<0,2>(gP_epi)); // (TMA,k) and (TMA,PIPE)
// avoid "warning #3357-D: capturing structured bindings is a C++20 feature"
auto bSG_gY = get<0>(oprands);
auto bSG_sY = get<1>(oprands);
#if 0
if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) {
print("sP_epi : ");print(sP_epi);print("\n");
print("gP_epi : ");print(gP_epi);print("\n");
print("bSG_sY : ");print(bSG_sY);print("\n");
print("bSG_gY : ");print(bSG_gY);print("\n");
}
#endif
// Thread synchronizer for previously issued waits or fences
// to ensure visibility of smem reads/writes to threads or TMA unit
// Use the reserved named barrier `streamkbarrier` since this kernel doesn't support streamk
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::StreamkBarrier0); };
// Predication for TMA store (one warp issues TMA store)
bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0;
auto tma_store_fn = [&] () {
// Write the tile from smem to gmem with TMA
cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
synchronize(); // ensure all threads have issued their async fence
if (issue_tma_store) {
copy(params.tma_store_p, bSG_sY, bSG_gY);
}
// Commit the TMA stores for this stage
if (issue_tma_store) {
store_pipeline.producer_commit(store_pipe_producer_state);
}
++store_pipe_producer_state;
// Wait for the next smem buffer to be available
if (issue_tma_store) {
store_pipeline.producer_acquire(store_pipe_producer_state);
}
// don't need pipeline
synchronize();
};
tma_store_fn();
}
template<
class ElementSrc, class ElementDst,
class TensorSrc, class TensorDst
>
CUTLASS_DEVICE
auto type_convert(
TensorSrc& tS,
TensorDst& tD) {
static constexpr int FragmentSize = 2;
NumericArrayConverter<ElementDst, ElementSrc, FragmentSize> converter;
auto tS_frg = recast<Array<ElementSrc, FragmentSize>>(tS);
auto tD_frg = recast<Array<ElementDst, FragmentSize>>(tD);
CUTLASS_PRAGMA_UNROLL
for (int ii = 0; ii < size(tS_frg); ++ii) {
tD_frg(ii) = converter(tS_frg(ii));
}
}
};
} // namespace cutlass::fmha::collective
examples/111_hopper_ssd/collective/sm90_ssd_gemm_tma_warpspecialized.hpp
+1146
View File
File diff suppressed because it is too large Load Diff
examples/111_hopper_ssd/device/ssd.hpp
+273
View File
@@ -0,0 +1,273 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
// common
#include "cutlass/cutlass.h"
#include "cutlass/device_kernel.h"
#if !defined(__CUDACC_RTC__)
#include "cutlass/cluster_launch.hpp"
#include "cutlass/trace.h"
#endif // !defined(__CUDACC_RTC__)
////////////////////////////////////////////////////////////////////////////////
namespace cutlass::ssd::device {
////////////////////////////////////////////////////////////////////////////////
////////////////////////////// CUTLASS 3.x API /////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template <class Kernel_>
class SSD {
public:
using Kernel = Kernel_;
static int const kThreadCount = Kernel::MaxThreadsPerBlock;
/// Argument structure: User API
using Arguments = typename Kernel::Arguments;
/// Argument structure: Kernel API
using Params = typename Kernel::Params;
private:
/// Kernel API parameters object
Params params_;
bool is_initialized(bool set = false) {
static bool initialized = false;
if (set) initialized = true;
return initialized;
}
public:
/// Access the Params structure
Params const& params() const {
return params_;
}
/// Determines whether the GEMM can execute the given problem.
static Status
can_implement(Arguments const& args) {
if (Kernel::can_implement(args)) {
return Status::kSuccess;
}
else {
return Status::kInvalid;
}
}
/// Gets the workspace size
static size_t
get_workspace_size(Arguments const& args) {
size_t workspace_bytes = 0;
workspace_bytes += Kernel::get_workspace_size(args);
return workspace_bytes;
}
/// Computes the grid shape
static dim3
get_grid_shape(Params const& params) {
return Kernel::get_grid_shape(params);
}
/// Computes the maximum number of active blocks per multiprocessor
static int maximum_active_blocks(int /* smem_capacity */ = -1) {
CUTLASS_TRACE_HOST("Universal::maximum_active_blocks()");
int max_active_blocks = -1;
int smem_size = Kernel::SharedStorageSize;
// first, account for dynamic smem capacity if needed
cudaError_t result;
if (smem_size >= (48 << 10)) {
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
result = cudaFuncSetAttribute(
device_kernel<Kernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(
" cudaFuncSetAttribute() returned error: "
<< cudaGetErrorString(result));
return -1;
}
}
// query occupancy after setting smem size
result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&max_active_blocks,
device_kernel<Kernel>,
Kernel::MaxThreadsPerBlock,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(
" cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: "
<< cudaGetErrorString(result));
return -1;
}
CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks);
return max_active_blocks;
}
/// Initializes GEMM state from arguments.
Status
initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("Universal::initialize() - workspace "
<< workspace << ", stream: " << (stream ? "non-null" : "null"));
// Initialize the workspace
Status status = Kernel::initialize_workspace(args, workspace, stream);
if (status != Status::kSuccess) {
return status;
}
// Initialize the Params structure
params_ = Kernel::to_underlying_arguments(args, workspace);
if (is_initialized()) return Status::kSuccess;
// account for dynamic smem capacity if needed
int smem_size = Kernel::SharedStorageSize;
printf("[Usage] smem : %d\n", smem_size);
if (smem_size >= (48 << 10)) {
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
cudaError_t result = cudaFuncSetAttribute(
device_kernel<Kernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
return Status::kErrorInternal;
}
}
is_initialized(true);
return Status::kSuccess;
}
/// Update API is preserved in 3.0, but does not guarantee a lightweight update of params.
Status
update(Arguments const& args, void* workspace = nullptr) {
CUTLASS_TRACE_HOST("Universal()::update() - workspace: " << workspace);
size_t workspace_bytes = get_workspace_size(args);
if (workspace_bytes > 0 && nullptr == workspace) {
return Status::kErrorWorkspaceNull;
}
params_ = Kernel::to_underlying_arguments(args, workspace);
return Status::kSuccess;
}
/// Primary run() entry point API that is static allowing users to create and manage their own params.
/// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
static Status
run(Params& params, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("Universal::run()");
dim3 const block = Kernel::get_block_shape();
dim3 const grid = get_grid_shape(params);
// configure smem size and carveout
int smem_size = Kernel::SharedStorageSize;
Status launch_result;
// Use extended launch API only for mainloops that use it
if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) {
dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}),
cute::size<1>(typename Kernel::ClusterShape{}),
cute::size<2>(typename Kernel::ClusterShape{}));
void const* kernel = (void const*) device_kernel<Kernel>;
void* kernel_params[] = {&params};
launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params);
}
else {
launch_result = Status::kSuccess;
device_kernel<Kernel><<<grid, block, smem_size, stream>>>(params);
}
cudaError_t result = cudaGetLastError();
if (cudaSuccess == result && Status::kSuccess == launch_result) {
return Status::kSuccess;
}
else {
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
return Status::kErrorInternal;
}
}
//
// Non-static launch overloads that first create and set the internal params struct of this kernel handle.
//
/// Launches the kernel after first constructing Params internal state from supplied arguments.
Status
run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (Status::kSuccess == status) {
status = run(params_, stream);
}
return status;
}
/// Launches the kernel after first constructing Params internal state from supplied arguments.
Status
operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
return run(args, workspace, stream);
}
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
Status
run(cudaStream_t stream = nullptr) {
return run(params_, stream);
}
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
Status
operator()(cudaStream_t stream = nullptr) {
return run(params_, stream);
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::device
////////////////////////////////////////////////////////////////////////////////
examples/111_hopper_ssd/kernel/sm90_ssd_kernel_builder.hpp
+116
View File
@@ -0,0 +1,116 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "../collective/sm90_ssd_epilogue.hpp"
#include "../collective/sm90_ssd_gemm_tma_warpspecialized.hpp"
#include "../kernel/sm90_ssd_kernel_tma_warpspecialized.hpp"
#include "../kernel/sm90_ssd_tile_scheduler.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/collective/collective_builder.hpp"
namespace cutlass::ssd::kernel {
template<
class Element_,
class ElementDA_,
class ElementAcc_,
class ElementY_,
class TileShape_,
bool HAS_D_,
bool D_HAS_HDIM_,
bool HAS_Z_
>
struct Sm90SsdBuilder {
using Element = Element_;
using ElementDA = ElementDA_;
using ElementAcc = ElementAcc_;
using ElementY = ElementY_;
using TileShape = TileShape_;
static constexpr bool HAS_D = HAS_D_;
static constexpr bool D_HAS_HDIM = D_HAS_HDIM_;
static constexpr bool HAS_Z = HAS_Z_;
static constexpr int StagesY = 2;
static constexpr int StagesX = 2;
static constexpr int StagesZ = 1; // smem size limitation
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using Schedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTile = decltype(cutlass::epilogue::collective::detail::sm90_compute_tile_shape_or_override<
ElementY, EpilogueTileType, Schedule, TileShape>());
using SmemLayoutAtomY = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
cute::GMMA::Major::MN, ElementY, decltype(get<0>(EpilogueTile{})), decltype(get<1>(EpilogueTile{}))>());
using SmemLayoutY = decltype(tile_to_shape(
SmemLayoutAtomY{},
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesY>{}),
Step<_2,_1,_3>{}));
using SmemLayoutAtomX = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
cute::GMMA::Major::MN, Element, decltype(get<0>(TileShape{})), decltype(get<1>(TileShape{}))>());
using SmemLayoutX = decltype(tile_to_shape(
SmemLayoutAtomY{},
make_shape(size<0>(TileShape{}), size<1>(TileShape{}), Int<StagesX>{}),
Step<_2,_1,_3>{}));
using SmemLayoutAtomZ = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
cute::GMMA::Major::MN, Element, decltype(get<0>(TileShape{})), decltype(get<1>(TileShape{}))>());
using SmemLayoutZ = decltype(tile_to_shape(
SmemLayoutAtomZ{},
make_shape(size<0>(TileShape{}), size<1>(TileShape{}), Int<StagesZ>{}),
Step<_2,_1,_3>{}));
static constexpr auto epi_tile_m = size<0>(EpilogueTile{});
static constexpr auto epi_tile_n = size<1>(EpilogueTile{});
static constexpr auto partial_m = Int<128>{};
static constexpr auto partial_n = Int<epi_tile_m * epi_tile_n / 128>{};
using SmemLayoutAtomPartialY = typename GMMA::Layout_K_SW64_Atom<ElementAcc>;
using SmemLayoutPartialY = decltype(tile_to_shape(
SmemLayoutAtomPartialY{},
make_shape(partial_m, partial_n, Int<StagesY>{})));
using CollectiveMainloop = cutlass::ssd::collective::SsdMainloopTmaWarpSpecialized<Element, ElementDA, ElementAcc, ElementY, TileShape, StagesX>;
using CollectiveEpilogue = cutlass::ssd::collective::SsdEpilogue<
ElementAcc, ElementY, TileShape,
EpilogueTile, SmemLayoutX, SmemLayoutY, SmemLayoutPartialY, typename CollectiveMainloop::SmemLayoutP, SmemLayoutZ,
StagesX, StagesY, StagesZ,
HAS_D, D_HAS_HDIM, HAS_Z>;
using TileScheduler = cutlass::ssd::kernel::PersistentTileScheduler;
using Kernel = cutlass::ssd::kernel::SsdKernelTmaWarpSpecialized<CollectiveMainloop, CollectiveEpilogue, TileScheduler>;
};
}
examples/111_hopper_ssd/kernel/sm90_ssd_kernel_tma_warpspecialized.hpp
+553
View File
@@ -0,0 +1,553 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/arch/reg_reconfig.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/arch.h"
namespace cutlass::ssd::kernel {
using namespace cute;
template<
class CollectiveMainloop,
class CollectiveEpilogue,
class TileScheduler
>
struct SsdKernelTmaWarpSpecialized {
static const int NumLoadWarpGroups = 1;
// hard code
static constexpr int NumMmaWarpGroups = 2;
// TileShape: LDN
using TileShape = typename CollectiveMainloop::TileShape;
// Force to use 1x1x1
using ClusterShape = typename CollectiveMainloop::ClusterShape;
// Pipeline for Tensor X
using MainloopPipelineX = typename CollectiveMainloop::MainloopPipelineX;
using PipelineParamsX = typename MainloopPipelineX::Params;
using PipelineStateX = typename cutlass::PipelineState<MainloopPipelineX::Stages>;
// Pipeline for Tensor Delta && DeltaA
using MainloopPipelineDelta = typename CollectiveMainloop::MainloopPipelineDelta;
using PipelineParamsDelta = typename MainloopPipelineDelta::Params;
using PipelineStateDelta = typename cutlass::PipelineState<MainloopPipelineDelta::Stages>;
// Pipeline for Tensor X
using MainloopPipelineB = typename CollectiveMainloop::MainloopPipelineB;
using PipelineParamsB = typename MainloopPipelineB::Params;
using PipelineStateB = typename cutlass::PipelineState<MainloopPipelineB::Stages>;
// Pipeline for Tensor X
using MainloopPipelineC = typename CollectiveMainloop::MainloopPipelineC;
using PipelineParamsC = typename MainloopPipelineC::Params;
using PipelineStateC = typename cutlass::PipelineState<MainloopPipelineC::Stages>;
// Pipeline for cooperate-warps
using CooperatePipeline = typename CollectiveEpilogue::CooperatePipeline;
using PipelineParamsCo = typename CooperatePipeline::Params;
using PipelineStateCo = typename cutlass::PipelineState<CooperatePipeline::Stages>;
// Pipeline for Tensor D
using EpiloadPipelineD = typename CollectiveEpilogue::EpiloadPipelineD;
using PipelineParamsD = typename EpiloadPipelineD::Params;
using PipelineStateD = typename cutlass::PipelineState<EpiloadPipelineD::Stages>;
// Pipeline for Tensor Z
using EpiloadPipelineZ = typename CollectiveEpilogue::EpiloadPipelineZ;
using PipelineParamsZ = typename EpiloadPipelineZ::Params;
using PipelineStateZ = typename cutlass::PipelineState<EpiloadPipelineZ::Stages>;
struct TensorStorage {
typename CollectiveMainloop::SharedStorage mainloop;
typename CollectiveEpilogue::TensorStorage epilogue;
};
struct SharedStorage {
TensorStorage tensors;
using PipelineStorageX = typename MainloopPipelineX::SharedStorage;
using PipelineStorageDelta = typename MainloopPipelineDelta::SharedStorage;
using PipelineStorageB = typename MainloopPipelineB::SharedStorage;
using PipelineStorageC = typename MainloopPipelineC::SharedStorage;
using PipelineStorageCo = typename CooperatePipeline::SharedStorage;
using PipelineStorageD = typename EpiloadPipelineD::SharedStorage;
using PipelineStorageZ = typename EpiloadPipelineZ::SharedStorage;
// pipeline
alignas(16) PipelineStorageX pipeline_storage_x;
alignas(16) PipelineStorageDelta pipeline_storage_delta;
alignas(16) PipelineStorageB pipeline_storage_b;
alignas(16) PipelineStorageC pipeline_storage_c;
alignas(16) PipelineStorageCo pipeline_storage_co;
alignas(16) PipelineStorageD pipeline_storage_d;
alignas(16) PipelineStorageZ pipeline_storage_z;
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
// [G, B, EH, C, L, D, N]
using ProblemShape = cute::tuple<int, int, int, int, int, int, int>;
struct Arguments {
ProblemShape problem_size;
typename CollectiveMainloop::Arguments mainloop;
typename CollectiveEpilogue::Arguments epilogue;
KernelHardwareInfo hw_info;
};
struct Params {
ProblemShape problem_size;
typename CollectiveMainloop::Params mainloop;
typename CollectiveEpilogue::Params epilogue;
typename TileScheduler::Params tile_scheduler;
};
static const int MinBlocksPerMultiprocessor = 1;
static const int MaxThreadsPerBlock = (NumMmaWarpGroups + NumLoadWarpGroups) * cutlass::NumThreadsPerWarpGroup;
using ArchTag = cutlass::arch::Sm90;
// CTA reconfig (TBD)
static constexpr uint32_t LoadRegisterRequirement = 40 - 2 * 8;
static constexpr uint32_t TotalRegisterSupply = (64*1024 / MaxThreadsPerBlock / MinBlocksPerMultiprocessor / 8) * 8 * MaxThreadsPerBlock / cutlass::NumThreadsPerWarpGroup;
static constexpr uint32_t MmaRegisterRequirement = ((TotalRegisterSupply - LoadRegisterRequirement) / NumMmaWarpGroups / 8) * 8;
// static constexpr uint32_t LoadRegisterRequirement = 40;
// static constexpr uint32_t MmaRegisterRequirement = 232;
static size_t get_workspace_size(Arguments const& args) { return 0; }
static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t) {
return cutlass::Status::kSuccess;
}
static bool can_implement(Arguments const& args) {
return CollectiveMainloop::can_implement(args.problem_size, args.mainloop);
}
static dim3 get_grid_shape(Params const& params) {
return TileScheduler::get_grid_shape(params.tile_scheduler);
}
static dim3 get_block_shape() {
dim3 block(MaxThreadsPerBlock, 1, 1);
return block;
}
static Params to_underlying_arguments(Arguments const& args, void* workspace) {
return Params{
args.problem_size,
CollectiveMainloop::to_underlying_arguments(args.problem_size, args.mainloop, workspace),
CollectiveEpilogue::to_underlying_arguments(args.problem_size, args.epilogue, workspace),
TileScheduler::to_underlying_arguments(args.problem_size, args.hw_info, ClusterShape{}, TileShape{})
};
}
CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
// TBD
enum class WarpGroupRole {
Producer = 0,
Consumer0 = 1,
Consumer1 = 2
};
// TBD
enum class ProducerWarpRole {
LoadX = 0,
LoadDelta = 1,
LoadBC = 2,
LoadZ = 3
};
// Parameters
// [G, B, EH, C, L, D, N]
auto C = get<3>(params.problem_size);
// Shared memory.
auto& storage = *reinterpret_cast<SharedStorage*>(smem);
int lane_idx = cutlass::canonical_lane_idx();
int warp_idx = cutlass::canonical_warp_idx_sync();
int warp_idx_in_warp_group = warp_idx % cutlass::NumWarpsPerWarpGroup;
int warp_group_idx = cutlass::canonical_warp_group_idx();
auto warp_group_role = WarpGroupRole(warp_group_idx);
auto producer_warp_role = ProducerWarpRole(warp_idx_in_warp_group);
int lane_predicate = cute::elect_one_sync();
uint32_t block_rank_in_cluster = cute::block_rank_in_cluster();
// Issue Tma Descriptor Prefetch from a single thread
if ((warp_idx == 0) && lane_predicate) {
CollectiveMainloop::prefetch_tma_descriptors(params.mainloop);
}
// Pipeline (TBD)
PipelineParamsX pipeline_params_x;
pipeline_params_x.transaction_bytes = CollectiveMainloop::kXLoadBytes;
pipeline_params_x.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadX);
pipeline_params_x.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups;
pipeline_params_x.initializing_warp = 4;
PipelineParamsDelta pipeline_params_delta;
pipeline_params_delta.transaction_bytes = CollectiveMainloop::kDeltaLoadBytes + CollectiveMainloop::kDeltaALoadBytes;
pipeline_params_delta.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadDelta);
pipeline_params_delta.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups;
pipeline_params_delta.initializing_warp = 5;
PipelineParamsB pipeline_params_b;
pipeline_params_b.transaction_bytes = CollectiveMainloop::kBLoadBytes;
pipeline_params_b.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadBC);
pipeline_params_b.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups;
pipeline_params_b.initializing_warp = 6;
PipelineParamsC pipeline_params_c;
pipeline_params_c.transaction_bytes = CollectiveMainloop::kCLoadBytes;
pipeline_params_c.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadBC);
pipeline_params_c.num_consumers = cutlass::NumThreadsPerWarpGroup * NumMmaWarpGroups;
pipeline_params_c.initializing_warp = 7;
PipelineParamsCo pipeline_params_co;
pipeline_params_co.producer_arv_count = cutlass::NumThreadsPerWarpGroup;
pipeline_params_co.consumer_arv_count = cutlass::NumThreadsPerWarpGroup;
pipeline_params_co.initializing_warp = 8;
PipelineParamsD pipeline_params_d;
pipeline_params_d.transaction_bytes = CollectiveEpilogue::kEpiloadDBytes;
pipeline_params_d.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadDelta);
pipeline_params_d.num_consumers = cutlass::NumThreadsPerWarpGroup;
pipeline_params_d.initializing_warp = 9;
PipelineParamsZ pipeline_params_z;
pipeline_params_z.transaction_bytes = CollectiveEpilogue::kEpiloadZBytes;
pipeline_params_z.is_leader = lane_predicate && (producer_warp_role == ProducerWarpRole::LoadZ);
pipeline_params_z.num_consumers = cutlass::NumThreadsPerWarpGroup;
pipeline_params_z.initializing_warp = 10;
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadX) {
pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadDelta) {
pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Producer;
pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadBC) {
pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Producer;
pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Producer && producer_warp_role == ProducerWarpRole::LoadZ) {
pipeline_params_z.role = EpiloadPipelineZ::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Consumer0 || warp_group_role == WarpGroupRole::Consumer1) {
pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer;
pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Consumer;
pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer;
pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Consumer;
}
if (warp_group_role == WarpGroupRole::Consumer0) {
pipeline_params_co.role = CooperatePipeline::ThreadCategory::Producer;
}
if (warp_group_role == WarpGroupRole::Consumer1) {
pipeline_params_co.role = CooperatePipeline::ThreadCategory::Consumer;
pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Consumer;
pipeline_params_z.role = EpiloadPipelineZ::ThreadCategory::Consumer;
}
MainloopPipelineX pipeline_x(storage.pipeline_storage_x, pipeline_params_x, Shape<_1,_1,_1>{});
PipelineStateX mainloop_pipe_x_consumer;
PipelineStateX mainloop_pipe_x_producer = cutlass::make_producer_start_state<MainloopPipelineX>();
MainloopPipelineDelta pipeline_delta(storage.pipeline_storage_delta, pipeline_params_delta, Shape<_1,_1,_1>{});
PipelineStateDelta mainloop_pipe_delta_consumer;
PipelineStateDelta mainloop_pipe_delta_producer = cutlass::make_producer_start_state<MainloopPipelineDelta>();
MainloopPipelineB pipeline_b(storage.pipeline_storage_b, pipeline_params_b, Shape<_1,_1,_1>{});
PipelineStateB mainloop_pipe_b_consumer;
PipelineStateB mainloop_pipe_b_producer = cutlass::make_producer_start_state<MainloopPipelineB>();
MainloopPipelineC pipeline_c(storage.pipeline_storage_c, pipeline_params_c, Shape<_1,_1,_1>{});
PipelineStateC mainloop_pipe_c_consumer;
PipelineStateC mainloop_pipe_c_producer = cutlass::make_producer_start_state<MainloopPipelineC>();
CooperatePipeline pipeline_co(storage.pipeline_storage_co, pipeline_params_co);
PipelineStateCo cooperate_pipe_consumer_state;
PipelineStateCo cooperate_pipe_producer_state = cutlass::make_producer_start_state<CooperatePipeline>();
EpiloadPipelineD pipeline_d(storage.pipeline_storage_d, pipeline_params_d, Shape<_1,_1,_1>{});
PipelineStateD epi_load_pipe_d_consumer;
PipelineStateD epi_load_pipe_d_producer = cutlass::make_producer_start_state<EpiloadPipelineD>();
EpiloadPipelineZ pipeline_z(storage.pipeline_storage_z, pipeline_params_z, Shape<_1,_1,_1>{});
PipelineStateZ epi_load_pipe_z_consumer;
PipelineStateZ epi_load_pipe_z_producer = cutlass::make_producer_start_state<EpiloadPipelineZ>();
// Epilogue Store pipeline
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
typename EpiStorePipeline::Params epi_store_pipeline_params;
epi_store_pipeline_params.always_wait = true;
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
// Epilogue Store P pipeline
using EpiStorePPipeline = typename CollectiveEpilogue::StorePPipeline;
typename EpiStorePPipeline::Params epi_store_p_pipeline_params;
epi_store_p_pipeline_params.always_wait = true;
EpiStorePPipeline epi_store_p_pipeline(epi_store_p_pipeline_params);
PipelineState epi_store_p_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePPipeline>();
// We need this to guarantee that the Pipeline init is visible
// To all producers and consumer blocks in the Cluster
// and to finish smem init
if constexpr (size(ClusterShape{}) > 1) {
cute::cluster_arrive_relaxed();
cute::cluster_wait();
}
else {
__syncthreads();
}
// Kernel implement(TBD)
CollectiveMainloop collective_mainloop;
CollectiveEpilogue collective_epilogue;
if (warp_group_role == WarpGroupRole::Producer) {
// disable reg dealloc to enable print
cutlass::arch::warpgroup_reg_dealloc<LoadRegisterRequirement>();
if (producer_warp_role == ProducerWarpRole::LoadX) {
// use local variable to avoid STL/LDL
TileScheduler tile_scheduler{params.tile_scheduler};
auto load_input = collective_mainloop.load_x_init(params.mainloop, params.problem_size);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
// Load X
collective_mainloop.load_x(
blk_coord, params.mainloop, params.problem_size,
pipeline_x, mainloop_pipe_x_producer,
load_input,
storage.tensors.mainloop
);
}
collective_mainloop.load_x_tail(
pipeline_x, mainloop_pipe_x_producer
);
}
else if (producer_warp_role == ProducerWarpRole::LoadDelta) {
TileScheduler tile_scheduler{params.tile_scheduler};
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
auto blk_coord_eh = tile_scheduler.get_block_coord_eh();
// Epiload
collective_epilogue.load_d(
blk_coord_eh, params.epilogue, params.problem_size,
pipeline_d, epi_load_pipe_d_producer,
storage.tensors.mainloop
);
// Load Delta
// Load DeltaA
collective_mainloop.load_delta(
blk_coord, params.mainloop, params.problem_size,
pipeline_delta, mainloop_pipe_delta_producer,
storage.tensors.mainloop
);
}
collective_mainloop.load_delta_tail(
pipeline_delta, mainloop_pipe_delta_producer
);
}
else if (producer_warp_role == ProducerWarpRole::LoadBC) {
TileScheduler tile_scheduler{params.tile_scheduler};
auto load_input_b = collective_mainloop.load_b_init(params.mainloop, params.problem_size);
auto load_input_c = collective_mainloop.load_c_init(params.mainloop, params.problem_size);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord_b();
// Load B
collective_mainloop.load_b_c(
blk_coord, params.mainloop, params.problem_size,
pipeline_b, mainloop_pipe_b_producer,
pipeline_c, mainloop_pipe_c_producer,
load_input_b,
load_input_c,
storage.tensors.mainloop
);
}
collective_mainloop.load_b_c_tail(
pipeline_b, mainloop_pipe_b_producer,
pipeline_c, mainloop_pipe_c_producer
);
}
else if (producer_warp_role == ProducerWarpRole::LoadZ) {
// use local variable to avoid STL/LDL
TileScheduler tile_scheduler{params.tile_scheduler};
auto load_input = collective_epilogue.load_z_init(params.epilogue, params.problem_size);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
// Load X
collective_epilogue.load_z(
blk_coord, params.epilogue, params.problem_size,
pipeline_z, epi_load_pipe_z_producer,
load_input,
storage.tensors.epilogue
);
}
collective_epilogue.load_z_tail(
pipeline_z, epi_load_pipe_z_producer
);
}
}
// Warpgroup1 for Intra
// Warpgroup2 for Inter
else if (warp_group_role == WarpGroupRole::Consumer0) {
TileScheduler tile_scheduler{params.tile_scheduler};
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
for (int chunk = 0; chunk < C; ++chunk) {
auto blk_coord = tile_scheduler.get_block_coord();
// IntraBMM1
// Wait B
// Wait C
auto [tIntra1] = collective_mainloop.mma_intra_1(
chunk,
pipeline_b, mainloop_pipe_b_consumer,
pipeline_c, mainloop_pipe_c_consumer,
storage.tensors.mainloop
);
// Pre Intra2
auto [tPreIntra2] = collective_mainloop.pre_intra_2(
chunk,
pipeline_delta, mainloop_pipe_delta_consumer,
tIntra1,
storage.tensors.mainloop
);
// IntraBMM2
auto [tIntra2] = collective_mainloop.mma_intra_2(
chunk,
pipeline_x, mainloop_pipe_x_consumer,
tPreIntra2,
storage.tensors.mainloop
);
collective_epilogue.store_intra(
chunk, blk_coord, params.epilogue, params.problem_size,
pipeline_co, cooperate_pipe_producer_state,
tIntra2,
typename CollectiveMainloop::TiledMmaIntra2{},
storage.tensors.epilogue
);
}
}
}
else if (warp_group_role == WarpGroupRole::Consumer1) {
TileScheduler tile_scheduler{params.tile_scheduler};
cutlass::arch::warpgroup_reg_alloc<MmaRegisterRequirement>();
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto [tState] = collective_mainloop.state_init(storage.tensors.mainloop);
auto blk_coord_eh = tile_scheduler.get_block_coord_eh();
bool is_first_iteration = true;
for (int chunk = 0; chunk < C; ++chunk) {
auto blk_coord = tile_scheduler.get_block_coord();
// Pre Inter1
// Wait delta
// Wait deltaA
auto [tPreInter1, last_column] = collective_mainloop.pre_inter_1(
chunk,
pipeline_b, mainloop_pipe_b_consumer,
pipeline_delta, mainloop_pipe_delta_consumer,
storage.tensors.mainloop
);
// InterBMM1
// Wait X
auto [tInter1] = collective_mainloop.mma_inter_1(
chunk,
pipeline_x, mainloop_pipe_x_consumer,
tPreInter1,
storage.tensors.mainloop
);
collective_mainloop.pre_inter_2(
last_column,
tInter1,
tState
);
// InterBMM2
auto [tInter2, tDelta] = collective_mainloop.mma_inter_2(
chunk,
pipeline_c, mainloop_pipe_c_consumer,
pipeline_delta, mainloop_pipe_delta_consumer,
storage.tensors.mainloop
);
// Pre Inter2
collective_mainloop.post_inter_2(
tState,
storage.tensors.mainloop
);
// update_d
auto [tD] = collective_epilogue.update_d(
blk_coord_eh, params.epilogue,
is_first_iteration,
pipeline_d, epi_load_pipe_d_consumer,
typename CollectiveMainloop::TiledMmaInter2{},
storage.tensors.mainloop
);
is_first_iteration = false;
// Epilogue TensorY store
collective_epilogue.store(
chunk, blk_coord, params.epilogue, params.problem_size,
epi_store_pipeline, epi_store_pipe_producer_state,
pipeline_co, cooperate_pipe_consumer_state,
pipeline_x, mainloop_pipe_x_consumer,
pipeline_z, epi_load_pipe_z_consumer,
tInter2, tDelta, tD,
typename CollectiveMainloop::TiledMmaInter2{},
storage.tensors.epilogue, storage.tensors.mainloop
);
}
if constexpr (CollectiveEpilogue::D_HAS_HDIM) {
// update the barrier
pipeline_d.consumer_release(epi_load_pipe_d_consumer);
++epi_load_pipe_d_consumer;
}
auto blk_coord = tile_scheduler.get_block_coord();
// Epilogue Fstate store
collective_epilogue.store_p(
blk_coord, params.epilogue, params.problem_size,
epi_store_p_pipeline, epi_store_p_pipe_producer_state,
storage.tensors.mainloop
);
}
}
}
};
} // namespace cutlass::fmha::kernel
examples/111_hopper_ssd/kernel/sm90_ssd_tile_scheduler.hpp
+132
View File
@@ -0,0 +1,132 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/kernel_hardware_info.h"
namespace cutlass::ssd::kernel {
////////////////////////////////////////////////////////////////////////////////
struct PersistentTileScheduler {
struct Params {
int num_blocks;
int num_groups;
FastDivmod divmod_eh;
FastDivmod divmod_ngroup_ratio;
KernelHardwareInfo hw_info;
};
int block_idx = 0;
Params params;
CUTLASS_DEVICE
PersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {}
template<class ProblemSize, class ClusterShape, class TileShape>
static Params to_underlying_arguments(
ProblemSize const& problem_size, KernelHardwareInfo hw_info,
ClusterShape const& cluster_shape, TileShape const& tile_shape)
{
using namespace cute;
auto [G, B, EH, C, L, D, N] = problem_size;
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = hw_info.sm_count;
if (sm_count <= 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
}
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
hw_info.sm_count = sm_count;
int num_blocks = B * EH;
int ngroup_ratio = EH / G;
return Params {
num_blocks,
G,
{EH},
{ngroup_ratio},
hw_info
};
}
static dim3 get_grid_shape(Params const& params) {
dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1);
return grid;
}
CUTLASS_DEVICE
bool is_valid() {
return block_idx < params.num_blocks;
}
CUTLASS_DEVICE
auto get_block_coord() {
return block_idx;
}
CUTLASS_DEVICE
auto get_block_coord_b() {
using namespace cute;
int eh_idx, b_idx;
int g_idx, rest_idx;
params.divmod_eh(b_idx, eh_idx, block_idx);
params.divmod_ngroup_ratio(g_idx, rest_idx, eh_idx);
return (params.num_groups * b_idx + g_idx);
}
CUTLASS_DEVICE
auto get_block_coord_eh() {
using namespace cute;
int eh_idx, b_idx;
params.divmod_eh(b_idx, eh_idx, block_idx);
return eh_idx;
}
CUTLASS_DEVICE
PersistentTileScheduler& operator++() {
block_idx += gridDim.x;
return *this;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::ssd::kernel
examples/111_hopper_ssd/reference/reference_ssd.hpp
+345
View File
@@ -0,0 +1,345 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cute/tensor.hpp"
// training or inference phase (not used yet)
// PHASE 0 : training
// PHASE 1 : inference
#define PHASE 0
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
bool transA,
bool transB,
class Element,
class TensorA,
class TensorB,
class TensorC
>
void mma(
TensorA tA,
TensorB tB,
TensorC tC) {
using namespace cute;
int M = transA ? int(shape<1>(tA)) : int(shape<0>(tA));
int N = transB ? int(shape<1>(tB)) : int(shape<0>(tB));
int K = transA ? int(shape<0>(tA)) : int(shape<1>(tA));
for (int mi = 0; mi < M; ++mi) {
for (int ni = 0; ni < N; ++ni) {
for (int ki = 0; ki < K; ++ki) {
float a = static_cast<float>(Element(transA ? tA(ki, mi) : tA(mi, ki)));
float b = static_cast<float>(Element(transB ? tB(ki, ni) : tB(ni, ki)));
tC(mi, ni) += a * b;
}
}
}
}
template<
class Element,
class Tensor
>
auto segsum(Tensor tensor) {
using namespace cute;
auto C = shape<0>(tensor);
auto L = shape<1>(tensor);
auto cum_sum = make_tensor<float>(make_shape(C,L));
// cum_sum
for (int ci = 0; ci < C; ++ci) {
for (int li = 0; li < L; ++li) {
if (li == 0) {
cum_sum(ci, li) = tensor(ci, li);
}
else {
cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li);
}
}
}
auto seg_sum_out = make_tensor<float>(make_shape(C, L, L));
// seg_sum
// [ 1, 0, 0]
// [ e^y, 1, 0]
// [e^(y+z), e^z, 1]
CUTLASS_PRAGMA_UNROLL
for (int ci = 0; ci < C; ++ci) {
for (int i = 0; i < L; ++i) {
for (int j = 0; j < L; ++j) {
if (j < i) {
float tmp = static_cast<float>(cum_sum(ci, i)) - static_cast<float>(cum_sum(ci, j));
seg_sum_out(ci, i, j) = expf(tmp);
}
else if (j == i) {
seg_sum_out(ci, i, j) = 1.f;
}
else {
seg_sum_out(ci, i, j) = 0.f;
}
}
}
}
return seg_sum_out;
}
template<
class Element,
class Tensor
>
auto cumsum(
Tensor tensor) {
using namespace cute;
auto C = shape<0>(tensor);
auto L = shape<1>(tensor);
auto cum_sum = make_tensor<float>(make_shape(C,L));
auto cum_sum_out = make_tensor<Element>(make_shape(C,L));
auto cum_sum_exp_out = make_tensor<float>(make_shape(C, L));
auto cum_sum_exp_out_last = make_tensor<float>(make_shape(C, L));
auto last_column = make_tensor<float>(make_shape(C));
// [x, x+y, x+y+z, ..]
CUTLASS_PRAGMA_UNROLL
for (int ci = 0; ci < C; ++ci) {
for (int li = 0; li < L; ++li) {
if (li == 0) {
cum_sum(ci, li) = tensor(ci, li);
}
else {
cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li);
}
// cum_sum_out(ci, li) = static_cast<Element>(cum_sum(ci, li));
}
}
CUTLASS_PRAGMA_UNROLL
for (int ci = 0; ci < C; ++ci) {
last_column(ci) = static_cast<float>(cum_sum(ci, L-1));
CUTLASS_PRAGMA_UNROLL
for (int li = 0; li < L; ++li) {
cum_sum_exp_out_last(ci, li) = expf(static_cast<float>(last_column(ci) - cum_sum(ci, li)));
cum_sum_exp_out(ci, li) = expf(static_cast<float>(cum_sum(ci, li)));
}
}
return make_tuple(cum_sum_exp_out_last, last_column, cum_sum_exp_out);
}
template<
bool HAS_D,
bool D_HAS_HDIM,
bool HAS_Z,
class TensorY,
class TensorF,
class TensorX,
class TensorDelta,
class TensorDeltaA,
class TensorB,
class TensorC,
class TensorD,
class TensorZ,
class Params
>
void ssd_reference_impl(
TensorY mY, TensorF mF,
TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA,
TensorB mB, TensorC mC, TensorD mD, TensorZ mZ,
Params params) {
using namespace cute;
using Element = typename Params::Element;
using ElementAcc = typename Params::ElementAcc;
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
// d [ eh, d]
auto [G, B, EH, C, L, D, N] = params.get_problem_shape();
int group_ratio = EH / G;
for (int b = 0; b < B; ++b) {
for (int eh = 0; eh < EH; ++eh) {
int g = eh / group_ratio;
auto tY = mY(b,eh,_,_,_);
auto tF = mF(b,eh,_,_);
auto tX = mX(b,eh,_,_,_);
auto tDelta = mDelta(b,eh,_,_);
auto tDeltaA = mDeltaA(b,eh,_,_);
auto tB = mB(b,g,_,_,_);
auto tC = mC(b,g,_,_,_);
auto tD = mD(eh,_);
auto tZ = mZ(b,eh,_,_,_);
// IntraBMM1 BxC, LxLxN, NT
// B: [n, c, l]
// C: [n, c, l]
// O: [c, l, l]
auto tIntraBMM1_out = make_tensor<float>(make_shape(C,L,L));
for (int ci = 0; ci < C; ++ci) {
mma<true,true,Element>(tC(_,ci,_), tB(_,ci,_), tIntraBMM1_out(ci,_,_));
}
// Pre_IntraBMM2 DeltaA_IntraBMM2 x Delta x IntraBMM_out
// DeltaA_xxx : [c, l, l]
// Delta : [c, l, _]
// IntraBMM1_out: [c, l, l]
auto tDeltaA_IntraBMM2 = segsum<Element>(tDeltaA);
auto tIntraBMM2_inp = make_tensor<float>(make_shape(C, L, L));
for (int ci = 0; ci < C; ++ci) {
for (int i = 0; i < L; ++i) {
for (int j = 0; j < L; ++j) {
tIntraBMM2_inp(ci, i, j) = tDeltaA_IntraBMM2(ci, i, j) * tDelta(ci, j) * tIntraBMM1_out(ci, i, j);
}
}
}
// IntraBMM2 IntraBMM2_inp x X, LxDxL, TT
// IntraBMM2_inp: [c, l, l]
// X : [d, c, l]
// IntraBMM2_out: [c, l, d]
auto tIntraBMM2_out = make_tensor<float>(make_shape(C,L,D));
for (int ci = 0; ci < C; ++ci) {
mma<false,false,Element>(tIntraBMM2_inp(ci,_,_), tX(_,ci,_), tIntraBMM2_out(ci,_,_));
}
// Pre_InterBMM1 DeltaA_InterBMM1 x Delta x B
// DeltaA_xxx : [c, l]
// Delta : [c, l]
// IntraBMM1_out: [c, n, l]
auto [tDeltaA_InterBMM1, tLast, tCumsum_exp] = cumsum<Element>(tDeltaA);
auto tInterBMM1_inp = make_tensor<float>(make_shape(C, N, L));
for (int ci = 0; ci < C; ++ci) {
for (int i = 0; i < N; ++i) {
for (int j = 0; j < L; ++j) {
tInterBMM1_inp(ci, i, j) = tDeltaA_InterBMM1(ci, j) * tDelta(ci, j) * tB(i, ci, j);
}
}
}
// InterBMM1 InterBMM1_inp x X, NxDxL, swapAB, TT
// InterBMM1_inp: [c, n, l]
// X : [d, c, l]
// InterBMM1_out: [c, n, d]
auto tInterBMM1_out = make_tensor<float>(make_shape(C,N,D));
for (int ci = 0; ci < C; ++ci) {
mma<false,false,Element>(tInterBMM1_inp(ci,_,_), tX(_,ci,_), tInterBMM1_out(ci,_,_));
}
// Initialize state
// PreInterBMM2
// InterBMM1_out: [c, n, d]
// Last : [c]
auto tInterBMM2_inp = make_tensor<float>(make_shape(C, N, D));
for (int ci = 0; ci < C; ++ci) {
for (int ni = 0; ni < N; ++ ni){
for (int di = 0; di < D; ++di) {
if (ci == 0) {
tInterBMM2_inp(ci, ni, di) = 0;
}
else {
tInterBMM2_inp(ci, ni, di) = tInterBMM1_out(ci - 1, ni, di) + expf(tLast(ci - 1)) * tInterBMM2_inp(ci - 1, ni, di);
}
}
}
}
// InterBMM2 InterBMM2_inp x C, LxDxN, NT
// C : [n, c, l]
// InterBMM2_inp: [c, n, d]
// InterBMM2_out: [c, l, d]
auto tInterBMM2_out = make_tensor<float>(make_shape(C,L,D));
for (int ci = 0; ci < C; ++ci) {
mma<true,true,Element>(tC(_,ci,_), tInterBMM2_inp(ci,_,_), tInterBMM2_out(ci,_,_));
}
// Epilogue Cumsum_exp x InterBMM2_out + IntraBMM2_out
// InterBMM2_out: [c, l, d]
// IntraBMM2_out: [c, l, d]
// Cumsum_exp : [c, l]
for (int ci = 0; ci < C; ++ci) {
for (int li = 0; li < L; ++li) {
for (int di = 0; di < D; ++di) {
float y = tInterBMM2_out(ci, li, di) * tCumsum_exp(ci, li) + tIntraBMM2_out(ci, li, di);
float scale;
if constexpr (D_HAS_HDIM) {
scale = static_cast<float>(tD(di));
}
else {
scale = static_cast<float>(tD(_0{}));
}
if constexpr (HAS_D) {
y = y + static_cast<float>(tX(di, ci, li)) * scale;
}
else {
y = y;
}
if constexpr (HAS_Z) {
float z = static_cast<float>(tZ(di, ci, li));
// y = y * z * (1 / (1 + exp(-z)));
y = y * z * (1 / (1 + exp(-z)));
}
tY(di, ci, li) = static_cast<typename Params::Element>(y);
}
}
}
// Epilogue Fstate(last C)
for (int ni = 0; ni < N; ++ ni){
for (int di = 0; di < D; ++di) {
tF(di, ni) = static_cast<typename Params::Element>(tInterBMM1_out(C - 1, ni, di) + expf(tLast(C - 1)) * tInterBMM2_inp(C - 1, ni, di));
}
}
}
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
bool HAS_D,
bool D_HAS_HDIM,
bool HAS_Z,
class TensorY,
class TensorF,
class TensorX,
class TensorDelta,
class TensorDeltaA,
class TensorB,
class TensorC,
class TensorD,
class TensorZ,
class Params
>
void ssd_reference(
TensorY mY, TensorF mF,
TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA,
TensorB mB, TensorC mC, TensorD mD, TensorZ mZ,
Params params) {
ssd_reference_impl<HAS_D, D_HAS_HDIM, HAS_Z>(mY, mF, mX, mDelta, mDeltaA, mB, mC, mD, mZ, params);
}
/////////////////////////////////////////////////////////////////////////////////////////////////
examples/111_hopper_ssd/reference/reference_ssd_cumsum.hpp
+194
View File
@@ -0,0 +1,194 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include <algorithm>
#include <random>
#include "cutlass/coord.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/tensor_view.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/host/gemm.h"
#include "cutlass/arch/arch.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/cuda_host_adapter.hpp"
#include "cute/int_tuple.hpp"
#include "cute/atom/mma_traits_sm100.hpp"
#include "cute/util/debug.hpp"
#include "cute/config.hpp"
namespace cutlass::ssd::kernel {
using namespace cute;
template<
class Element_,
class ElementD_,
class TileShape_>
struct CumsumKernel {
using Element = Element_;
using ElementD = ElementD_;
using TileShape = TileShape_; // L,D,N
// Required by `device_kernel`
static constexpr int MaxThreadsPerBlock = 128;
static constexpr int MinBlocksPerMultiprocessor = 1;
using ArchTag = arch::Sm90;
static constexpr int AlignmentBytes = 16;
struct SharedStorage {
/* empty, no smem needed */
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
struct TransformArguments {
const Element* ptr_DeltaA;
ElementD* ptr_Cumsum;
};
struct TransformParams {
const Element* ptr_DeltaA;
ElementD* ptr_Cumsum;
};
using ProblemShape = cute::tuple<int, int, int, int>; // b, eh, c, l
struct Arguments {
ProblemShape problem_shape{};
TransformArguments transform{};
KernelHardwareInfo hw_info{};
};
struct Params {
ProblemShape problem_shape{};
TransformParams transform{};
KernelHardwareInfo hw_info{};
};
static Params
to_underlying_arguments(Arguments const& args, void* workspace) {
return Params{
ProblemShape{args.problem_shape},
TransformParams{args.transform.ptr_DeltaA, args.transform.ptr_Cumsum},
KernelHardwareInfo{args.hw_info}};
}
static Status
can_implement(Arguments const& args) {
return Status::kSuccess;
}
static size_t
get_workspace_size(Arguments const& args) {
return size_t(0);
}
static Status
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
CudaHostAdapter *cuda_adapter = nullptr) {
return Status::kSuccess;
}
static dim3
get_grid_shape(Params const& params) {
auto [B, EH, C, L] = params.problem_shape;
return dim3(B*EH, 1, 1);
}
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
CUTE_HOST_DEVICE
void
operator()(Params params, [[maybe_unused]] char* smem_buf = nullptr) {
auto [B, EH, C, L] = params.problem_shape;
auto layout = make_layout(make_shape(L, C, EH*B));
auto mD_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_DeltaA), make_layout(reverse(layout.shape()), reverse(layout.stride())));
auto mC_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_Cumsum), make_layout(reverse(layout.shape()), reverse(layout.stride())));
auto cD_bcl = make_identity_tensor(shape(mD_bcl));
int blk_idx = blockIdx.x;
int thread_idx = threadIdx.x;
auto tD = logical_divide(mD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_);
auto tC = logical_divide(mC_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_);
auto cD = logical_divide(cD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_);
static constexpr int NumPacked = AlignmentBytes / sizeof(ElementD);
using PackedTypeDeltaA = uint_bit_t<sizeof_bits_v<Element> * NumPacked>;
using PackedTypeCumsum = uint_bit_t<sizeof_bits_v<ElementD> * NumPacked>;
#if 0
if (thread_idx % 128 == 0 && blk_idx == 0) {
print("tD : ");print(tD);print("\n");
print("tC : ");print(tC);print("\n");
print("cD : ");print(cD);print("\n");
}
#endif
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < shape<0>(tD); ++i) {
float last_element = 0.f;
auto crd = cD(i,_0{});
auto tD_recast = recast<PackedTypeDeltaA>(tD);
auto tC_recast = recast<PackedTypeCumsum>(tC);
if (elem_less(crd, shape(mD_bcl))) {
for (int j = 0; j < shape<1>(tD_recast); ++j) {
auto tD_slice = make_tensor<Element>(make_shape(Int<NumPacked>{}));
auto tC_slice = make_tensor<ElementD>(make_shape(Int<NumPacked>{}));
auto tD_slice_recast = recast<PackedTypeDeltaA>(tD_slice);
auto tC_slice_recast = recast<PackedTypeCumsum>(tC_slice);
tD_slice_recast(_0{}) = tD_recast(i,j);
for (int k = 0; k < NumPacked; ++ k) {
last_element += static_cast<float>(tD_slice(k));
tC_slice(k) = static_cast<ElementD>(last_element);
}
tC_recast(i,j) = tC_slice_recast(_0{});
}
}
}
}
private:
};
} // End namespace cutlass
examples/112_blackwell_ssd/112_blackwell_ssd.cu
+850
View File
@@ -0,0 +1,850 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#include <iostream>
#include "cutlass/util/command_line.h"
#include "cutlass/cutlass.h"
#include "cute/tensor.hpp"
#include "cute/layout.hpp"
#include "cutlass/kernel_hardware_info.hpp"
#include "thrust/universal_vector.h"
#include "cutlass/util/distribution.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/util/reference/host/tensor_fill.h"
#include "cutlass/util/reference/host/tensor_copy.h"
#include "cutlass/util/reference/host/tensor_compare.h"
#include "cutlass/util/reference/device/tensor_fill.h"
#include "cutlass/util/reference/device/tensor_compare.h"
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
#include "reference/reference_ssd_cumsum.hpp"
#include "reference/reference_ssd.hpp"
#include "cutlass/transform/device/transform_universal_adapter.hpp"
#include "device/ssd.hpp"
#include "kernel/sm100_ssd_kernel_builder.hpp"
using namespace cute;
// Command line options parsing
struct Options {
using Element = cutlass::bfloat16_t;
using ElementAcc = float;
using ElementDA = float;
static constexpr bool D_HAS_HDIM = true;
static constexpr bool HAS_D = true;
// Blackwell SSD doesn't support Z now(huge perf drop).
static constexpr bool HAS_Z = false;
bool help;
bool error;
// All static number now
int G = 2;
int B = 3;
int E = 2;
int H = 2;
// Reference kernel doesn't support dynamic C now.
static constexpr auto C = Int<8>{};
static constexpr auto D = Int<64>{};
static constexpr auto L = Int<128>{};
static constexpr auto N = Int<128>{};
int EH = E * H;
int iterations;
bool verify;
bool verbose;
int warmups;
bool measure;
Options():
help(false),
error(false),
iterations(1), verify(true),
measure(false), warmups(3)
{}
// Parses the command line
void parse(int argc, char const **args) {
cutlass::CommandLine cmd(argc, args);
Options defaults;
if (cmd.check_cmd_line_flag("help")) {
help = true;
return;
}
cmd.get_cmd_line_argument("iterations", iterations, defaults.iterations);
cmd.get_cmd_line_argument("G", G, defaults.G);
cmd.get_cmd_line_argument("B", B, defaults.B);
cmd.get_cmd_line_argument("E", E, defaults.E);
cmd.get_cmd_line_argument("H", H, defaults.H);
verbose = cmd.check_cmd_line_flag("verbose");
verify = !(cmd.check_cmd_line_flag("without_verify"));
EH = E*H;
if (iterations > 1) {
measure = true;
verbose = true;
}
auto problem_shape = cute::make_tuple(G, B, EH, C, L, D, N);
cute::print("problem_shape : "); cute::print(problem_shape); cute::print("\n");
}
/// Prints the usage statement.
std::ostream & print_usage(std::ostream &out) const {
out << "112_blackwell_ssd\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --iterations=<int> Benchmarking iterations.\n"
<< " --without_verify Don't verify the results.\n"
<< " --verbose Print execution time per kernel\n"
<< " --G=<int> Group\n"
<< " --B=<int> Batch\n"
<< " --E=<int> Expanded factor\n"
<< " --H=<int> Number of heads\n"
<< "\n";
return out;
}
auto get_problem_shape() const {
return cute::make_tuple(G, B, EH, C, L, D, N);
}
// acceptable layout by cuDNN
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
auto layoutX() const {
auto layout = make_layout(make_shape(L, C, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutDelta() const {
auto layout = make_layout(make_shape(L, C, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutDeltaA() const {
auto layout = make_layout(make_shape(L, C, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutB() const {
auto layout = make_layout(make_shape(L, C, N, G, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutC() const {
auto layout = make_layout(make_shape(L, C, N, G, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutY() const {
auto layout = make_layout(make_shape(L, C, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutF() const {
auto layout = make_layout(make_shape(N, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
auto layoutD() const {
if constexpr (D_HAS_HDIM) {
auto layout = make_layout(make_shape(D, EH));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
else {
auto layout = make_layout(make_shape(Int<1>{}, EH));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
}
auto layoutZ() const {
auto layout = make_layout(make_shape(L, C, D, EH, B));
return make_layout(reverse(layout.shape()), reverse(layout.stride()));
}
// transformed layout for kernel parameters
auto layoutX_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B));
return make_layout(
make_shape(D,L,int32_t(C),EH*B),
make_stride(
stride<2>(layout),
stride<0>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutB_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),N,G*B));
return make_layout(
make_shape(L,N,int32_t(C),G*B),
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutC_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),N,G*B));
return make_layout(
make_shape(L,N,int32_t(C),G*B),
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutDelta_transformed() const {
return make_layout(make_shape(L,int32_t(C),EH*B));
}
auto layoutY_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B));
return make_layout(
make_shape(L,D,int32_t(C),EH*B), // (M,K,L,...)
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
auto layoutF_transformed() const {
auto layout = make_layout(make_shape(N,D,EH*B));
return make_layout(
make_shape(D,N,EH*B),
make_stride(
stride<1>(layout),
stride<0>(layout),
stride<2>(layout)
)
);
}
auto layoutD_transformed() const {
if constexpr (D_HAS_HDIM) {
return make_layout(make_shape(D, EH));
}
else {
return make_layout(make_shape(Int<1>{}, EH));
}
}
auto layoutZ_transformed() const {
auto layout = make_layout(make_shape(L,int32_t(C),D,EH*B));
return make_layout(
make_shape(L,D,int32_t(C),EH*B),
make_stride(
stride<0>(layout),
stride<2>(layout),
stride<1>(layout),
stride<3>(layout)
)
);
}
};
template <typename Element>
static void
initialize_values(
thrust::universal_vector<Element>& dst_ptr,
cutlass::Distribution::Kind dist_kind,
uint64_t seed,
Element var = Element(1.f)) {
if (cutlass::Distribution::Uniform == dist_kind) {
int scope = 2;
cutlass::reference::host::BlockFillRandomUniform(
dst_ptr.data().get(), dst_ptr.size(), seed, scope, -scope, 0);
}
else if (cutlass::Distribution::AllZeros == dist_kind) {
cutlass::reference::host::BlockFillRandomUniform(
dst_ptr.data().get(), dst_ptr.size(), seed, 0, 0, 0);
}
else if (cutlass::Distribution::AllOnes == dist_kind) {
cutlass::reference::host::BlockFillRandomUniform(
dst_ptr.data().get(), dst_ptr.size(), seed, 1, 1, 0);
}
else if (cutlass::Distribution::Gaussian == dist_kind) {
cutlass::reference::device::BlockFillRandomGaussian(
dst_ptr.data().get(), dst_ptr.size(), seed, (Element) 0, var);
}
else if (cutlass::Distribution::Sequential == dist_kind) {
cutlass::reference::host::BlockFillSequential(dst_ptr.data().get(), dst_ptr.size());
}
else {
std::cerr << "Invalid distribution kind!\n.";
exit(1);
}
}
template <
class Options_
>
struct TestBed {
using Option = Options_;
using Element = typename Option::Element;
using ElementDA = typename Option::ElementDA;
using ElementAcc = typename Option::ElementAcc;
thrust::universal_vector<Element> tensor_X;
thrust::universal_vector<Element> tensor_DeltaA;
thrust::universal_vector<ElementDA> tensor_DeltaA_cumsum;
thrust::universal_vector<Element> tensor_Delta;
thrust::universal_vector<Element> tensor_B;
thrust::universal_vector<Element> tensor_C;
thrust::universal_vector<Element> tensor_D;
thrust::universal_vector<Element> tensor_Y;
thrust::universal_vector<Element> tensor_Z;
thrust::universal_vector<Element> tensor_Y_ref_0;
thrust::universal_vector<Element> tensor_Y_ref_1;
thrust::universal_vector<Element> tensor_F;
thrust::universal_vector<Element> tensor_F_ref_0;
thrust::universal_vector<Element> tensor_F_ref_1;
cutlass::Distribution::Kind init_X = cutlass::Distribution::Uniform;
cutlass::Distribution::Kind init_DeltaA = cutlass::Distribution::Gaussian;
cutlass::Distribution::Kind init_Delta = cutlass::Distribution::Gaussian;
cutlass::Distribution::Kind init_B = cutlass::Distribution::Uniform;
cutlass::Distribution::Kind init_C = cutlass::Distribution::Uniform;
using TileShape = decltype(make_shape(Options::L, Options::D, Options::N)); // (L, D, N)
using SsdOperation = cutlass::ssd::device::SSD<
typename cutlass::ssd::kernel::Sm100SsdBuilder<
Element, ElementDA, ElementAcc, Element,
TileShape,
Option::HAS_D, Option::D_HAS_HDIM
>::Kernel>;
using CumsumKenrel = cutlass::ssd::kernel::CumsumKernel<Element, ElementDA, TileShape>;
using CumsumOperation = cutlass::transform::device::TransformUniversalAdapter<CumsumKenrel>;
bool initialize(Options const& options, const cutlass::KernelHardwareInfo& hw_info, uint64_t seed = 2024) {
auto [g, b, eh, c, l, d, n] = options.get_problem_shape();
assert(g == 1 && "Only group size == 1 is supported") ;
auto size_X = b * eh * c * l * d;
auto size_DeltaA = b * eh * c * l;
auto size_Delta = b * eh * c * l;
auto size_B = g * b * c * n * l;
auto size_C = g * b * c * n * l;
auto size_Y = b * eh * c * l * d;
auto size_F = b * eh * d * n;
tensor_X .resize(sizeof(Element) * size(options.layoutX()));
tensor_DeltaA .resize(sizeof(Element) * size(options.layoutDeltaA()));
tensor_Delta .resize(sizeof(Element) * size(options.layoutDelta()));
tensor_B .resize(sizeof(Element) * size(options.layoutB()));
tensor_C .resize(sizeof(Element) * size(options.layoutC()));
tensor_D .resize(sizeof(Element) * size(options.layoutD()));
tensor_Z .resize(sizeof(Element) * size(options.layoutZ()));
tensor_Y .resize(sizeof(Element) * size(options.layoutY()));
tensor_Y_ref_0.resize(sizeof(Element) * size(options.layoutY()));
tensor_Y_ref_1.resize(sizeof(Element) * size(options.layoutY()));
tensor_F .resize(sizeof(Element) * size(options.layoutF()));
tensor_F_ref_0.resize(sizeof(Element) * size(options.layoutF()));
tensor_F_ref_1.resize(sizeof(Element) * size(options.layoutF()));
tensor_DeltaA_cumsum.resize(sizeof(ElementDA) * size(options.layoutDeltaA()));
// Limit distribution to reduce skew between hosts and devices
initialize_values(tensor_X, init_X, seed);
initialize_values(tensor_DeltaA, init_DeltaA, seed + 1, Element(0.05f));
initialize_values(tensor_Delta, init_Delta, seed + 3, Element(0.05f));
initialize_values(tensor_B, init_B, seed + 5);
initialize_values(tensor_C, init_C, seed + 7);
initialize_values(tensor_D, init_C, seed + 9);
initialize_values(tensor_Z, init_X, seed);
cudaError_t result;
result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the Initialization kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
}
// apply cumsum(device) before kernel launch
typename CumsumOperation::Arguments arguments{
make_shape(int(b), int(eh), int(c), int(l)),
{
tensor_DeltaA.data().get(),
tensor_DeltaA_cumsum.data().get(),
},
hw_info
};
CumsumOperation op;
size_t workspace_size = CumsumOperation::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
// may be used uninitialized
cudaEvent_t start;
cudaEvent_t end;
cudaEventCreate(&start);
cudaEventCreate(&end);
// warm up
if (options.measure) {
for (int i = 0; i < options.warmups; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
}
result = cudaEventRecord(start);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
// Run
for (int i = 0; i < options.iterations; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
result = cudaEventRecord(end);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
float runtime_ms = 0;
result = cudaEventElapsedTime(&runtime_ms, start, end);
if (result != cudaSuccess) {
std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
runtime_ms /= static_cast<float>(options.iterations);
if (options.verbose) {
printf("[iters = %d, warmups = %d] cumsum kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms);
}
return true;
}
bool sufficient() const {
int device_idx;
cudaError_t result = cudaGetDevice(&device_idx);
if (result != cudaSuccess) {
throw std::runtime_error("cudaGetDevice() API call failed.");
}
int max_smem_size;
result = cudaDeviceGetAttribute(&max_smem_size, cudaDevAttrMaxSharedMemoryPerBlockOptin, device_idx);
if (result != cudaSuccess) {
throw std::runtime_error("cudaDeviceGetAttribute() failed");
}
return true;
}
bool run(Options const& options, const cutlass::KernelHardwareInfo& hw_info) {
if (!sufficient()) {
std::cerr << "Test waived due to insufficient CUDA device.\n";
return true;
}
if (!initialize(options, hw_info)) {
std::cerr << "Failed to initialize the test.\n";
return true;
};
auto [g, b, eh, c, l, d, n] = options.get_problem_shape();
typename SsdOperation::Arguments arguments{
make_shape(int(g), int(b), int(eh), int(c), int(l), int(d), int(n)),
{
tensor_X.data().get(),
tensor_DeltaA_cumsum.data().get(),
tensor_Delta.data().get(),
tensor_B.data().get(),
tensor_C.data().get(),
options.layoutX_transformed(),
options.layoutB_transformed(),
options.layoutC_transformed(),
options.layoutDelta_transformed()
},
{
tensor_Y.data().get(),
tensor_F.data().get(),
tensor_D.data().get(),
// tensor_Z.data().get(),
options.layoutY_transformed(),
options.layoutF_transformed(),
options.layoutD_transformed(),
// options.layoutZ_transformed()
},
hw_info
};
SsdOperation op;
size_t workspace_size = SsdOperation::get_workspace_size(arguments);
cutlass::device_memory::allocation<uint8_t> workspace(workspace_size);
cutlass::Status status = op.can_implement(arguments);
if (status != cutlass::Status::kSuccess) {
std::cerr << "This kernel is not supported. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
status = op.initialize(arguments, workspace.get());
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to initialize the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
cudaError_t result;
// may be used uninitialized
cudaEvent_t start;
cudaEvent_t end;
cudaEventCreate(&start);
cudaEventCreate(&end);
// warm up
if (options.measure) {
for (int i = 0; i < options.warmups; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
}
result = cudaEventRecord(start);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
// Run
for (int i = 0; i < options.iterations; i++) {
status = op.run();
if (status != cutlass::Status::kSuccess) {
std::cerr << "Failed to launch the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(cudaGetLastError()) << std::endl;
return false;
}
}
result = cudaEventRecord(end);
if (result != cudaSuccess) {
std::cerr << "cudaEventRecord() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
result = cudaDeviceSynchronize();
if (result != cudaSuccess) {
std::cerr << "Error running the CUTLASS kernel. Last CUDA error is: "
<< cudaGetErrorString(result) << std::endl;
return false;
}
float runtime_ms = 0;
result = cudaEventElapsedTime(&runtime_ms, start, end);
if (result != cudaSuccess) {
std::cerr << "cudaEventElapsed() failed: " << cudaGetErrorString(result) << std::endl;
return false;
}
runtime_ms /= static_cast<float>(options.iterations);
if (options.verbose) {
printf("[iters = %d, warmups = %d] ssd kernel runtime_ms = %.4f\n", options.iterations, options.warmups, runtime_ms);
printf("smem size = %d\n", SsdOperation::Kernel::SharedStorageSize);
}
// Matrix
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
auto mY_ref_0 = cute::make_tensor(tensor_Y_ref_0.data().get(), options.layoutY());
auto mY_ref_1 = cute::make_tensor(tensor_Y_ref_1.data().get(), options.layoutY());
auto mY_res = cute::make_tensor(tensor_Y.data().get(), options.layoutY());
auto mF_ref_0 = cute::make_tensor(tensor_F_ref_0.data().get(), options.layoutF());
auto mF_ref_1 = cute::make_tensor(tensor_F_ref_1.data().get(), options.layoutF());
auto mF_res = cute::make_tensor(tensor_F.data().get(), options.layoutF());
auto mX = cute::make_tensor(tensor_X.data().get(), options.layoutX());
auto mB = cute::make_tensor(tensor_B.data().get(), options.layoutB());
auto mC = cute::make_tensor(tensor_C.data().get(), options.layoutC());
auto mD = cute::make_tensor(tensor_D.data().get(), options.layoutD());
auto mZ = cute::make_tensor(tensor_Z.data().get(), options.layoutZ());
auto mDelta = cute::make_tensor(tensor_Delta.data().get(), options.layoutDelta());
auto mDeltaA = cute::make_tensor(tensor_DeltaA.data().get(), options.layoutDeltaA());
// Reference Device kernel
if (options.verify) {
ssd_reference<Option::HAS_D, Option::D_HAS_HDIM, Option::HAS_Z>(
mY_ref_1,
mF_ref_1,
mX,
mDelta,
mDeltaA,
mB,
mC,
mD,
mZ,
options
);
}
bool passed = true;
if (options.verify) {
printf("[TensorY]verifying...\n");
passed &= compare_reference<5>(mY_ref_1, mY_res);
printf("[TensorF]verifying...\n");
passed &= compare_reference<4>(mF_ref_1, mF_res);
}
return passed;
}
template<
int TensorDim,
class Engine, class Layout
>
static constexpr bool
compare_reference(
cute::Tensor<Engine, Layout> const& reference,
cute::Tensor<Engine, Layout> const& computed,
float epsilon = 0.05f) {
if (size(reference) != size(computed)) {
return false;
}
bool passed = true;
if (epsilon == 0.0f) {
// fast refcheck w/o epsilon
for (size_t i = 0; i < size_t(size(reference)); ++i) {
if (reference(i) != computed(i)) {
passed = false;
printf("[%llu] %f, %f\n", static_cast<unsigned long long>(i),
float(reference(i)), float(computed(i)));
break;
}
}
}
else {
// refcheck with epsilon
for (size_t i = 0; i < size_t(size(reference)); ++i) {
auto ref = static_cast<float>(reference(i));
auto act = static_cast<float>(computed(i));
auto abs_error = std::abs(act - ref);
auto rel_error = abs_error / (std::max(std::abs(act), std::abs(ref)) + 0.00001f);
if (std::isnan(abs_error) || std::isnan(rel_error) ||
std::min(rel_error, abs_error) > epsilon) {
passed = false;
printf("[%llu] %f, %f\n", static_cast<unsigned long long>(i),
float(reference(i)), float(computed(i)));
break;
}
}
}
if (not passed) {
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
auto m = cute::shape<2>(reference);
auto n = cute::shape<TensorDim-1>(reference);
printf("reference:\n");
for (int mi = 0; mi < m; ++mi) {
for (int ni = 0; ni < n; ++ni) {
if constexpr (TensorDim == 5) {
printf("%.4f ", static_cast<float>(reference(0,0,mi,0,ni)));
}
else {
printf("%.4f ", static_cast<float>(reference(0,0,mi,ni)));
}
}
printf("\n");
}
printf("\n");
printf("computed:\n");
for (int mi = 0; mi < m; ++mi) {
for (int ni = 0; ni < n; ++ni) {
if constexpr (TensorDim == 5) {
printf("%.4f ", static_cast<float>(computed(0,0,mi,0,ni)));
}
else {
printf("%.4f ", static_cast<float>(computed(0,0,mi,ni)));
}
}
printf("\n");
}
printf("\n");
}
return passed;
}
};
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
int main(int argc, char const **args) {
cudaDeviceProp props;
cudaError_t error = cudaGetDeviceProperties(&props, 0);
if (error != cudaSuccess) {
std::cerr << "cudaGetDeviceProperties() returned an error: " << cudaGetErrorString(error) << std::endl;
return -1;
}
if (__CUDACC_VER_MAJOR__ < 12 || props.major < 10) {
std::cout
<< "This example requires a GPU of NVIDIA's Blackwell Architecture or "
<< "later (compute capability 100 or greater) and CUDA 12.0 or greater.\n";
return 0;
}
else if (__CUDACC_VER_MAJOR__ < 12 || (props.major != 10 || props.minor != 0)) {
std::cout
<< "This example requires a GPU of NVIDIA's Blackwell Architecture "
<< "(compute capability 100) and CUDA 12.0 or greater.\n";
return 0;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
//
// Parse options
//
Options options;
options.parse(argc, args);
if (options.help) {
options.print_usage(std::cout) << std::endl;
return 0;
}
if (options.error) {
std::cerr << "Aborting execution." << std::endl;
return -1;
}
// Execute kernel
printf("start testing....\n");
// The KernelHardwareInfo struct holds the number of SMs on the GPU with a given device ID. This
// information is used by the underlying kernel.
cutlass::KernelHardwareInfo hw_info;
// Change device_id to another value if you are running on a machine with multiple GPUs and wish
// to use a GPU other than that with device ID 0.
hw_info.device_id = 0;
hw_info.sm_count = cutlass::KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
// Check Device/Host ref kernel
TestBed<Options> testbed{};
bool passed = testbed.run(options, hw_info);
if (passed) {
printf("everything is ok.\n");
}
else {
printf("something is wrong!!!!!\n");
}
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
return 0;
}
examples/112_blackwell_ssd/CMakeLists.txt
+45
View File
@@ -0,0 +1,45 @@
# 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.
if (CUTLASS_NVCC_ARCHS MATCHES 100a)
set_property(
SOURCE 112_blackwell_ssd.cu
PROPERTY COMPILE_FLAGS "--use_fast_math"
)
cutlass_example_add_executable(
112_blackwell_ssd
112_blackwell_ssd.cu
)
if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")))
endif()
endif()
examples/112_blackwell_ssd/README.md
+65
View File
@@ -0,0 +1,65 @@
# NVIDIA Blackwell SSD (State Space Decomposition) CUDA Example
## Overview
This example demonstrates the implementation of State Space Decomposition (SSD) operations on NVIDIA's Blackwell GPU architecture. It showcases the use of CUTLASS library components for high-performance tensor computations that efficiently leverage Blackwell's advanced hardware capabilities.
## System Requirements
+ NVIDIA GPU with Blackwell Architecture (compute capability 10.0)
+ CUDA Toolkit 12.8 or newer
+ C++17 compatible compiler
## Build the example
Follow the cutlass example building.
## Command Line Options
The example supports the following command line options:
--help: Display the usage statement
--iterations=<int>: Number of iterations for benchmarking (default: 1)
--without_verify: Skip result verification
--verbose: Print execution time per kernel
--G=<int>: Group size (default: 2)
--B=<int>: Batch size (default: 3)
--E=<int>: Expanded factor (default: 2)
--H=<int>: Number of heads (default: 2)
## Limitation
+ Only support LxDxN = 128x64x128
+ Require all TMEM at once and no more cta on the same SM
## Performance
+ Limited by the SEGSUM part.
+ Low MMA utils
+ ALU bound.
# Copyright
Copyright (c) 2024 - 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.
```
examples/112_blackwell_ssd/collective/sm100_ssd_epilogue.hpp
+535
View File
@@ -0,0 +1,535 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/collective/collective_builder.hpp"
namespace cutlass::ssd::collective {
using namespace cute;
template<
class ElementAcc_,
class Element_,
class ElementDA_,
class TileShape_,
class EpilogueTile_,
class SmemLayoutY_,
class SmemLayoutP_,
class SmemLayoutX_,
int StagesInput,
int StagesOutput,
bool HasScaleD_,
bool HasBlockScaleD_>
struct SsdEpilogue {
using TileShape = TileShape_;
using ElementAcc = ElementAcc_;
using Element = Element_;
using ElementD = Element_;
using ElementDA = ElementDA_;
using ElementY = Element_;
using ElementP = Element_;
using EpilogueTile = EpilogueTile_;
using SmemLayoutY = SmemLayoutY_;
using SmemLayoutP = SmemLayoutP_;
using SmemLayoutX = SmemLayoutX_;
static constexpr bool HasScaleD = HasScaleD_;
static constexpr bool HasBlockScaleD = HasBlockScaleD_;
// avoid "warning #3357-D: capturing structured bindings is a C++20 feature"
static constexpr auto L = get<0>(TileShape{});
static constexpr auto D = get<1>(TileShape{});
static constexpr auto N = get<2>(TileShape{});
constexpr static int ThreadCount = 128;
constexpr static size_t SmemAlignmentY = cutlass::detail::alignment_for_swizzle(SmemLayoutY{});
struct CollectiveStorage {
alignas(SmemAlignmentY) ArrayEngine<ElementY , cosize_v<SmemLayoutY>> smem_y;
};
// TMA pipeline for storing D
using StorePipeline = cutlass::PipelineTmaStore<StagesOutput>;
using StorePipelineState = cutlass::PipelineState<StagesOutput>;
using StorePPipeline = cutlass::PipelineTmaStore<1>;
using StorePPipelineState = cutlass::PipelineState<1>;
using CooperatePipeline = cutlass::PipelineAsync<StagesInput>;
using CooperatePipelineState = cutlass::PipelineState<StagesInput>;
using EpiloadPipelineD = cutlass::PipelineTmaAsync<StagesInput>;
static constexpr int kEpiloadDBytes = HasBlockScaleD ? D * sizeof(ElementD) : 0;
struct SharedStorage {
using TensorStorage = CollectiveStorage;
TensorStorage tensors;
};
using TensorStorage = typename SharedStorage::TensorStorage;
using LayoutY = decltype(make_layout(make_shape(L, D, int32_t(0), int32_t(0)),
make_stride(_1{}, int32_t(0), L, int32_t(0)))); // (L,D,C,B)
using LayoutP = decltype(make_layout(make_shape(D, N, int32_t(0)), make_stride(N, _1{}, D*N))); // (D,N,B)
using LayoutD_2D = decltype(make_layout(make_shape(D, int32_t(0)), make_stride(_1{}, D))); // (D,EH)
using LayoutD_1D = decltype(make_layout(make_shape(_1{}, int32_t(0)), make_stride(_0{}, _1{}))); // (D,EH)
using LayoutD = cute::conditional_t<
HasBlockScaleD,
LayoutD_2D,
LayoutD_1D
>;
struct Arguments {
ElementY* ptr_Y{nullptr};
ElementP* ptr_P{nullptr};
const ElementD* ptr_D{nullptr};
LayoutY layout_Y{};
LayoutP layout_P{};
LayoutD layout_D{};
};
using StrideY = cute::tuple<_1, int, int, int>; // (L,D,C,B)
using StrideP = cute::tuple<int, _1, int>; // (D,N,B)
using CopyOpS2G = SM90_TMA_STORE;
struct Params {
using TMA_Y = decltype(make_tma_copy(
CopyOpS2G{},
make_tensor(make_gmem_ptr(static_cast<ElementY const*>(nullptr)),
repeat_like(StrideY{}, int32_t(0)), StrideY{}),
take<0,2>(SmemLayoutY{}),
EpilogueTile{},
_1{}));
using TMA_P = decltype(make_tma_copy(
CopyOpS2G{},
make_tensor(make_gmem_ptr(static_cast<ElementP const*>(nullptr)),
repeat_like(StrideP{}, int32_t(0)), StrideP{}),
take<0,3>(SmemLayoutP{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})), // (D,N)
_1{}));
using TensorD = decltype(make_tensor(
make_gmem_ptr(static_cast<ElementD const*>(nullptr)),
LayoutD{}));
TMA_Y tma_store_y;
TMA_P tma_store_p;
TensorD tensor_d;
};
template<class ProblemShape>
static Params to_underlying_arguments(ProblemShape const& problem_size, Arguments const& args, void* workspace = nullptr) {
using X = Underscore;
auto [G, B, EH, C, L, D, N] = problem_size;
auto tensor_y = make_tensor(make_gmem_ptr(args.ptr_Y), args.layout_Y);
auto tensor_p = make_tensor(make_gmem_ptr(args.ptr_P), args.layout_P);
auto tensor_d = make_tensor(make_gmem_ptr(args.ptr_D), args.layout_D);
auto tma_store_y = make_tma_copy_C_sm90(
CopyOpS2G{},
tensor_y,
take<0,2>(SmemLayoutY{}),
EpilogueTile{});
auto tma_store_p = make_tma_copy_C_sm90(
CopyOpS2G{},
tensor_p,
take<0,2>(SmemLayoutP{}),
make_shape(shape<1>(TileShape{}), shape<2>(TileShape{})));
return Params{
tma_store_y,
tma_store_p,
tensor_d
};
}
template<
class Params, class ProblemShape,
class EpiloadPipeline, class PipelineState,
class TensorStorage
>
CUTLASS_DEVICE
void load_d(
int const& blk_coord, Params const& params, ProblemShape const& problem_size,
EpiloadPipeline& pipeline, PipelineState& pipeline_d_producer_state,
TensorStorage& shared_tensors) {
if constexpr (HasBlockScaleD) {
int lane_predicate = cute::elect_one_sync();
if (lane_predicate) {
auto& gD = params.tensor_d;
ElementD* ptr_d = shared_tensors.smem_d.data();
auto smem_layout = make_layout(make_shape(get<1>(TileShape{}), Int<StagesInput>{})); // (D,)
auto sD = cute::as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(ptr_d), smem_layout));
auto bulk_atom = Copy_Atom<SM90_BULK_COPY_AUTO, ElementD>{};
int write_stage = pipeline_d_producer_state.index();
// LOCK pipeline_state for _writing_
pipeline.producer_acquire(pipeline_d_producer_state);
using BarrierType = typename EpiloadPipeline::ProducerBarrierType;
BarrierType* tma_barrier = pipeline.producer_get_barrier(pipeline_d_producer_state);
copy(bulk_atom.with(*tma_barrier), gD(_,blk_coord), sD(_,write_stage));
// Advance pipeline_state
++pipeline_d_producer_state;
}
}
}
template<
class Params, class ProblemShape,
class MainloopPipelineIntra, class PipelineStateIntra,
class MainloopPipelineAcc, class PipelineStateAcc,
class MainloopPipelineDelta, class PipelineStateDelta,
class MainloopPipelineX, class PipelineStateX,
class EpiloadPipelineD, class PipelineStateD,
class FragmentC_Intra_1, class FragmentC_Intra_2,
class FragmentC_Inter_1, class FragmentC_Inter_2,
class MainloopStorage, class EpilogueStorage
>
CUTLASS_DEVICE
auto store(
int& chunk, int const& blk_coord, int const& blk_coord_eh,
Params const& params, ProblemShape const& problem_size,
MainloopPipelineIntra& pipeline_intra, PipelineStateIntra& pipeline_intra_consumer_state,
MainloopPipelineAcc& pipeline_acc, PipelineStateAcc& pipeline_acc_consumer_state,
MainloopPipelineDelta& pipeline_delta, PipelineStateDelta& pipeline_delta_consumer_state,
MainloopPipelineX& pipeline_x, PipelineStateX& pipeline_x_consumer_state,
EpiloadPipelineD& pipeline_d, PipelineStateD& pipeline_d_producer_state,
StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state,
cute::tuple<FragmentC_Intra_1, FragmentC_Intra_2>& acc_intra,
cute::tuple<FragmentC_Inter_1, FragmentC_Inter_2>& acc_inter,
MainloopStorage& mainloop_tensors, EpilogueStorage& epilogue_tensors,
bool is_first_iteration) {
using CopyOpT2R = SM100_TMEM_LOAD_16dp256b4x;
using CopyOpR2S = SM90_U16x8_STSM_T;
using CopyOpS2R = SM75_U16x8_LDSM_T;
// Epilogue
int thread_idx = int(threadIdx.x % 128);
int warp_idx = thread_idx / 32;
auto [G, B, EH, C, L, D, N] = problem_size;
Tensor mY_mn = params.tma_store_y.get_tma_tensor(make_shape(L,D,C,B*EH));
Tensor gY_mn = local_tile(mY_mn, take<0,2>(TileShape{}), make_coord(_,_,_))(_,_,_0{},_0{},chunk,blk_coord);
auto ptr_sY = epilogue_tensors.smem_y.begin();
Tensor sY_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sY), SmemLayoutY{}));
auto ptr_sX = mainloop_tensors.smem_x.data();
Tensor sX_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sX), SmemLayoutX{}));
auto ts0 = size<0>(TileShape{});
auto ts1 = size<1>(TileShape{});
Layout col_layout = make_layout(make_shape ( ts0, ts1, Int<StagesInput>{}),
make_stride(_1{}, _0{}, ts0));
Tensor sDeltaA = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(mainloop_tensors.smem_delta_a.data()), col_layout));
auto tDeltaA = sDeltaA(_,_,pipeline_delta_consumer_state.index());
Layout row_layout = make_layout(make_shape ( ts0, ts1, Int<StagesInput>{}),
make_stride(_0{}, _1{}, ts1));
Tensor sD = as_position_independent_swizzle_tensor(make_tensor(make_smem_ptr(mainloop_tensors.smem_d.data()), row_layout));
auto tD = sD(_,_,pipeline_d_producer_state.index());
auto accumulator_intra_2 = get<1>(acc_intra);
auto accumulator_inter_2 = get<1>(acc_inter);
auto tIntra = accumulator_intra_2(make_coord(_,_),_0{},_0{},_0{});
auto tInter = accumulator_inter_2(make_coord(_,_),_0{},_0{},_0{});
Tensor tIntra_epi = flat_divide(tIntra, EpilogueTile{});
Tensor tInter_epi = flat_divide(tInter, EpilogueTile{});
Tensor tDeltaA_epi = flat_divide(tDeltaA, EpilogueTile{});
Tensor tD_epi = flat_divide(tD, EpilogueTile{});
Tensor gY_epi = flat_divide(gY_mn, EpilogueTile{});
TiledCopy tiled_t2r = make_tmem_copy(CopyOpT2R{}, tIntra_epi(_,_,_0{},_0{}));
ThrCopy thread_t2r = tiled_t2r.get_slice(thread_idx);
Tensor tTR_tIntra = thread_t2r.partition_S(tIntra_epi);
Tensor tTR_tInter = thread_t2r.partition_S(tInter_epi);
Tensor tTR_tDeltaA = thread_t2r.partition_D(tDeltaA_epi);
Tensor tTR_tD = thread_t2r.partition_D(tD_epi);
Tensor tTR_sY = thread_t2r.partition_D(sY_epi(_,_,_0{}));
Tensor tTR_rIntra = make_tensor<ElementAcc>(shape(thread_t2r.partition_D(gY_epi)));
// Tensor tTR_rInter = make_tensor<ElementAcc>(shape(thread_t2r.partition_D(gY_epi)));
// Tensor tTR_rIntra = make_tensor<ElementAcc>(shape(tTR_sY));
Tensor tTR_rInter = make_tensor<ElementAcc>(shape(tTR_sY));
Tensor tTR_rY = make_tensor<Element>(shape(tTR_sY));
Tensor tTR_rDeltaA = make_tensor<ElementDA>(shape(tTR_sY));
Tensor tTR_rD = make_tensor<Element>(shape(tTR_sY));
Tensor tTR_rD_Acc = make_tensor<ElementAcc>(shape(tTR_sY));
Tensor tTR_rX = make_tensor<Element>(shape(tTR_sY));
Tensor tTR_rX_Acc = make_tensor<ElementAcc>(shape(tTR_sY));
// (t)hread-partition for (r)egister to (s)mem copy (tRS_)
TiledCopy tiled_r2s = make_tiled_copy_D(Copy_Atom<CopyOpR2S, Element>{}, tiled_t2r);
ThrCopy thread_r2s = tiled_r2s.get_slice(thread_idx);
Tensor tRS_sY = thread_r2s.partition_D(sY_epi);
Tensor tRS_rY = thread_r2s.retile_S(tTR_rY);
Tensor tRS_rDeltaA = thread_r2s.retile_S(tTR_rDeltaA);
Tensor tRS_rD = thread_r2s.retile_S(tTR_rD);
Tensor tRS_rCompute = make_tensor<ElementAcc>(shape(tRS_rY));
TiledCopy tiled_s2r = make_tiled_copy_D(Copy_Atom<CopyOpS2R, Element>{}, tiled_t2r);
ThrCopy thread_s2r = tiled_s2r.get_slice(thread_idx);
Tensor tSR_sX = thread_s2r.partition_S(flat_divide(sX_epi, EpilogueTile{}));
Tensor tSR_rX = thread_s2r.retile_D(tTR_rX);
constexpr int FragmentSize = 4;
// Tensor tTR_rIntra_frg = recast<Array<ElementAcc, FragmentSize>>(coalesce(tTR_rIntra));
// Tensor tTR_rInter_frg = recast<Array<ElementAcc, FragmentSize>>(coalesce(tTR_rInter));
Tensor tTR_rY_frg = recast<Array<Element , FragmentSize>>(coalesce(tTR_rY));
Tensor tRS_rY_frg = recast<Array<Element , FragmentSize>>(coalesce(tRS_rY));
Tensor tRS_rDeltaA_frg = recast<Array<ElementDA , FragmentSize>>(coalesce(tRS_rDeltaA));
Tensor tRS_rD_frg = recast<Array<Element , FragmentSize>>(coalesce(tRS_rD));
Tensor tRS_rCompute_frg = recast<Array<ElementAcc, FragmentSize>>(coalesce(tRS_rCompute));
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
ThrCopy thrblk_s2g = params.tma_store_y.get_slice(Int<0>{});
Tensor bSG_sY = thrblk_s2g.partition_S(sY_epi);
Tensor bSG_gY = thrblk_s2g.partition_D(gY_epi);
auto synchronize = [] () { cutlass::arch::NamedBarrier::sync(ThreadCount, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
[[maybe_unused]] bool issue_smem_store = true;
[[maybe_unused]] bool issue_tma_store = warp_idx == 0;
// The TMA store sequence for one subtile iteration
auto tma_store_fn = [&] (int epi_m, int epi_n) {
// Write the tile from smem to gmem with TMA
cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
synchronize(); // ensure all threads have issued their async fence
if (issue_tma_store) {
copy(params.tma_store_y, bSG_sY(_,_,_,store_pipe_producer_state.index()), bSG_gY(_,_,_,epi_m,epi_n));
}
// Commit the TMA stores for this stage
if (issue_tma_store) {
store_pipeline.producer_commit(store_pipe_producer_state);
}
++store_pipe_producer_state;
// Wait for the next smem buffer to be available
if (issue_tma_store) {
store_pipeline.producer_acquire(store_pipe_producer_state);
}
synchronize();
};
// Require Acc
pipeline_intra.consumer_wait(pipeline_intra_consumer_state);
copy(tiled_t2r, tTR_tIntra, tTR_rIntra);
cutlass::arch::fence_view_async_tmem_load();
pipeline_intra.consumer_release(pipeline_intra_consumer_state);
++pipeline_intra_consumer_state;
// Require Inter Acc
pipeline_acc.consumer_wait(pipeline_acc_consumer_state);
CUTLASS_PRAGMA_UNROLL
for (int iter_n = 0; iter_n < size<3>(gY_epi); ++iter_n) {
CUTLASS_PRAGMA_UNROLL
for (int iter_m = 0; iter_m < size<2>(gY_epi); ++iter_m) {
int epi_m = iter_m, epi_n = iter_n;
Tensor tTR_rIntra_mn = tTR_rIntra(_,_,_,epi_m,epi_n);
Tensor tTR_tInter_mn = tTR_tInter(_,_,_,epi_m,epi_n);
Tensor tTR_tDeltaA_mn = tTR_tDeltaA(_,_,_,epi_m,epi_n);
Tensor tTR_tD_mn = tTR_tD(_,_,_,epi_m,epi_n);
copy(tiled_t2r, tTR_tInter_mn, tTR_rInter);
cutlass::arch::fence_view_async_tmem_load();
copy(tTR_tDeltaA_mn, tTR_rDeltaA);
if constexpr (HasScaleD) {
copy(tiled_s2r, tSR_sX(_,_,_,epi_m,epi_n,pipeline_x_consumer_state.index()), tSR_rX);
type_convert<Element, ElementAcc>(tTR_rX, tTR_rX_Acc);
}
if constexpr (HasBlockScaleD) {
if (is_first_iteration) {
pipeline_d.consumer_wait(pipeline_d_producer_state);
}
copy(tTR_tD_mn, tTR_rD);
type_convert<ElementD, ElementAcc>(tTR_rD, tTR_rD_Acc);
}
else if constexpr (HasScaleD) {
auto& gD = params.tensor_d;
auto value = static_cast<ElementAcc>(gD(_0{}, blk_coord_eh));
CUTLASS_PRAGMA_UNROLL
for (int ii = 0; ii < size(tTR_rD_Acc); ++ii) {
tTR_rD_Acc(ii) = value;
}
}
// ?
synchronize();
NumericArrayConverter<Element, ElementAcc, FragmentSize> converter;
CUTLASS_PRAGMA_UNROLL
for (int ii = 0; ii < size(tRS_rCompute); ++ii) {
tRS_rCompute(ii) = tTR_rIntra_mn(ii) + tTR_rInter(ii) * expf(tTR_rDeltaA(ii));
if constexpr (HasScaleD) {
tRS_rCompute(ii) += tTR_rD_Acc(ii) * tTR_rX_Acc(ii);
}
}
CUTLASS_PRAGMA_UNROLL
for (int epi_v = 0; epi_v < size(tTR_rY_frg); ++epi_v) {
tRS_rY_frg(epi_v) = converter(tRS_rCompute_frg(epi_v));
}
copy(tiled_r2s, tRS_rY, tRS_sY(_,_,_,store_pipe_producer_state.index()));
tma_store_fn(epi_m, epi_n);
}
}
pipeline_acc.consumer_release(pipeline_acc_consumer_state);
++pipeline_acc_consumer_state;
pipeline_x.consumer_release_from_threads(pipeline_x_consumer_state);
++pipeline_x_consumer_state;
pipeline_delta.consumer_release_from_threads(pipeline_delta_consumer_state);
++pipeline_delta_consumer_state;
}
template<
class Params, class ProblemShape,
class StorePipeline, class StorePipelineState,
class TensorStorage
>
CUTLASS_DEVICE
auto store_p(
int const& blk_coord, Params const& params, ProblemShape const& problem_size,
StorePipeline& store_pipeline, StorePipelineState& store_pipe_producer_state,
TensorStorage& shared_tensors) {
int thread_idx = int(threadIdx.x % 128);
auto [G, B, EH, C, L, D, N] = problem_size;
Tensor mP_mn = params.tma_store_p.get_tma_tensor(make_shape(D,N,B*EH));
Tensor gP_mn = local_tile(mP_mn, take<1,3>(TileShape{}), make_coord(_,_,blk_coord));
auto gP_epi = gP_mn(_,_,_0{},_0{});
// Construct the corresponding pipelined smem tensors
auto ptr_sP = shared_tensors.smem_p.begin();
Tensor sP_epi = cute::as_position_independent_swizzle_tensor(
make_tensor(make_smem_ptr(ptr_sP), SmemLayoutP{}))(_,_,_0{}); // (EPI_TILE_M,EPI_TILE_N)
// thread(b)lock-partition for (s)mem to (g)mem copy (bSG_)
auto oprands = tma_partition(params.tma_store_p, Int<0>{}, Layout<_1>{},
group_modes<0,2>(sP_epi), group_modes<0,2>(gP_epi)); // (TMA,k) and (TMA,PIPE)
// avoid "warning #3357-D: capturing structured bindings is a C++20 feature"
auto bSG_gY = get<0>(oprands);
auto bSG_sY = get<1>(oprands);
#if 0
if (threadIdx.x % 128 == 0 && (blockIdx.x + blockIdx.y + blockIdx.z == 0)) {
print("mP_mn : ");print(mP_mn);print("\n");
print("gP_mn : ");print(gP_mn);print("\n");
print("sP_epi : ");print(sP_epi);print("\n");
print("gP_epi : ");print(gP_epi);print("\n");
print("bSG_sY : ");print(bSG_sY);print("\n");
print("bSG_gY : ");print(bSG_gY);print("\n");
}
#endif
// Thread synchronizer for previously issued waits or fences
// to ensure visibility of smem reads/writes to threads or TMA unit
// Use the reserved named barrier `streamkbarrier` since this kernel doesn't support streamk
auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(128, cutlass::arch::ReservedNamedBarriers::StreamkBarrier0); };
// Predication for TMA store (one warp issues TMA store)
bool issue_tma_store = (thread_idx / NumThreadsPerWarp) == 0;
auto tma_store_fn = [&] () {
// Write the tile from smem to gmem with TMA
cutlass::arch::fence_view_async_shared(); // ensure smem writes are visible to TMA
synchronize(); // ensure all threads have issued their async fence
if (issue_tma_store) {
copy(params.tma_store_p, bSG_sY, bSG_gY);
}
// Commit the TMA stores for this stage
if (issue_tma_store) {
store_pipeline.producer_commit(store_pipe_producer_state);
}
++store_pipe_producer_state;
// Wait for the next smem buffer to be available
if (issue_tma_store) {
store_pipeline.producer_acquire(store_pipe_producer_state);
}
// don't need pipeline
synchronize();
};
tma_store_fn();
}
template<
class ElementSrc, class ElementDst,
class TensorSrc, class TensorDst
>
CUTLASS_DEVICE
auto type_convert(
TensorSrc& tS,
TensorDst& tD) {
static constexpr int FragmentSize = 2;
NumericArrayConverter<ElementDst, ElementSrc, FragmentSize> converter;
auto tS_frg = recast<Array<ElementSrc, FragmentSize>>(tS);
auto tD_frg = recast<Array<ElementDst, FragmentSize>>(tD);
CUTLASS_PRAGMA_UNROLL
for (int ii = 0; ii < size(tS_frg); ++ii) {
tD_frg(ii) = converter(tS_frg(ii));
}
}
};
} // namespace cutlass::fmha::collective
examples/112_blackwell_ssd/collective/sm100_ssd_gemm_tma_warpspecialized.hpp
+1145
View File
File diff suppressed because it is too large Load Diff
examples/112_blackwell_ssd/device/ssd.hpp
+272
View File
@@ -0,0 +1,272 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
// common
#include "cutlass/cutlass.h"
#include "cutlass/device_kernel.h"
#if !defined(__CUDACC_RTC__)
#include "cutlass/cluster_launch.hpp"
#include "cutlass/trace.h"
#endif // !defined(__CUDACC_RTC__)
////////////////////////////////////////////////////////////////////////////////
namespace cutlass::ssd::device {
////////////////////////////////////////////////////////////////////////////////
////////////////////////////// CUTLASS 3.x API /////////////////////////////////
////////////////////////////////////////////////////////////////////////////////
template <class Kernel_>
class SSD {
public:
using Kernel = Kernel_;
static int const kThreadCount = Kernel::MaxThreadsPerBlock;
/// Argument structure: User API
using Arguments = typename Kernel::Arguments;
/// Argument structure: Kernel API
using Params = typename Kernel::Params;
private:
/// Kernel API parameters object
Params params_;
bool is_initialized(bool set = false) {
static bool initialized = false;
if (set) initialized = true;
return initialized;
}
public:
/// Access the Params structure
Params const& params() const {
return params_;
}
/// Determines whether the GEMM can execute the given problem.
static Status
can_implement(Arguments const& args) {
if (Kernel::can_implement(args)) {
return Status::kSuccess;
}
else {
return Status::kInvalid;
}
}
/// Gets the workspace size
static size_t
get_workspace_size(Arguments const& args) {
size_t workspace_bytes = 0;
workspace_bytes += Kernel::get_workspace_size(args);
return workspace_bytes;
}
/// Computes the grid shape
static dim3
get_grid_shape(Params const& params) {
return Kernel::get_grid_shape(params);
}
/// Computes the maximum number of active blocks per multiprocessor
static int maximum_active_blocks(int /* smem_capacity */ = -1) {
CUTLASS_TRACE_HOST("Universal::maximum_active_blocks()");
int max_active_blocks = -1;
int smem_size = Kernel::SharedStorageSize;
// first, account for dynamic smem capacity if needed
cudaError_t result;
if (smem_size >= (48 << 10)) {
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
result = cudaFuncSetAttribute(
device_kernel<Kernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(
" cudaFuncSetAttribute() returned error: "
<< cudaGetErrorString(result));
return -1;
}
}
// query occupancy after setting smem size
result = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
&max_active_blocks,
device_kernel<Kernel>,
Kernel::MaxThreadsPerBlock,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(
" cudaOccupancyMaxActiveBlocksPerMultiprocessor() returned error: "
<< cudaGetErrorString(result));
return -1;
}
CUTLASS_TRACE_HOST(" max_active_blocks: " << max_active_blocks);
return max_active_blocks;
}
/// Initializes GEMM state from arguments.
Status
initialize(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("Universal::initialize() - workspace "
<< workspace << ", stream: " << (stream ? "non-null" : "null"));
// Initialize the workspace
Status status = Kernel::initialize_workspace(args, workspace, stream);
if (status != Status::kSuccess) {
return status;
}
// Initialize the Params structure
params_ = Kernel::to_underlying_arguments(args, workspace);
if (is_initialized()) return Status::kSuccess;
// account for dynamic smem capacity if needed
int smem_size = Kernel::SharedStorageSize;
if (smem_size >= (48 << 10)) {
CUTLASS_TRACE_HOST(" Setting smem size to " << smem_size);
cudaError_t result = cudaFuncSetAttribute(
device_kernel<Kernel>,
cudaFuncAttributeMaxDynamicSharedMemorySize,
smem_size);
if (cudaSuccess != result) {
result = cudaGetLastError(); // to clear the error bit
CUTLASS_TRACE_HOST(" cudaFuncSetAttribute() returned error: " << cudaGetErrorString(result));
return Status::kErrorInternal;
}
}
is_initialized(true);
return Status::kSuccess;
}
/// Update API is preserved in 3.0, but does not guarantee a lightweight update of params.
Status
update(Arguments const& args, void* workspace = nullptr) {
CUTLASS_TRACE_HOST("Universal()::update() - workspace: " << workspace);
size_t workspace_bytes = get_workspace_size(args);
if (workspace_bytes > 0 && nullptr == workspace) {
return Status::kErrorWorkspaceNull;
}
params_ = Kernel::to_underlying_arguments(args, workspace);
return Status::kSuccess;
}
/// Primary run() entry point API that is static allowing users to create and manage their own params.
/// Supplied params struct must be construct by calling Kernel::to_underling_arguments()
static Status
run(Params& params, cudaStream_t stream = nullptr) {
CUTLASS_TRACE_HOST("Universal::run()");
dim3 const block = Kernel::get_block_shape();
dim3 const grid = get_grid_shape(params);
// configure smem size and carveout
int smem_size = Kernel::SharedStorageSize;
Status launch_result;
// Use extended launch API only for mainloops that use it
if constexpr(Kernel::ArchTag::kMinComputeCapability >= 90) {
dim3 cluster(cute::size<0>(typename Kernel::ClusterShape{}),
cute::size<1>(typename Kernel::ClusterShape{}),
cute::size<2>(typename Kernel::ClusterShape{}));
void const* kernel = (void const*) device_kernel<Kernel>;
void* kernel_params[] = {&params};
launch_result = ClusterLauncher::launch(grid, cluster, block, smem_size, stream, kernel, kernel_params);
}
else {
launch_result = Status::kSuccess;
device_kernel<Kernel><<<grid, block, smem_size, stream>>>(params);
}
cudaError_t result = cudaGetLastError();
if (cudaSuccess == result && Status::kSuccess == launch_result) {
return Status::kSuccess;
}
else {
CUTLASS_TRACE_HOST(" Kernel launch failed. Reason: " << result);
return Status::kErrorInternal;
}
}
//
// Non-static launch overloads that first create and set the internal params struct of this kernel handle.
//
/// Launches the kernel after first constructing Params internal state from supplied arguments.
Status
run(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
Status status = initialize(args, workspace, stream);
if (Status::kSuccess == status) {
status = run(params_, stream);
}
return status;
}
/// Launches the kernel after first constructing Params internal state from supplied arguments.
Status
operator()(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr) {
return run(args, workspace, stream);
}
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
Status
run(cudaStream_t stream = nullptr) {
return run(params_, stream);
}
/// Overload that allows a user to re-launch the same kernel without updating internal params struct.
Status
operator()(cudaStream_t stream = nullptr) {
return run(params_, stream);
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::device
////////////////////////////////////////////////////////////////////////////////
examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_builder.hpp
+259
View File
@@ -0,0 +1,259 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "../collective/sm100_ssd_epilogue.hpp"
#include "../collective/sm100_ssd_gemm_tma_warpspecialized.hpp"
#include "../kernel/sm100_ssd_kernel_tma_warpspecialized.hpp"
#include "../kernel/sm100_ssd_tile_scheduler.hpp"
#include "cutlass/cutlass.h"
#include "cutlass/epilogue/collective/collective_builder.hpp"
namespace cutlass::ssd::kernel::detail {
template<
class ElementA,
class ElementB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB
>
constexpr auto
sm100_make_ts_tiled_mma() {
return cutlass::gemm::collective::detail::sm100_make_1sm_ts_trivial_tiled_mma<
ElementA, ElementB, ElementAccumulator,
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB>();
}
template<
class ElementA,
class ElementB,
class ElementAccumulator,
class TileShape_MNK,
class ClusterShape_MNK,
UMMA::Major UmmaMajorA,
UMMA::Major UmmaMajorB
>
constexpr auto
sm100_make_ss_tiled_mma() {
return cutlass::gemm::collective::detail::sm100_make_1sm_trivial_tiled_mma<
ElementA, ElementB, ElementAccumulator,
TileShape_MNK, ClusterShape_MNK, UmmaMajorA, UmmaMajorB>();
}
}
namespace cutlass::ssd::kernel {
template<
class Element_,
class ElementDA_,
class ElementAcc_,
class ElementY_,
class TileShape_,
bool HAS_D_,
bool D_HAS_HDIM_
>
struct Sm100SsdBuilder {
using Element = Element_;
using ElementDA = ElementDA_;
using ElementAcc = ElementAcc_;
using ElementY = ElementY_;
using TileShape = TileShape_;
using ArchTag = cutlass::arch::Sm100;
// hard-code
using ClusterShape = Shape<_1,_1,_1>;
static constexpr int StagesInput = 2;
static constexpr int StagesOutput = 2;
using TileShapeIntraBMM1 = decltype(make_shape(get<0>(TileShape{}), get<0>(TileShape{}), get<2>(TileShape{}))); // (L,L,N)
using TileShapeIntraBMM2 = decltype(make_shape(get<0>(TileShape{}), get<1>(TileShape{}), get<0>(TileShape{}))); // (L,D,L)
using TileShapeInterBMM1 = decltype(make_shape(get<2>(TileShape{}), get<1>(TileShape{}), get<0>(TileShape{}))); // (N,D,L)
using TileShapeInterBMM2 = decltype(make_shape(get<0>(TileShape{}), get<1>(TileShape{}), get<2>(TileShape{}))); // (L,D,N)
// LxLxN, NT
using TiledMmaIntra1 = decltype(detail::sm100_make_ss_tiled_mma<Element, Element, ElementAcc,
TileShapeIntraBMM1, ClusterShape,
cute::UMMA::Major::MN, cute::UMMA::Major::MN>());
// LxNxL, TN
using TiledMmaIntra2 = decltype(detail::sm100_make_ts_tiled_mma<Element, Element, ElementAcc,
TileShapeIntraBMM2, ClusterShape,
cute::UMMA::Major::K, cute::UMMA::Major::K>());
// NxDxL, TN
using TiledMmaInter1 = decltype(detail::sm100_make_ts_tiled_mma<Element, Element, ElementAcc,
TileShapeInterBMM1, ClusterShape,
cute::UMMA::Major::K, cute::UMMA::Major::K>());
// LxDxN, NN
using TiledMmaInter2 = decltype(detail::sm100_make_ss_tiled_mma<Element, Element, ElementAcc,
TileShapeInterBMM2, ClusterShape,
cute::UMMA::Major::MN, cute::UMMA::Major::K>());
// ((MMA_TILE_M,MMA_TILE_K), MMA_M, MMA_K)
using MmaShapeC_MK = decltype(partition_shape_A(TiledMmaIntra1{}, make_shape(cute::size<0>(TileShapeIntraBMM1{}),
cute::size<2>(TileShapeIntraBMM1{}))));
using MmaShapeB_NK = decltype(partition_shape_B(TiledMmaIntra1{}, make_shape(cute::size<1>(TileShapeIntraBMM1{}),
cute::size<2>(TileShapeIntraBMM1{}))));
using MmaShapeQ_MK = decltype(partition_shape_A(TiledMmaIntra2{}, make_shape(cute::size<0>(TileShapeIntraBMM2{}),
cute::size<2>(TileShapeIntraBMM2{}))));
using MmaShapeX_NK = decltype(partition_shape_B(TiledMmaIntra2{}, make_shape(cute::size<1>(TileShapeIntraBMM2{}),
cute::size<2>(TileShapeIntraBMM2{}))));
using MmaShapeB_MK = decltype(partition_shape_A(TiledMmaInter1{}, make_shape(cute::size<0>(TileShapeInterBMM1{}),
cute::size<2>(TileShapeInterBMM1{}))));
using MmaShapeP_NK = decltype(partition_shape_B(TiledMmaInter2{}, make_shape(cute::size<1>(TileShapeInterBMM2{}),
cute::size<2>(TileShapeInterBMM2{}))));
using GmemTiledCopyX = cute::SM90_TMA_LOAD;
using GmemTiledCopyB = cute::SM90_TMA_LOAD;
using GmemTiledCopyC = cute::SM90_TMA_LOAD;
using BlockTileX_N = decltype(cute::size<0,0>(MmaShapeX_NK{}) * cute::size<1>(MmaShapeX_NK{}));
using BlockTileX_K = decltype(cute::size<0,1>(MmaShapeX_NK{}) * cute::size<2>(MmaShapeX_NK{}));
using SmemLayoutAtomX = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::K, Element, BlockTileX_N, BlockTileX_K>());
using BlockTileB_N = decltype(cute::size<0,0>(MmaShapeB_NK{}) * cute::size<1>(MmaShapeB_NK{}));
using BlockTileB_K = decltype(cute::size<0,1>(MmaShapeB_NK{}) * cute::size<2>(MmaShapeB_NK{}));
using SmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::MN, Element, BlockTileB_N, BlockTileB_K>());
using BlockTileBT_M = decltype(cute::size<0,0>(MmaShapeB_MK{}) * cute::size<1>(MmaShapeB_MK{}));
using BlockTileBT_K = decltype(cute::size<0,1>(MmaShapeB_MK{}) * cute::size<2>(MmaShapeB_MK{}));
using SmemLayoutAtomBT = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::K, Element, BlockTileBT_M, BlockTileBT_K>());
using TmemLayoutAtomB = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::K, Element, BlockTileBT_M, BlockTileBT_K>());
using BlockTileC_M = decltype(cute::size<0,0>(MmaShapeC_MK{}) * cute::size<1>(MmaShapeC_MK{}));
using BlockTileC_K = decltype(cute::size<0,1>(MmaShapeC_MK{}) * cute::size<2>(MmaShapeC_MK{}));
using SmemLayoutAtomC = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::MN, Element, BlockTileC_M, BlockTileC_K>());
using BlockTileP_N = decltype(cute::size<0,0>(MmaShapeP_NK{}) * cute::size<1>(MmaShapeP_NK{}));
using BlockTileP_K = decltype(cute::size<0,1>(MmaShapeP_NK{}) * cute::size<2>(MmaShapeP_NK{}));
using SmemLayoutAtomP = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::K, Element, BlockTileP_N, BlockTileP_K>());
using SmemLayoutAtomPT = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::MN, Element, BlockTileP_K, BlockTileP_N>());
using BlockTileQ_M = decltype(cute::size<0,0>(MmaShapeQ_MK{}) * cute::size<1>(MmaShapeQ_MK{}));
using BlockTileQ_K = decltype(cute::size<0,1>(MmaShapeQ_MK{}) * cute::size<2>(MmaShapeQ_MK{}));
using SmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::K, Element, BlockTileP_N, BlockTileP_K>());
using TmemLayoutAtomQ = decltype(cutlass::gemm::collective::detail::sm100_smem_selector<
cute::UMMA::Major::K, Element, BlockTileQ_M, BlockTileQ_K>());
using SmemLayoutX = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomX{},
append(MmaShapeX_NK{}, Int<StagesInput>{}),
Step<_2,_1,_3>{}));
using SmemLayoutB = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomB{},
append(MmaShapeB_NK{}, Int<StagesInput>{}),
Step<_2,_1,_3>{}));
// Be consistent with SmemLayoutB
using SmemLayoutBT = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomBT{},
append(MmaShapeB_MK{}, Int<StagesInput>{}),
Step<_1,_2,_3>{}));
using TmemLayoutB = decltype(UMMA::tile_to_mma_shape(
TmemLayoutAtomB{},
append(MmaShapeB_MK{}, Int<1>{}),
Step<_2,_1,_3>{}));
using SmemLayoutC = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomC{},
append(MmaShapeC_MK{}, Int<StagesInput>{}),
Step<_2,_1,_3>{}));
// P only need 1 stage in this case
using SmemLayoutPT = decltype(tile_to_shape(
SmemLayoutAtomPT{},
append(make_shape(get<2>(TileShape{}), get<1>(TileShape{})), Int<1>{})));
using SmemLayoutP = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomP{},
append(MmaShapeP_NK{}, Int<1>{}),
Step<_2,_1,_3>{}));
using TmemLayoutQ = decltype(UMMA::tile_to_mma_shape(
TmemLayoutAtomQ{},
append(MmaShapeQ_MK{}, Int<1>{}),
Step<_2,_1,_3>{}));
using SmemLayoutQ = decltype(UMMA::tile_to_mma_shape(
SmemLayoutAtomQ{},
append(MmaShapeQ_MK{}, Int<2>{}),
Step<_2,_1,_3>{}));
using SmemLayoutAtomXT = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
cute::GMMA::Major::MN, Element, decltype(get<0>(TileShape{})), decltype(get<1>(TileShape{}))>());
using SmemLayoutXT = decltype(tile_to_shape(
SmemLayoutAtomXT{},
make_shape(size<0>(TileShape{}), size<1>(TileShape{}), Int<StagesInput>{}),
Step<_1,_2,_3>{}));
using EpilogueTileType = cutlass::epilogue::collective::EpilogueTileAuto;
using Schedule = cutlass::epilogue::TmaWarpSpecialized;
using EpilogueTile = Shape<Int<128>, Int<32>>;
using SmemLayoutAtomY = decltype(cutlass::gemm::collective::detail::ss_smem_selector<
cute::GMMA::Major::MN, ElementY, decltype(get<0>(EpilogueTile{})), decltype(get<1>(EpilogueTile{}))>());
using SmemLayoutY = decltype(tile_to_shape(
SmemLayoutAtomY{},
make_shape(size<0>(EpilogueTile{}), size<1>(EpilogueTile{}), Int<StagesOutput>{}),
Step<_2,_1,_3>{}));
using SmemLayoutStoreP = decltype(tile_to_shape(
SmemLayoutAtomP{},
append(make_shape(get<1>(TileShape{}), get<2>(TileShape{})), Int<1>{}),
Step<_2,_1,_3>{}));
using CollectiveMainloop = cutlass::ssd::collective::SsdMainloopTmaWarpSpecialized<
Element, ElementDA, ElementAcc, ElementY, TileShape,
StagesInput, StagesOutput,
TiledMmaIntra1, TiledMmaIntra2,
TiledMmaInter1, TiledMmaInter2,
SmemLayoutX, SmemLayoutB, SmemLayoutC, SmemLayoutP,
SmemLayoutBT, SmemLayoutPT, SmemLayoutQ,
TmemLayoutB, TmemLayoutQ>;
using CollectiveEpilogue = cutlass::ssd::collective::SsdEpilogue<
ElementAcc, Element, ElementDA, TileShape,
EpilogueTile, SmemLayoutY, SmemLayoutStoreP, SmemLayoutXT, StagesInput, StagesOutput,
HAS_D_, D_HAS_HDIM_>;
using TileScheduler = cutlass::ssd::kernel::PersistentTileScheduler;
using Kernel = cutlass::ssd::kernel::SsdKernelTmaWarpSpecialized<CollectiveMainloop, CollectiveEpilogue, TileScheduler>;
};
}
examples/112_blackwell_ssd/kernel/sm100_ssd_kernel_tma_warpspecialized.hpp
+519
View File
@@ -0,0 +1,519 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/arch/reg_reconfig.h"
#include "cutlass/pipeline/pipeline.hpp"
#include "cutlass/arch/arch.h"
namespace cutlass::ssd::kernel {
using namespace cute;
template<
class CollectiveMainloop_,
class CollectiveEpilogue_,
class TileScheduler_
>
struct SsdKernelTmaWarpSpecialized {
using CollectiveMainloop = CollectiveMainloop_;
using CollectiveEpilogue = CollectiveEpilogue_;
using TileScheduler = TileScheduler_;
// TileShape: LDN
using TileShape = typename CollectiveMainloop::TileShape;
// Force to use 1x1x1
using ClusterShape = typename CollectiveMainloop::ClusterShape;
using ArchTag = cutlass::arch::Sm100;
static constexpr uint32_t NumMmaThreads = NumThreadsPerWarp; // 1 warp
static constexpr uint32_t NumEpilogueThreads = CollectiveEpilogue::ThreadCount;
static constexpr uint32_t NumEpilogueWarps = NumEpilogueThreads / NumThreadsPerWarp;
// Pipeline for Tensor X
using MainloopPipelineX = typename CollectiveMainloop::MainloopPipelineX;
using PipelineStateX = typename cutlass::PipelineState<MainloopPipelineX::Stages>;
using PipelineParamsX = typename MainloopPipelineX::Params;
// Pipeline for Tensor Delta && DeltaA
using MainloopPipelineDelta = typename CollectiveMainloop::MainloopPipelineDelta;
using PipelineStateDelta = typename cutlass::PipelineState<MainloopPipelineDelta::Stages>;
using PipelineParamsDelta = typename MainloopPipelineDelta::Params;
// Pipeline for Tensor X
using MainloopPipelineB = typename CollectiveMainloop::MainloopPipelineB;
using PipelineStateB = typename cutlass::PipelineState<MainloopPipelineB::Stages>;
using PipelineParamsB = typename MainloopPipelineB::Params;
// Pipeline for Tensor X
using MainloopPipelineC = typename CollectiveMainloop::MainloopPipelineC;
using PipelineStateC = typename cutlass::PipelineState<MainloopPipelineC::Stages>;
using PipelineParamsC = typename MainloopPipelineC::Params;
// Pipeline for Intra Fusion
using MainloopPipelineIntra = typename CollectiveMainloop::MainloopPipelineIntra;
using PipelineStateIntra = typename cutlass::PipelineState<MainloopPipelineIntra::Stages>;
using PipelineParamsIntra = typename MainloopPipelineIntra::Params;
// Pipeline for Inter Fusion
using MainloopPipelineInter = typename CollectiveMainloop::MainloopPipelineInter;
using PipelineStateInter = typename cutlass::PipelineState<MainloopPipelineInter::Stages>;
using PipelineParamsInter = typename MainloopPipelineInter::Params;
// Pipeline for Accumulator
using AccumulatorPipeline = typename CollectiveMainloop::AccumulatorPipeline;
using AccumulatorPipelineState = typename CollectiveMainloop::AccumulatorPipelineState;
using PipelineParamsAcc = typename AccumulatorPipeline::Params;
// Pipeline for Epilgoue store
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
using EpiStorePipelineState = typename CollectiveEpilogue::StorePipelineState;
// Pipeline for Epilgoue load D
using EpiloadPipelineD = typename CollectiveEpilogue::EpiloadPipelineD;
using PipelineParamsD = typename EpiloadPipelineD::Params;
using PipelineStateD = typename cutlass::PipelineState<EpiloadPipelineD::Stages>;
using TmemAllocator = typename cute::TMEM::Allocator1Sm;
struct SharedStorage {
struct PipelineStorage : cute::aligned_struct<16, _1> {
using PipelineStorageX = typename MainloopPipelineX::SharedStorage;
using PipelineStorageDelta = typename MainloopPipelineDelta::SharedStorage;
using PipelineStorageB = typename MainloopPipelineB::SharedStorage;
using PipelineStorageC = typename MainloopPipelineC::SharedStorage;
using PipelineStorageIntra = typename MainloopPipelineIntra::SharedStorage;
using PipelineStorageInter = typename MainloopPipelineInter::SharedStorage;
using PipelineStorageAcc = typename AccumulatorPipeline::SharedStorage;
using PipelineStorageD = typename EpiloadPipelineD::SharedStorage;
alignas(16) PipelineStorageX pipeline_storage_x;
alignas(16) PipelineStorageDelta pipeline_storage_delta;
alignas(16) PipelineStorageB pipeline_storage_b;
alignas(16) PipelineStorageC pipeline_storage_c;
alignas(16) PipelineStorageD pipeline_storage_d;
alignas(16) PipelineStorageIntra pipeline_storage_intra;
alignas(16) PipelineStorageInter pipeline_storage_inter;
alignas(16) PipelineStorageAcc pipeline_storage_acc;
} pipelines;
uint32_t tmem_base_ptr;
struct TensorStorage : cute::aligned_struct<128, _1> {
using EpilogueTensorStorage = typename CollectiveEpilogue::TensorStorage;
using MainloopTensorStorage = typename CollectiveMainloop::TensorStorage;
EpilogueTensorStorage epilogue;
MainloopTensorStorage mainloop;
} tensors;
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
// [G, B, EH, C, L, D, N]
using ProblemShape = cute::tuple<int, int, int, int, int, int, int>;
struct Arguments {
ProblemShape problem_size;
typename CollectiveMainloop::Arguments mainloop;
typename CollectiveEpilogue::Arguments epilogue;
KernelHardwareInfo hw_info;
};
struct Params {
ProblemShape problem_size;
typename CollectiveMainloop::Params mainloop;
typename CollectiveEpilogue::Params epilogue;
typename TileScheduler::Params tile_scheduler;
};
static const int MinBlocksPerMultiprocessor = 1;
static const int MaxThreadsPerBlock = 384;
static size_t get_workspace_size(Arguments const& args) { return 0; }
static cutlass::Status initialize_workspace(Arguments const&, void*, cudaStream_t) {
return cutlass::Status::kSuccess;
}
static bool can_implement(Arguments const& args) {
return CollectiveMainloop::can_implement(args.problem_size, args.mainloop);
}
static dim3 get_grid_shape(Params const& params) {
return TileScheduler::get_grid_shape(params.tile_scheduler);
}
static dim3 get_block_shape() {
dim3 block(MaxThreadsPerBlock, 1, 1);
return block;
}
static Params to_underlying_arguments(Arguments const& args, void* workspace) {
return Params{
args.problem_size,
CollectiveMainloop::to_underlying_arguments(args.problem_size, args.mainloop, workspace),
CollectiveEpilogue::to_underlying_arguments(args.problem_size, args.epilogue, workspace),
TileScheduler::to_underlying_arguments(args.problem_size, args.hw_info, ClusterShape{}, TileShape{})
};
}
CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
using namespace cute;
using X = Underscore;
// TBD
enum class WarpCategory : int32_t {
MMAInter = 0,
MMAIntra = 1,
DMA0 = 2, // Delta, X
DMA1 = 3, // B, C
PreInter = 4,
PreIntra = 8
};
// Parameters
// [G, B, EH, C, L, D, N]
auto C = get<3>(params.problem_size);
// Shared memory.
auto& storage = *reinterpret_cast<SharedStorage*>(smem);
int lane_predicate = cute::elect_one_sync();
int lane_idx = cutlass::canonical_lane_idx();
int warp_idx = cutlass::canonical_warp_idx_sync();
auto warp_category = (WarpCategory(warp_idx) < WarpCategory::PreInter) ?
WarpCategory(warp_idx) :
(WarpCategory(warp_idx) < WarpCategory::PreIntra) ?
WarpCategory::PreInter : WarpCategory::PreIntra;
// Issue Tma Descriptor Prefetch from a single thread
if ((warp_idx == 0) && lane_predicate) {
CollectiveMainloop::prefetch_tma_descriptors(params.mainloop);
}
// Pipeline (TBD)
PipelineParamsX pipeline_params_x;
pipeline_params_x.transaction_bytes = CollectiveMainloop::kXLoadBytes;
pipeline_params_x.is_leader = lane_predicate && (warp_category == WarpCategory::DMA0);
pipeline_params_x.num_consumers = 1 + 1 + cutlass::NumThreadsPerWarpGroup; // MMA0 + MMA1 + PreInter
pipeline_params_x.initializing_warp = 4;
PipelineParamsDelta pipeline_params_delta;
pipeline_params_delta.transaction_bytes = CollectiveMainloop::kDeltaLoadBytes + CollectiveMainloop::kDeltaALoadBytes;
pipeline_params_delta.is_leader = lane_predicate && (warp_category == WarpCategory::DMA0);
pipeline_params_delta.num_consumers = cutlass::NumThreadsPerWarpGroup + cutlass::NumThreadsPerWarpGroup; // PreInter + PreIntra
pipeline_params_delta.initializing_warp = 5;
PipelineParamsB pipeline_params_b;
pipeline_params_b.transaction_bytes = CollectiveMainloop::kBLoadBytes;
pipeline_params_b.is_leader = lane_predicate && (warp_category == WarpCategory::DMA1);
pipeline_params_b.num_consumers = 1 + cutlass::NumThreadsPerWarpGroup;
pipeline_params_b.initializing_warp = 6;
PipelineParamsC pipeline_params_c;
pipeline_params_c.transaction_bytes = CollectiveMainloop::kCLoadBytes;
pipeline_params_c.is_leader = lane_predicate && (warp_category == WarpCategory::DMA1);
pipeline_params_c.num_consumers = 1 + 1;
pipeline_params_c.initializing_warp = 7;
PipelineParamsIntra pipeline_params_intra;
pipeline_params_intra.producer_arv_count = 1;
pipeline_params_intra.consumer_arv_count = cutlass::NumThreadsPerWarpGroup;
pipeline_params_intra.initializing_warp = 8;
PipelineParamsIntra pipeline_params_inter;
pipeline_params_inter.producer_arv_count = 1;
pipeline_params_inter.consumer_arv_count = cutlass::NumThreadsPerWarpGroup;
pipeline_params_inter.initializing_warp = 9;
PipelineParamsAcc pipeline_params_acc;
pipeline_params_acc.producer_arv_count = 1;
pipeline_params_acc.consumer_arv_count = cutlass::NumThreadsPerWarpGroup;
pipeline_params_acc.initializing_warp = 10;
PipelineParamsD pipeline_params_d;
pipeline_params_d.transaction_bytes = CollectiveEpilogue::kEpiloadDBytes;
pipeline_params_d.is_leader = lane_predicate && (warp_category == WarpCategory::DMA0);
pipeline_params_d.num_consumers = cutlass::NumThreadsPerWarpGroup;
pipeline_params_d.initializing_warp = 11;
if (warp_category == WarpCategory::DMA0) {
pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Producer;
pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Producer;
pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Producer;
}
if (warp_category == WarpCategory::DMA1) {
pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Producer;
pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Producer;
}
// TBD
if (warp_category == WarpCategory::MMAInter) {
pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer;
// pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer;
pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Consumer;
pipeline_params_inter.role = MainloopPipelineInter::ThreadCategory::Producer;
pipeline_params_acc.role = AccumulatorPipeline::ThreadCategory::Producer;
}
if (warp_category == WarpCategory::MMAIntra) {
pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer;
pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer;
pipeline_params_c.role = MainloopPipelineC::ThreadCategory::Consumer;
pipeline_params_intra.role = MainloopPipelineIntra::ThreadCategory::Producer;
}
if (warp_category == WarpCategory::PreInter) {
pipeline_params_b.role = MainloopPipelineB::ThreadCategory::Consumer;
pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Consumer;
pipeline_params_inter.role = MainloopPipelineInter::ThreadCategory::Consumer;
}
if (warp_category == WarpCategory::PreIntra) {
pipeline_params_x.role = MainloopPipelineX::ThreadCategory::Consumer;
pipeline_params_delta.role = MainloopPipelineDelta::ThreadCategory::Consumer;
pipeline_params_acc.role = AccumulatorPipeline::ThreadCategory::Consumer;
pipeline_params_intra.role = MainloopPipelineIntra::ThreadCategory::Consumer;
pipeline_params_d.role = EpiloadPipelineD::ThreadCategory::Consumer;
}
MainloopPipelineX pipeline_x(storage.pipelines.pipeline_storage_x, pipeline_params_x, Shape<_1,_1,_1>{});
PipelineStateX mainloop_pipe_x_consumer;
PipelineStateX mainloop_pipe_x_producer = cutlass::make_producer_start_state<MainloopPipelineX>();
MainloopPipelineDelta pipeline_delta(storage.pipelines.pipeline_storage_delta, pipeline_params_delta, Shape<_1,_1,_1>{});
PipelineStateDelta mainloop_pipe_delta_consumer;
PipelineStateDelta mainloop_pipe_delta_producer = cutlass::make_producer_start_state<MainloopPipelineDelta>();
MainloopPipelineB pipeline_b(storage.pipelines.pipeline_storage_b, pipeline_params_b, Shape<_1,_1,_1>{});
PipelineStateB mainloop_pipe_b_consumer;
PipelineStateB mainloop_pipe_b_producer = cutlass::make_producer_start_state<MainloopPipelineB>();
MainloopPipelineC pipeline_c(storage.pipelines.pipeline_storage_c, pipeline_params_c, Shape<_1,_1,_1>{});
PipelineStateC mainloop_pipe_c_consumer;
PipelineStateC mainloop_pipe_c_producer = cutlass::make_producer_start_state<MainloopPipelineC>();
EpiloadPipelineD pipeline_d(storage.pipelines.pipeline_storage_d, pipeline_params_d, Shape<_1,_1,_1>{});
PipelineStateD epi_load_pipe_d_consumer;
PipelineStateD epi_load_pipe_d_producer = cutlass::make_producer_start_state<EpiloadPipelineD>();
MainloopPipelineIntra pipeline_intra(storage.pipelines.pipeline_storage_intra, pipeline_params_intra, Shape<_1,_1,_1>{});
PipelineStateIntra mainloop_pipe_intra_consumer;
PipelineStateIntra mainloop_pipe_intra_producer = cutlass::make_producer_start_state<MainloopPipelineIntra>();
MainloopPipelineInter pipeline_inter(storage.pipelines.pipeline_storage_inter, pipeline_params_inter, Shape<_1,_1,_1>{});
// Opposite pipeline state
PipelineStateInter mainloop_pipe_inter_consumer = cutlass::make_producer_start_state<MainloopPipelineIntra>();
PipelineStateInter mainloop_pipe_inter_producer;
AccumulatorPipeline pipeline_acc(storage.pipelines.pipeline_storage_acc, pipeline_params_acc, Shape<_1,_1,_1>{});
AccumulatorPipelineState mainloop_pipe_acc_consumer;
AccumulatorPipelineState mainloop_pipe_acc_producer = cutlass::make_producer_start_state<AccumulatorPipeline>();
// Epilogue Store pipeline
using EpiStorePipeline = typename CollectiveEpilogue::StorePipeline;
typename EpiStorePipeline::Params epi_store_pipeline_params;
epi_store_pipeline_params.always_wait = true;
EpiStorePipeline epi_store_pipeline(epi_store_pipeline_params);
PipelineState epi_store_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePipeline>();
// Epilogue Store P pipeline
using EpiStorePPipeline = typename CollectiveEpilogue::StorePPipeline;
typename EpiStorePPipeline::Params epi_store_p_pipeline_params;
epi_store_p_pipeline_params.always_wait = true;
EpiStorePPipeline epi_store_p_pipeline(epi_store_p_pipeline_params);
PipelineState epi_store_p_pipe_producer_state = cutlass::make_producer_start_state<EpiStorePPipeline>();
// Tmem allocator
TmemAllocator tmem_allocator{};
// We need this to guarantee that the Pipeline init is visible
// To all producers and consumer blocks in the Cluster
// and to finish smem init
if constexpr (size(ClusterShape{}) > 1) {
cute::cluster_arrive_relaxed();
cute::cluster_wait();
}
else {
__syncthreads();
}
// ignore the tmem alloc
arch::NamedBarrier tmem_allocation_result_barrier(NumMmaThreads + NumMmaThreads + NumEpilogueThreads + NumEpilogueThreads,
cutlass::arch::ReservedNamedBarriers::TmemAllocBarrier);
if (warp_category == WarpCategory::MMAIntra) {
tmem_allocator.allocate(TmemAllocator::Sm100TmemCapacityColumns, &storage.tmem_base_ptr);
__syncwarp();
tmem_allocation_result_barrier.arrive();
}
if (warp_category == WarpCategory::MMAInter || warp_category == WarpCategory::PreIntra || warp_category == WarpCategory::PreInter) {
tmem_allocation_result_barrier.arrive_and_wait();
}
// Kernel implement(TBD)
CollectiveMainloop collective_mainloop;
CollectiveEpilogue collective_epilogue;
TileScheduler tile_scheduler{params.tile_scheduler};
auto mma_output_intra = collective_mainloop.get_mma_intra_acc();
auto mma_output_inter = collective_mainloop.get_mma_inter_acc();
if (warp_category == WarpCategory::DMA0) {
auto load_input = collective_mainloop.load_x_init(params.mainloop, params.problem_size);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
auto blk_coord_eh = tile_scheduler.get_block_coord_eh();
// Load D
collective_epilogue.load_d(
blk_coord_eh, params.epilogue, params.problem_size,
pipeline_d, epi_load_pipe_d_producer,
storage.tensors.mainloop
);
// Load X
// Load Delta
// Load DeltaA
collective_mainloop.load_x_delta(
blk_coord, params.mainloop, params.problem_size,
pipeline_x, mainloop_pipe_x_producer,
pipeline_delta, mainloop_pipe_delta_producer,
load_input,
storage.tensors.mainloop
);
}
}
else if (warp_category == WarpCategory::DMA1) {
auto load_input_b = collective_mainloop.load_b_init(params.mainloop, params.problem_size);
auto load_input_c = collective_mainloop.load_c_init(params.mainloop, params.problem_size);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord_b();
// Load B
// Load C
collective_mainloop.load_b_c(
blk_coord, params.mainloop, params.problem_size,
pipeline_b, mainloop_pipe_b_producer,
pipeline_c, mainloop_pipe_c_producer,
load_input_b, load_input_c,
storage.tensors.mainloop
);
}
}
else if (warp_category == WarpCategory::MMAIntra) {
auto [mma_inputs_1, mma_inputs_2] = collective_mainloop.mma_intra_init(storage.tensors.mainloop);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
for (int chunk = 0; chunk < C; ++chunk) {
collective_mainloop.mma_intra(
pipeline_b, mainloop_pipe_b_consumer,
pipeline_c, mainloop_pipe_c_consumer,
pipeline_x, mainloop_pipe_x_consumer,
pipeline_intra, mainloop_pipe_intra_producer,
mma_inputs_1, mma_inputs_2,
mma_output_intra
);
}
}
}
else if (warp_category == WarpCategory::MMAInter) {
auto [mma_inputs_1, mma_inputs_2] = collective_mainloop.mma_inter_init(storage.tensors.mainloop);
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
for (int chunk = 0; chunk < C; ++chunk) {
collective_mainloop.mma_inter(
pipeline_c, mainloop_pipe_c_consumer,
pipeline_x, mainloop_pipe_x_consumer,
pipeline_inter, mainloop_pipe_inter_producer,
pipeline_acc, mainloop_pipe_acc_producer,
mma_inputs_1, mma_inputs_2,
mma_output_inter
);
}
}
}
else if (warp_category == WarpCategory::PreIntra) {
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto blk_coord = tile_scheduler.get_block_coord();
auto blk_coord_eh = tile_scheduler.get_block_coord_eh();
bool is_first_iteration = true;
for (int chunk = 0; chunk < C; ++chunk) {
collective_mainloop.pre_intra(
pipeline_delta, mainloop_pipe_delta_consumer,
pipeline_intra, mainloop_pipe_intra_consumer,
mma_output_intra,
storage.tensors.mainloop
);
collective_epilogue.store(
chunk, blk_coord, blk_coord_eh, params.epilogue, params.problem_size,
pipeline_intra, mainloop_pipe_intra_consumer,
pipeline_acc, mainloop_pipe_acc_consumer,
pipeline_delta, mainloop_pipe_delta_consumer,
pipeline_x, mainloop_pipe_x_consumer,
pipeline_d, epi_load_pipe_d_consumer,
epi_store_pipeline, epi_store_pipe_producer_state,
mma_output_intra,
mma_output_inter,
storage.tensors.mainloop, storage.tensors.epilogue,
is_first_iteration
);
is_first_iteration = false;
}
if constexpr (CollectiveEpilogue::HasBlockScaleD) {
// update the barrier
pipeline_d.consumer_release(epi_load_pipe_d_consumer);
++epi_load_pipe_d_consumer;
}
}
uint32_t free_stage_ptr = storage.tmem_base_ptr;
tmem_allocator.free(free_stage_ptr, TmemAllocator::Sm100TmemCapacityColumns);
}
else if (warp_category == WarpCategory::PreInter) {
for (; tile_scheduler.is_valid(); ++tile_scheduler) {
auto [tState] = collective_mainloop.state_init(mma_output_inter, storage.tensors.mainloop);
for (int chunk = 0; chunk < C; ++chunk) {
collective_mainloop.pre_inter(
pipeline_b, mainloop_pipe_b_consumer,
pipeline_delta, mainloop_pipe_delta_consumer,
pipeline_inter, mainloop_pipe_inter_consumer,
mma_output_inter,
tState,
storage.tensors.mainloop
);
}
auto blk_coord = tile_scheduler.get_block_coord();
// Epilogue Fstate store
collective_epilogue.store_p(
blk_coord, params.epilogue, params.problem_size,
epi_store_p_pipeline, epi_store_p_pipe_producer_state,
storage.tensors.mainloop
);
}
}
}
};
} // namespace cutlass::ssd::kernel
examples/112_blackwell_ssd/kernel/sm100_ssd_tile_scheduler.hpp
+132
View File
@@ -0,0 +1,132 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cutlass/cutlass.h"
#include "cutlass/fast_math.h"
#include "cutlass/kernel_hardware_info.h"
namespace cutlass::ssd::kernel {
////////////////////////////////////////////////////////////////////////////////
struct PersistentTileScheduler {
struct Params {
int num_blocks;
int num_groups;
FastDivmod divmod_eh;
FastDivmod divmod_ngroup_ratio;
KernelHardwareInfo hw_info;
};
int block_idx = 0;
Params params;
CUTLASS_DEVICE
PersistentTileScheduler(Params const& params) : block_idx(blockIdx.x), params(params) {}
template<class ProblemSize, class ClusterShape, class TileShape>
static Params to_underlying_arguments(
ProblemSize const& problem_size, KernelHardwareInfo hw_info,
ClusterShape const& cluster_shape, TileShape const& tile_shape)
{
using namespace cute;
auto [G, B, EH, C, L, D, N] = problem_size;
// Get SM count if needed, otherwise use user supplied SM count
int sm_count = hw_info.sm_count;
if (sm_count <= 0) {
CUTLASS_TRACE_HOST(" WARNING: Arguments do not include a valid SM count.\n"
" For optimal performance, populate the arguments KernelHardwareInfo struct with the SM count.");
sm_count = KernelHardwareInfo::query_device_multiprocessor_count(hw_info.device_id);
}
CUTLASS_TRACE_HOST("to_underlying_arguments(): Setting persistent grid SM count to " << sm_count);
hw_info.sm_count = sm_count;
int num_blocks = B * EH;
int ngroup_ratio = EH / G;
return Params {
num_blocks,
G,
{EH},
{ngroup_ratio},
hw_info
};
}
static dim3 get_grid_shape(Params const& params) {
dim3 grid(std::min(params.num_blocks, params.hw_info.sm_count), 1, 1);
return grid;
}
CUTLASS_DEVICE
bool is_valid() {
return block_idx < params.num_blocks;
}
CUTLASS_DEVICE
auto get_block_coord() {
return block_idx;
}
CUTLASS_DEVICE
auto get_block_coord_b() {
using namespace cute;
int eh_idx, b_idx;
int g_idx, rest_idx;
params.divmod_eh(b_idx, eh_idx, block_idx);
params.divmod_ngroup_ratio(g_idx, rest_idx, eh_idx);
return (params.num_groups * b_idx + g_idx);
}
CUTLASS_DEVICE
auto get_block_coord_eh() {
using namespace cute;
int eh_idx, b_idx;
params.divmod_eh(b_idx, eh_idx, block_idx);
return eh_idx;
}
CUTLASS_DEVICE
PersistentTileScheduler& operator++() {
block_idx += gridDim.x;
return *this;
}
};
////////////////////////////////////////////////////////////////////////////////
} // namespace cutlass::ssd::kernel
examples/112_blackwell_ssd/reference/reference_ssd.hpp
+341
View File
@@ -0,0 +1,341 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cute/tensor.hpp"
// training or inference phase (not used yet)
// PHASE 0 : training
// PHASE 1 : inference
#define PHASE 0
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
bool transA,
bool transB,
class Element,
class TensorA,
class TensorB,
class TensorC
>
void mma(TensorA tA, TensorB tB, TensorC tC) {
using namespace cute;
int M = transA ? int(shape<1>(tA)) : int(shape<0>(tA));
int N = transB ? int(shape<1>(tB)) : int(shape<0>(tB));
int K = transA ? int(shape<0>(tA)) : int(shape<1>(tA));
for (int mi = 0; mi < M; ++mi) {
for (int ni = 0; ni < N; ++ni) {
for (int ki = 0; ki < K; ++ki) {
float a = static_cast<float>(Element(transA ? tA(ki, mi) : tA(mi, ki)));
float b = static_cast<float>(Element(transB ? tB(ki, ni) : tB(ni, ki)));
tC(mi, ni) += a * b;
}
}
}
}
template<
class Element,
class Tensor
>
auto segsum(Tensor tensor) {
using namespace cute;
auto C = shape<0>(tensor);
auto L = shape<1>(tensor);
auto cum_sum = make_tensor<float>(make_shape(C,L));
// cum_sum
for (int ci = 0; ci < C; ++ci) {
for (int li = 0; li < L; ++li) {
if (li == 0) {
cum_sum(ci, li) = tensor(ci, li);
}
else {
cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li);
}
}
}
auto seg_sum_out = make_tensor<float>(make_shape(C, L, L));
// seg_sum
// [ 1, 0, 0]
// [ e^y, 1, 0]
// [e^(y+z), e^z, 1]
CUTLASS_PRAGMA_UNROLL
for (int ci = 0; ci < C; ++ci) {
for (int i = 0; i < L; ++i) {
for (int j = 0; j < L; ++j) {
if (j < i) {
float tmp = static_cast<float>(cum_sum(ci, i)) - static_cast<float>(cum_sum(ci, j));
seg_sum_out(ci, i, j) = expf(tmp);
}
else if (j == i) {
seg_sum_out(ci, i, j) = 1.f;
}
else {
seg_sum_out(ci, i, j) = 0.f;
}
}
}
}
return seg_sum_out;
}
template<
class Element,
class Tensor
>
auto cumsum(Tensor tensor) {
using namespace cute;
auto C = shape<0>(tensor);
auto L = shape<1>(tensor);
auto cum_sum = make_tensor<float>(make_shape(C,L));
auto cum_sum_out = make_tensor<Element>(make_shape(C,L));
auto cum_sum_exp_out = make_tensor<float>(make_shape(C, L));
auto cum_sum_exp_out_last = make_tensor<float>(make_shape(C, L));
auto last_column = make_tensor<float>(make_shape(C));
// [x, x+y, x+y+z, ..]
CUTLASS_PRAGMA_UNROLL
for (int ci = 0; ci < C; ++ci) {
for (int li = 0; li < L; ++li) {
if (li == 0) {
cum_sum(ci, li) = tensor(ci, li);
}
else {
cum_sum(ci, li) = cum_sum(ci, li - 1) + tensor(ci, li);
}
// cum_sum_out(ci, li) = static_cast<Element>(cum_sum(ci, li));
}
}
CUTLASS_PRAGMA_UNROLL
for (int ci = 0; ci < C; ++ci) {
last_column(ci) = static_cast<float>(cum_sum(ci, L-1));
CUTLASS_PRAGMA_UNROLL
for (int li = 0; li < L; ++li) {
cum_sum_exp_out_last(ci, li) = expf(static_cast<float>(last_column(ci) - cum_sum(ci, li)));
cum_sum_exp_out(ci, li) = expf(static_cast<float>(cum_sum(ci, li)));
}
}
return make_tuple(cum_sum_exp_out_last, last_column, cum_sum_exp_out);
}
template<
bool HAS_D,
bool D_HAS_HDIM,
bool HAS_Z,
class TensorY,
class TensorF,
class TensorX,
class TensorDelta,
class TensorDeltaA,
class TensorB,
class TensorC,
class TensorD,
class TensorZ,
class Params
>
void ssd_reference_impl(
TensorY mY, TensorF mF,
TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA,
TensorB mB, TensorC mC, TensorD mD, TensorZ mZ,
Params params) {
using namespace cute;
using Element = typename Params::Element;
using ElementAcc = typename Params::ElementAcc;
// x [b, eh, d, c, l]
// delta [b, eh, c, l]
// delta_A [b, eh, c, l]
// B [b, g, n, c, l]
// C [b, g, n, c, l]
// y [b, eh, d, c, l]
// fstate [b, eh, d, n]
// d [ eh, d]
auto [G, B, EH, C, L, D, N] = params.get_problem_shape();
int group_ratio = EH / G;
for (int b = 0; b < B; ++b) {
for (int eh = 0; eh < EH; ++eh) {
int g = eh / group_ratio;
auto tY = mY(b,eh,_,_,_);
auto tF = mF(b,eh,_,_);
auto tX = mX(b,eh,_,_,_);
auto tDelta = mDelta(b,eh,_,_);
auto tDeltaA = mDeltaA(b,eh,_,_);
auto tB = mB(b,g,_,_,_);
auto tC = mC(b,g,_,_,_);
auto tD = mD(eh,_);
auto tZ = mZ(b,eh,_,_,_);
// IntraBMM1 BxC, LxLxN, NT
// B: [n, c, l]
// C: [n, c, l]
// O: [c, l, l]
auto tIntraBMM1_out = make_tensor<float>(make_shape(C,L,L));
for (int ci = 0; ci < C; ++ci) {
mma<true,true,Element>(tC(_,ci,_), tB(_,ci,_), tIntraBMM1_out(ci,_,_));
}
// Pre_IntraBMM2 DeltaA_IntraBMM2 x Delta x IntraBMM_out
// DeltaA_xxx : [c, l, l]
// Delta : [c, l, _]
// IntraBMM1_out: [c, l, l]
auto tDeltaA_IntraBMM2 = segsum<Element>(tDeltaA);
auto tIntraBMM2_inp = make_tensor<float>(make_shape(C, L, L));
for (int ci = 0; ci < C; ++ci) {
for (int i = 0; i < L; ++i) {
for (int j = 0; j < L; ++j) {
tIntraBMM2_inp(ci, i, j) = tDeltaA_IntraBMM2(ci, i, j) * tDelta(ci, j) * tIntraBMM1_out(ci, i, j);
}
}
}
// IntraBMM2 IntraBMM2_inp x X, LxDxL, TT
// IntraBMM2_inp: [c, l, l]
// X : [d, c, l]
// IntraBMM2_out: [c, l, d]
auto tIntraBMM2_out = make_tensor<float>(make_shape(C,L,D));
for (int ci = 0; ci < C; ++ci) {
mma<false,false,Element>(tIntraBMM2_inp(ci,_,_), tX(_,ci,_), tIntraBMM2_out(ci,_,_));
}
// Pre_InterBMM1 DeltaA_InterBMM1 x Delta x B
// DeltaA_xxx : [c, l]
// Delta : [c, l]
// IntraBMM1_out: [c, n, l]
auto [tDeltaA_InterBMM1, tLast, tCumsum_exp] = cumsum<Element>(tDeltaA);
auto tInterBMM1_inp = make_tensor<float>(make_shape(C, N, L));
for (int ci = 0; ci < C; ++ci) {
for (int i = 0; i < N; ++i) {
for (int j = 0; j < L; ++j) {
tInterBMM1_inp(ci, i, j) = tDeltaA_InterBMM1(ci, j) * tDelta(ci, j) * tB(i, ci, j);
}
}
}
// InterBMM1 InterBMM1_inp x X, NxDxL, swapAB, TT
// InterBMM1_inp: [c, n, l]
// X : [d, c, l]
// InterBMM1_out: [c, n, d]
auto tInterBMM1_out = make_tensor<float>(make_shape(C,N,D));
for (int ci = 0; ci < C; ++ci) {
mma<false,false,Element>(tInterBMM1_inp(ci,_,_), tX(_,ci,_), tInterBMM1_out(ci,_,_));
}
// Initialize state
// PreInterBMM2
// InterBMM1_out: [c, n, d]
// Last : [c]
auto tInterBMM2_inp = make_tensor<float>(make_shape(C, N, D));
for (int ci = 0; ci < C; ++ci) {
for (int ni = 0; ni < N; ++ ni){
for (int di = 0; di < D; ++di) {
if (ci == 0) {
tInterBMM2_inp(ci, ni, di) = 0;
}
else {
tInterBMM2_inp(ci, ni, di) = tInterBMM1_out(ci - 1, ni, di) + expf(tLast(ci - 1)) * tInterBMM2_inp(ci - 1, ni, di);
}
}
}
}
// InterBMM2 InterBMM2_inp x C, LxDxN, NT
// C : [n, c, l]
// InterBMM2_inp: [c, n, d]
// InterBMM2_out: [c, l, d]
auto tInterBMM2_out = make_tensor<float>(make_shape(C,L,D));
for (int ci = 0; ci < C; ++ci) {
mma<true,true,Element>(tC(_,ci,_), tInterBMM2_inp(ci,_,_), tInterBMM2_out(ci,_,_));
}
// Epilogue Cumsum_exp x InterBMM2_out + IntraBMM2_out
// InterBMM2_out: [c, l, d]
// IntraBMM2_out: [c, l, d]
// Cumsum_exp : [c, l]
for (int ci = 0; ci < C; ++ci) {
for (int li = 0; li < L; ++li) {
for (int di = 0; di < D; ++di) {
float y = tInterBMM2_out(ci, li, di) * tCumsum_exp(ci, li) + tIntraBMM2_out(ci, li, di);
float scale;
if constexpr (D_HAS_HDIM) {
scale = static_cast<float>(tD(di));
}
else {
scale = static_cast<float>(tD(_0{}));
}
if constexpr (HAS_D) {
y = y + static_cast<float>(tX(di, ci, li)) * scale;
}
else {
y = y;
}
if constexpr (HAS_Z) {
float z = static_cast<float>(tZ(di, ci, li));
// y = y * z * (1 / (1 + exp(-z)));
y = y * z * (1 / (1 + exp(-z)));
}
tY(di, ci, li) = static_cast<typename Params::Element>(y);
}
}
}
// Epilogue Fstate(last C)
for (int ni = 0; ni < N; ++ ni){
for (int di = 0; di < D; ++di) {
tF(di, ni) = static_cast<typename Params::Element>(tInterBMM1_out(C - 1, ni, di) + expf(tLast(C - 1)) * tInterBMM2_inp(C - 1, ni, di));
}
}
}
}
}
/////////////////////////////////////////////////////////////////////////////////////////////////
template<
bool HAS_D,
bool D_HAS_HDIM,
bool HAS_Z,
class TensorY,
class TensorF,
class TensorX,
class TensorDelta,
class TensorDeltaA,
class TensorB,
class TensorC,
class TensorD,
class TensorZ,
class Params
>
void ssd_reference(
TensorY mY, TensorF mF,
TensorX mX, TensorDelta mDelta, TensorDeltaA mDeltaA,
TensorB mB, TensorC mC, TensorD mD, TensorZ mZ,
Params params) {
ssd_reference_impl<HAS_D, D_HAS_HDIM, HAS_Z>(mY, mF, mX, mDelta, mDeltaA, mB, mC, mD, mZ, params);
}
/////////////////////////////////////////////////////////////////////////////////////////////////
examples/112_blackwell_ssd/reference/reference_ssd_cumsum.hpp
+194
View File
@@ -0,0 +1,194 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include <algorithm>
#include <random>
#include "cutlass/coord.h"
#include "cutlass/util/host_tensor.h"
#include "cutlass/tensor_view.h"
#include "cutlass/util/tensor_view_io.h"
#include "cutlass/util/reference/host/gemm.h"
#include "cutlass/arch/arch.h"
#include "cutlass/layout/matrix.h"
#include "cutlass/util/packed_stride.hpp"
#include "cutlass/cuda_host_adapter.hpp"
#include "cute/int_tuple.hpp"
#include "cute/atom/mma_traits_sm100.hpp"
#include "cute/util/debug.hpp"
#include "cute/config.hpp"
namespace cutlass::ssd::kernel {
using namespace cute;
template<
class Element_,
class ElementD_,
class TileShape_>
struct CumsumKernel {
using Element = Element_;
using ElementD = ElementD_;
using TileShape = TileShape_; // L,D,N
// Required by `device_kernel`
static constexpr int MaxThreadsPerBlock = 128;
static constexpr int MinBlocksPerMultiprocessor = 1;
using ArchTag = arch::Sm90;
static constexpr int AlignmentBytes = 16;
struct SharedStorage {
/* empty, no smem needed */
};
static constexpr int SharedStorageSize = sizeof(SharedStorage);
struct TransformArguments {
const Element* ptr_DeltaA;
ElementD* ptr_Cumsum;
};
struct TransformParams {
const Element* ptr_DeltaA;
ElementD* ptr_Cumsum;
};
using ProblemShape = cute::tuple<int, int, int, int>; // b, eh, c, l
struct Arguments {
ProblemShape problem_shape{};
TransformArguments transform{};
KernelHardwareInfo hw_info{};
};
struct Params {
ProblemShape problem_shape{};
TransformParams transform{};
KernelHardwareInfo hw_info{};
};
static Params
to_underlying_arguments(Arguments const& args, void* workspace) {
return Params{
ProblemShape{args.problem_shape},
TransformParams{args.transform.ptr_DeltaA, args.transform.ptr_Cumsum},
KernelHardwareInfo{args.hw_info}};
}
static Status
can_implement(Arguments const& args) {
return Status::kSuccess;
}
static size_t
get_workspace_size(Arguments const& args) {
return size_t(0);
}
static Status
initialize_workspace(Arguments const& args, void* workspace = nullptr, cudaStream_t stream = nullptr,
CudaHostAdapter *cuda_adapter = nullptr) {
return Status::kSuccess;
}
static dim3
get_grid_shape(Params const& params) {
auto [B, EH, C, L] = params.problem_shape;
return dim3(B*EH, 1, 1);
}
static dim3
get_block_shape() {
return dim3(MaxThreadsPerBlock, 1, 1);
}
CUTE_HOST_DEVICE
void
operator()(Params params, [[maybe_unused]] char* smem_buf = nullptr) {
auto [B, EH, C, L] = params.problem_shape;
auto layout = make_layout(make_shape(L, C, EH*B));
auto mD_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_DeltaA), make_layout(reverse(layout.shape()), reverse(layout.stride())));
auto mC_bcl = make_tensor(make_gmem_ptr(params.transform.ptr_Cumsum), make_layout(reverse(layout.shape()), reverse(layout.stride())));
auto cD_bcl = make_identity_tensor(shape(mD_bcl));
int blk_idx = blockIdx.x;
int thread_idx = threadIdx.x;
auto tD = logical_divide(mD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_);
auto tC = logical_divide(mC_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_);
auto cD = logical_divide(cD_bcl(blk_idx,_,_), make_shape(Int<128>{},_))(make_coord(thread_idx,_),_);
static constexpr int NumPacked = AlignmentBytes / sizeof(ElementD);
using PackedTypeDeltaA = uint_bit_t<sizeof_bits_v<Element> * NumPacked>;
using PackedTypeCumsum = uint_bit_t<sizeof_bits_v<ElementD> * NumPacked>;
#if 0
if (thread_idx % 128 == 0 && blk_idx == 0) {
print("tD : ");print(tD);print("\n");
print("tC : ");print(tC);print("\n");
print("cD : ");print(cD);print("\n");
}
#endif
CUTLASS_PRAGMA_UNROLL
for (int i = 0; i < shape<0>(tD); ++i) {
float last_element = 0.f;
auto crd = cD(i,_0{});
auto tD_recast = recast<PackedTypeDeltaA>(tD);
auto tC_recast = recast<PackedTypeCumsum>(tC);
if (elem_less(crd, shape(mD_bcl))) {
for (int j = 0; j < shape<1>(tD_recast); ++j) {
auto tD_slice = make_tensor<Element>(make_shape(Int<NumPacked>{}));
auto tC_slice = make_tensor<ElementD>(make_shape(Int<NumPacked>{}));
auto tD_slice_recast = recast<PackedTypeDeltaA>(tD_slice);
auto tC_slice_recast = recast<PackedTypeCumsum>(tC_slice);
tD_slice_recast(_0{}) = tD_recast(i,j);
for (int k = 0; k < NumPacked; ++ k) {
last_element += static_cast<float>(tD_slice(k));
tC_slice(k) = static_cast<ElementD>(last_element);
}
tC_recast(i,j) = tC_slice_recast(_0{});
}
}
}
}
private:
};
} // End namespace cutlass
examples/112_blackwell_ssd/utils/pipeline.h
+225
View File
@@ -0,0 +1,225 @@
/***************************************************************************************************
* 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.
*
**************************************************************************************************/
#pragma once
#include "cute/numeric/integral_constant.hpp"
#include "cute/arch/cluster_sm90.hpp"
#include "cutlass/arch/barrier.h"
#include "cutlass/pipeline/sm90_pipeline.hpp"
////////////////////////////////////////////////////////////////////////////////////////////////////
namespace cutlass {
using namespace cute;
// Producer-consumer pipeline implementation
// for TMA producer. In this case, Multi-consumers
// (UMMAs, TransformWarps, ...)
// will arrive at the same empty barrier.
// A naive implement without mcast support.
template <int Stages_, class ClusterShape = Shape<int,int,_1>, class AtomThrShape_MNK_ = Shape<_1,_1,_1>>
class PipelineTmaMultiConsumersAsync {
public:
static constexpr uint32_t Stages = Stages_;
using AtomThrShape_MNK = AtomThrShape_MNK_;
private:
using Impl = PipelineTmaAsync<Stages>;
public:
using FullBarrier = typename Impl::FullBarrier;
using EmptyBarrier = typename Impl::EmptyBarrier;
using ProducerBarrierType = typename Impl::ProducerBarrierType;
using ConsumerBarrierType = typename Impl::ConsumerBarrierType;
using PipelineState = typename Impl::PipelineState;
using SharedStorage = typename Impl::SharedStorage;
using ThreadCategory = typename Impl::ThreadCategory;
using Params = typename Impl::Params;
// Helper function to initialize barriers
static
CUTLASS_DEVICE
void
init_barriers(SharedStorage& storage, Params params, ClusterShape cluster_shape) {
int warp_idx = canonical_warp_idx_sync();
if (warp_idx == params.initializing_warp) {
constexpr int producer_arv_cnt = 1;
int const consumer_arv_cnt = params.num_consumers;
cutlass::arch::detail::initialize_barrier_array_pair_aligned<decltype(storage.full_barrier_), decltype(storage.empty_barrier_), Stages>(
storage.full_barrier_, storage.empty_barrier_, producer_arv_cnt, consumer_arv_cnt);
}
cutlass::arch::fence_barrier_init();
}
CUTLASS_DEVICE
void init_masks(ClusterShape cluster_shape) {
// Calculate consumer mask
if (params_.role == ThreadCategory::Consumer) {
is_signalling_thread_ = 1;
dst_blockid_ = 0;
}
}
// Constructor by default initializes barriers and calculates masks.
// These operations can be explicity deferred by specifying InitBarriers and InitMasks.
// If deferred, user code needs to guarantee init_masks and/or init_barriers is/are called.
template<typename InitBarriers = cute::true_type, typename InitMasks = cute::true_type>
CUTLASS_DEVICE
PipelineTmaMultiConsumersAsync(SharedStorage& storage, Params params, ClusterShape cluster_shape, InitBarriers = {}, InitMasks = {})
: impl_(storage, params, cluster_shape)
, params_(params)
, empty_barrier_ptr_(&storage.empty_barrier_[0])
, full_barrier_ptr_(&storage.full_barrier_[0]) {
static_assert(cute::is_same_v<InitBarriers, cute::true_type> || cute::is_same_v<InitBarriers, cute::false_type>);
static_assert(size(cluster_shape) == 1, "PipelineTmaMultiConsumersAsync only supports 1x1x1 cluster shape now");
if constexpr (cute::is_same_v<InitBarriers, cute::true_type>) {
init_barriers(storage, params_, cluster_shape);
}
static_assert(cute::is_same_v<InitMasks, cute::true_type> || cute::is_same_v<InitMasks, cute::false_type>);
if constexpr (cute::is_same_v<InitMasks, cute::true_type>) {
init_masks(cluster_shape);
}
}
////////////////////
// Producer APIs
////////////////////
// Four member functions are always used in pairs:
//
// * producer_try_acquire and producer_acquire, and
// * consumer_try_wait and consumer_wait.
//
// The two functions with "try" in their names are called "try" functions,
// and the other two are conceptually "finalize" functions.
// The "try" function in each pair starts the process of waiting on the barrier to flip.
// It opportunistically waits for an implementation-dependent timeout.
// Whether or not the barrier has flipped yet, the try function will return a token.
// If the token indicates that the barrier has not flipped,
// then the token must be passed into the corresponding "finalize" function.
// The finalize function will then block until the barrier has flipped.
// If the token indicates that the barrier _has_ flipped,
// then it is still correct to pass it into the finalize function.
// The finalize function will return immediately in that case.
CUTLASS_DEVICE
ProducerToken producer_try_acquire(PipelineState state, uint32_t skip_wait = false) {
return impl_.producer_try_acquire(state, skip_wait);
}
CUTLASS_DEVICE
void producer_acquire(PipelineState state, ProducerToken barrier_token = {BarrierStatus::WaitAgain}) {
impl_.producer_acquire(state, barrier_token);
}
// NOP for TMA based mainloop
CUTLASS_DEVICE
void producer_commit(PipelineState state, uint32_t bytes) {
impl_.producer_commit(state, bytes);
}
// Prevents early exit of producer blocks in Cluster.
// This should be called once before kernel exits.
CUTLASS_DEVICE
void producer_tail(PipelineState state) {
impl_.producer_tail(state);
}
CUTLASS_DEVICE
ProducerBarrierType* producer_get_barrier(PipelineState state) {
return impl_.producer_get_barrier(state);
}
////////////////////
// Consumer APIs
////////////////////
CUTLASS_DEVICE
ConsumerToken consumer_try_wait(PipelineState state, uint32_t skip_wait = false) {
return impl_.consumer_try_wait(state, skip_wait);
}
CUTLASS_DEVICE
void consumer_wait(PipelineState state, ConsumerToken barrier_token = {BarrierStatus::WaitAgain}) {
impl_.consumer_wait(state, barrier_token);
}
CUTLASS_DEVICE
void consumer_release_from_umma(PipelineState state) {
consumer_release_from_umma(state.index(), false);
}
CUTLASS_DEVICE
void consumer_release_from_threads(PipelineState state) {
consumer_release_from_threads(state.index());
}
private:
Impl impl_;
Params params_;
uint32_t dst_blockid_ = 0;
uint32_t is_signalling_thread_ = 0;
EmptyBarrier *empty_barrier_ptr_;
FullBarrier *full_barrier_ptr_;
uint16_t block_id_mask_ = 0;
static constexpr bool is_2sm_mma = size(AtomThrShape_MNK{}) > 1;
// Consumer signalling Producer of completion
// Ensures all blocks in the Same Row and Column get notifed.
CUTLASS_DEVICE
void consumer_release_from_umma(uint32_t stage, uint32_t skip) {
uint64_t* smem_ptr = reinterpret_cast<uint64_t*>(&empty_barrier_ptr_[stage]);
if constexpr (is_2sm_mma) { // Mma cluster shape is 2x1
if (!skip) {
cutlass::arch::umma_arrive_multicast_2x1SM(smem_ptr, block_id_mask_);
}
}
else {
if (!skip) {
if constexpr (cute::is_static_v<ClusterShape> and size(ClusterShape{}) == 1) {
cutlass::arch::umma_arrive(smem_ptr);
}
else {
cutlass::arch::umma_arrive_multicast(smem_ptr, block_id_mask_);
}
}
}
}
CUTLASS_DEVICE
void consumer_release_from_threads(uint32_t stage, uint32_t skip = false) {
empty_barrier_ptr_[stage].arrive(dst_blockid_, is_signalling_thread_ & (!skip));
#ifndef NDEBUG
if (params_.role == ThreadCategory::Producer || params_.role == ThreadCategory::NonParticipant) {
asm volatile ("brkpt;\n" ::);
}
#endif
}
};
}
examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_bf16_gemm.cu
+33 -5
View File
@@ -115,9 +115,29 @@ using namespace cute;
/////////////////////////////////////////////////////////////////////////////////////////////////
/// GEMM kernel configurations
/////////////////////////////////////////////////////////////////////////////////////////////////
using MmaType = cutlass::bfloat16_t;
using QuantType = cutlass::int4b_t;
constexpr int TileShapeK = 128 * 8 / sizeof_bits<MmaType>::value;
// Select MMA type via compile flag
#if defined(CUTLASS_USE_FP16)
using MmaType = cutlass::half_t; // FP16
#elif defined(CUTLASS_USE_TF32)
using MmaType = cutlass::tfloat32_t; // TF32 (FP32 format with reduced precision)
#else
using MmaType = cutlass::bfloat16_t; // BF16 (default)
#endif
// Select quantization type via compile flag for this example
#if defined(CUTLASS_MIXED_DTYPE_E2M1)
using QuantType = cutlass::float_e2m1_t; // E2M1 (FP4)
#else
using QuantType = cutlass::int4b_t; // INT4 Two's Complement (default)
#endif
// TF32 requires K to be multiple of 8; BF16/FP16 can go higher
#if defined(CUTLASS_USE_TF32)
constexpr int TileShapeK = 64; // TF32: K must be multiple of 8, use 64 for good performance
#else
constexpr int TileShapeK = 128 * 8 / sizeof_bits<MmaType>::value;
#endif
// A matrix configuration
using ElementA = MmaType; // Element type for A matrix operand
@@ -139,11 +159,19 @@ using StrideB = cutlass::detail::TagToStrideB_t<LayoutB>;
// Define the CuTe layout for reoredered quantized tensor B
// LayoutAtomQuant places values that will be read by the same thread in contiguous locations in global memory.
// It specifies the reordering within a single warp's fragment
//using ValueShuffle = Layout<_1>; // no value reordering
#if defined(CUTLASS_MIXED_DTYPE_E2M1) || defined(CUTLASS_USE_TF32)
// E2M1 & TF32: Use simpler layout without ValueShuffle (like FP8 example)
// ValueShuffle currrently isn't enabled for E2M1 until LayoutAwareConvertImpl specializations support shuffle reordering.
// and for TF32 until we support both the K=8 vs K=16 dimensions for tiles
using LayoutAtomQuant = decltype(cutlass::compute_memory_reordering_atom<MmaType>());
#else
// INT4: Use ValueShuffle for optimal performance with FP16/BF16
// using ValueShuffle = Layout<_1>; // no value reordering
using ValueShuffle = Layout<Shape<_2,_4>, Stride<_4,_1>>; // order [0,2,4,6,1,3,5,7]
int constexpr NumShuffleAtoms = 1;
using MmaAtomShape = Layout<Shape<_1,Int<NumShuffleAtoms>>>;
using LayoutAtomQuant = decltype(cutlass::compute_memory_reordering_atom<MmaType, MmaAtomShape, ValueShuffle>());
#endif
using LayoutB_Reordered = decltype(cute::tile_to_shape(LayoutAtomQuant{}, Layout<Shape<int,int,int>, StrideB>{}));
using ElementScale = MmaType;
@@ -151,7 +179,7 @@ using ElementZero = ElementScale;
using LayoutScale = cutlass::layout::RowMajor;
// C/D matrix configuration
using ElementC = cutlass::half_t; // Element type for C and D matrix operands
using ElementC = MmaType; // Element type for C and D matrix operands (matches MMA type)
using LayoutC = cutlass::layout::RowMajor; // Layout type for C and D matrix operands
constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value; // Memory access granularity/alignment of C matrix in units of elements (up to 16 bytes)
examples/55_hopper_mixed_dtype_gemm/55_hopper_int4_fp8_gemm.cu
+37 -5
View File
@@ -120,7 +120,15 @@ using namespace cute;
/// GEMM kernel configurations
/////////////////////////////////////////////////////////////////////////////////////////////////
using MmaType = cutlass::float_e4m3_t;
using QuantType = cutlass::int4b_t;
// Select quantization type via compile flag for this example
// templatized throughout code to enable instantiating both int4 and e2m1 versions in the same program
#if defined(CUTLASS_MIXED_DTYPE_E2M1)
using QuantType = cutlass::float_e2m1_t; // E2M1 (FP4)
#else
using QuantType = cutlass::int4b_t; // INT4 Two's Complement (default)
#endif
constexpr int TileShapeK = 128 * 8 / sizeof_bits<MmaType>::value;
// A matrix configuration
@@ -315,6 +323,7 @@ struct Options : MixedDtypeOptions {
/////////////////////////////////////////////////////////////////////////////////////////////////
/// Initialize operands to be used in the GEMM and reference GEMM
template <typename QuantType>
void initialize(Options const& options) {
auto shape_B = cute::make_shape(options.n, options.k, options.l);
@@ -330,7 +339,7 @@ void initialize(Options const& options) {
auto layout_B = make_layout(shape_B, stride_B);
auto a_coord = cutlass::make_Coord(options.m * options.l, options.k);
auto a_coord = cutlass::make_Coord(options.m * options.l, options.k);
auto b_coord = cutlass::make_Coord(options.k, options.n * options.l);
auto c_coord = cutlass::make_Coord(options.m * options.l, options.n);
@@ -346,14 +355,15 @@ void initialize(Options const& options) {
block_scale_packed.reset(scale_k * options.l * options.n);
block_zero.reset(scale_k * options.l * options.n);
// Initialize all base tensors
initialize_tensor(block_A, seed + 2022);
initialize_tensor(block_B, seed + 2021);
cutlass::unified_encode_int4b(block_B.get(), block_B_modified.get(), block_B.size());
initialize_tensor(block_C, seed + 2020);
initialize_scale(block_scale, options);
cutlass::pack_scale_fp8(block_scale.get(), block_scale_packed.get(), block_scale.size());
cutlass::pack_scale_fp8<ElementScale, QuantType>(block_scale.get(), block_scale_packed.get(), block_scale.size());
initialize_zero(block_zero, options);
// Compute dequantized reference for validation BEFORE formatting block_B
auto shape_scale_zero = cute::make_shape(options.n, scale_k, options.l);
stride_S = cutlass::make_cute_packed_stride(StrideS{}, cute::make_shape(options.n, scale_k, options.l));
stride_S_ref = cutlass::make_cute_packed_stride(StrideS_ref{}, cute::make_shape(options.n, scale_k, options.l));
@@ -362,6 +372,28 @@ void initialize(Options const& options) {
cudaStream_t stream = cudaStreamDefault;
cutlass::dequantize(block_B_dq.get(), block_B.get(), layout_B, block_scale.get(), block_zero.get(), layout_scale_zero, options.g, stream);
// Format B to separate buffer, preserving original block_B in case
// it's needed by the application.
#ifdef CUTLASS_MIXED_DTYPE_E2M1
// E2M1: Copy to formatting buffer
cutlass::device_memory::copy_device_to_device(block_B_modified.get(), block_B.get(), block_B.size());
#else
// INT4: Encode to formatting buffer
cutlass::unified_encode_int4b(block_B.get(), block_B_modified.get(), block_B.size());
#endif
// original code
// if constexpr (cutlass::platform::is_floating_point<QuantType>::value ||
// cute::is_same_v<QuantType, cutlass::float_e2m1_t>) {
// // E2M1: Copy to formatting buffer
// cutlass::device_memory::copy_device_to_device(block_B_modified.get(), block_B.get(), block_B.size());
// } else {
// // INT4: Encode to formatting buffer
// cutlass::unified_encode_int4b(block_B.get(), block_B_modified.get(), block_B.size());
// }
if (options.shuffle) {
// Repeat the reorder layout atom to tile the whole tensor shape
layout_B_reordered = cute::tile_to_shape(LayoutAtomQuant{}, shape_B);
@@ -464,7 +496,7 @@ bool verify(Options const& options) {
template <typename Gemm>
int run(Options &options)
{
initialize(options);
initialize<QuantType>(options);
// Instantiate CUTLASS kernel depending on templates
Gemm gemm;
examples/60_cutlass_import/CMakeLists.txt
+1 -1
View File
@@ -54,7 +54,7 @@ target_sources(example PRIVATE main.cpp)
target_include_directories(
example
PRIVATE
SYSTEM PRIVATE
${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
)
examples/77_blackwell_fmha/77_blackwell_fmha_gen.cu
+2 -2
View File
@@ -730,8 +730,8 @@ int main_single(int argc, char const **args) {
if (props.major != 10 || (props.minor != 0 && props.minor != 3)) {
std::cout
<< "This example requires a GPU of NVIDIA's Blackwell Architecture "
<< "(compute capability 90) and CUDA 12.0 or greater.\n";
<< "This example requires a GPU of NVIDIA's Blackwell Datacenter-class Architecture "
<< "(compute capability 100 or 103) and CUDA 12.0 or greater.\n";
return 0;
}
examples/77_blackwell_fmha/77_blackwell_mla.cu
+7 -3
View File
@@ -75,7 +75,11 @@ struct Options {
int split_kv = -1; // number of split along k dim.
bool is_var_split_kv = false;
int max_split_kv = 16;
#ifdef CPASYNC
int page = 1;
#else
int page = -1;
#endif
float spread = 0.2f;
int iterations = 3;
bool verify = false;
@@ -260,7 +264,7 @@ struct ExampleResult {
///////////////////////////////////////////////////////////////////////////////////////////////////
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED))
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -751,7 +755,7 @@ void run_mla(Options const & options, cutlass::KernelHardwareInfo const& hw_info
///////////////////////////////////////////////////////////////////////////////////////////////////
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
#endif // defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED)
///////////////////////////////////////////////////////////////////////////////////////////////////
@@ -796,7 +800,7 @@ int main_single(int argc, char const **args) {
return -1;
}
#if defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED)
#if (defined(CUTLASS_ARCH_MMA_SM100_SUPPORTED) || defined(CUTLASS_ARCH_MMA_SM103_SUPPORTED))
//
// Run examples
examples/77_blackwell_fmha/CMakeLists.txt
+1 -1
View File
@@ -112,7 +112,7 @@ set(TEST_MLA_FUSE_REDUCTION --b=1 --k=4096 --split_kv=8 --page=128 --fuse_reduct
set(TEST_MLA_LARGE_SPLIT_KV --verify --split_kv=20 --page=128)
if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC_ARCHS MATCHES 100a))
if(NOT WIN32 AND (NOT (CMAKE_CXX_COMPILER_ID MATCHES "Clang")) AND (CUTLASS_NVCC_ARCHS MATCHES 100a OR CUTLASS_NVCC_ARCHS MATCHES 103a))
foreach(PREC fp8 fp16)
string(TOUPPER "${PREC}" PREC_MACRO)
examples/77_blackwell_fmha/README.md
+2
View File
@@ -69,6 +69,8 @@ For detailed information on how to invoke them, check out either the tests in `C
* 4.3.0: For variable sequence length, the code requires a batch of valid (but never used) padding memory ahead of the first output batch. No padding is needed for the input tensor, but it requires that the input tensor contain no NaN or Inf values. Note that users should set `total_length` to the `problem_shape`.
* 4.4.0: Added support for Blackwell Ultra (Sm103).
# Copyright
Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
examples/77_blackwell_fmha/kernel/fmha_kernel_bwd_convert.hpp
+1 -1
View File
@@ -69,7 +69,7 @@ struct FmhaKernelBwdConvert {
static const int MinBlocksPerMultiprocessor = 1;
static const int MaxThreadsPerBlock = 128;
using ArchTag = cutlass::arch::Sm90;
using ArchTag = cutlass::arch::Sm100;
static const int kBlockSeq = 8;
examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_kernel_tma_warpspecialized.hpp
+2 -1
View File
@@ -1508,7 +1508,8 @@ struct Sm100FmhaBwdKernelTmaWarpSpecialized {
CUTLASS_DEVICE void operator()(Params const& params, char* smem) {
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED))
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \
! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED))
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
int warp_idx = cutlass::canonical_warp_idx_sync();
examples/77_blackwell_fmha/kernel/sm100_fmha_bwd_mla_kernel_tma_warpspecialized.hpp
+2 -1
View File
@@ -1480,7 +1480,8 @@ struct Sm100FmhaBwdMlaKernelTmaWarpSpecialized {
CUTLASS_DEVICE void operator()(Params const& params, char* smem) {
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED))
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \
! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED))
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
int warp_idx = cutlass::canonical_warp_idx_sync();
examples/77_blackwell_fmha/kernel/sm100_fmha_fwd_kernel_tma_warpspecialized.hpp
+2 -1
View File
@@ -251,7 +251,8 @@ struct Sm100FmhaFwdKernelTmaWarpspecialized {
}
CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED))
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \
! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED))
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
examples/77_blackwell_fmha/kernel/sm100_fmha_gen_kernel_warpspecialized.hpp
+2 -1
View File
@@ -247,7 +247,8 @@ struct Sm100FmhaGenKernelWarpspecialized {
}
CUTLASS_DEVICE void operator()(const Params &params, char* smem) {
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED))
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \
! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED))
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
examples/77_blackwell_fmha/kernel/sm100_fmha_mla_tma_warpspecialized.hpp
+2 -1
View File
@@ -507,7 +507,8 @@ struct Sm100FmhaMlaKernelTmaWarpspecialized {
CUTLASS_DEVICE void operator()(Params const& params, char* smem_raw) {
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED))
#if (! defined(CUTLASS_ARCH_MMA_SM100A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM100F_ENABLED) && \
! defined(CUTLASS_ARCH_MMA_SM103A_ENABLED) && ! defined(CUTLASS_ARCH_MMA_SM103F_ENABLED))
CUTE_INVALID_CONTROL_PATH("ERROR : Arch conditional MMA instruction used without targeting appropriate compute capability. Aborting.\n");
#else
examples/81_blackwell_gemm_blockwise/README.md
+1 -1
View File
@@ -50,7 +50,7 @@ To determine the most performance Blockwise/Groupwise GEMM or Grouped GEMM kerne
All Blockwise/Groupwise GEMMs and Group GEMMs with `f32` scaling of `e4m3` or runtime `f8` types can be selected by
selecting a subset of kernels when configuring with CMake by passing:
`-DCUTLASS_LIBRARY_KERNELS="cutlass3x*f32xe4m3_*f32xe4m3*,cutlass3x*f32xf8_*f32xf8*"` (you can further reduce the amount of kernels generated by specifying the SFA and SFB scale granularities e.g., `cutlass3x*1x128f32xe4m3_*128x128f32xe4m3*`).
`-DCUTLASS_LIBRARY_KERNELS="cutlass3x*f32xe4m3_*f32xe4m3*,cutlass3x*f32xf8_*f32xf8*"` you can further reduce the amount of kernels generated by specifying the SFA and SFB scale granularities e.g., `cutlass3x*1x128f32xe4m3_*128x128f32xe4m3*`).
The simplest way to use the profiler is to pass `m`, `n`, and `k` as well as your `scale_vec_size_m`,
`scale_vec_size_n`, and `scale_vec_size_k`. Passing `enable-best-kernel-for-fixed-shape` will do some autotuning
examples/89_sm103_fp4_ultra_gemm/89_sm103_fp4_ultra_gemm.cu
+2 -2
View File
@@ -32,7 +32,7 @@
/*! \file
\brief A GEMM example using CUTLASS for the NVIDIA Blackwell SM103 architecture.
This example demonstrates a simple way to instantiate and run a blockscaled 3xFP4 GEMM on the NVIDIA Blackwell SM103 architecture.
This example demonstrates a simple way to instantiate and run a blockscaled ultra FP4 GEMM on the NVIDIA Blackwell SM103 architecture.
Usage:
@@ -269,7 +269,7 @@ struct Options {
std::ostream & print_usage(std::ostream &out) const {
out << "89_sm103_fp4_ultra_gemm\n\n"
<< " Sm103 3xFP4 GEMM using a Warp Specialized kernel.\n\n"
<< " Sm103 ultra FP4 GEMM using a Warp Specialized kernel.\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --m=<int> Sets the M extent of the GEMM\n"
examples/90_sm103_fp4_ultra_grouped_gemm/90_sm103_fp4_ultra_grouped_gemm.cu
+2 -2
View File
@@ -441,7 +441,7 @@ struct Options {
std::ostream & print_usage(std::ostream &out) const {
out << "90_sm103_fp4_ultra_grouped_gemm\n\n"
<< " Sm103 3xFP4 Grouped GEMM using a Warp Specialized kernel.\n\n"
<< " Sm103 ultra FP4 Grouped GEMM using a Warp Specialized kernel.\n\n"
<< "Options:\n\n"
<< " --help If specified, displays this usage statement\n\n"
<< " --m=<int> Sets the M extent of the GEMM for all groups\n"
@@ -963,7 +963,7 @@ int run(Options &options, bool host_problem_shapes_available = true)
int main(int argc, char const **args) {
std::cout << "\n====================================================" << std::endl;
std::cout << "CUTLASS 3.0 Grouped GEMM Example - 3xfp4 Block Scaled" << std::endl;
std::cout << "CUTLASS 3.0 Grouped GEMM Example - ultra fp4 Block Scaled" << std::endl;
std::cout << "====================================================" << std::endl;
// CUTLASS must be compiled with CUDA 12.9 Toolkit to run this example
examples/93_blackwell_low_latency_gqa/CMakeLists.txt
+36
View File
@@ -0,0 +1,36 @@
# Copyright (c) 2014 - 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.
if (NOT MSVC AND CUTLASS_NVCC_ARCHS MATCHES "100a|100f|103a|103f")
cutlass_example_add_executable(
93_blackwell_low_latency_gqa
tgv_gqa.cu
)
endif()
examples/93_blackwell_low_latency_gqa/figures/acc2.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 966 KiB

examples/93_blackwell_low_latency_gqa/figures/cta.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 1.1 MiB

examples/93_blackwell_low_latency_gqa/figures/fmax.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 619 KiB

examples/93_blackwell_low_latency_gqa/figures/fsum.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 867 KiB

examples/93_blackwell_low_latency_gqa/figures/tgv_gqa.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 980 KiB

examples/93_blackwell_low_latency_gqa/readme.md
+73
View File
@@ -0,0 +1,73 @@
# Blackwell Low Latency GQA
This example introduces TGV GQA, a CuTe C++-based Blackwell kernel optimized for low latency (low batch) generation phase GQA.
To compile and run this example:
```bash
# in cutlass top level directory
mkdir build && cd build
cmake .. -DCUTLASS_NVCC_ARCHS=100a -DCUTLASS_ENABLE_TESTS=OFF -DCUTLASS_ENABLE_EXAMPLES=ON -DCUTLASS_ENABLE_LIBRARY=OFF
cd examples/93_blackwell_low_latency_gqa
make
./93_blackwell_low_latency_gqa --kvL 8192 --kvH 8 --qH 64 --BS 1
```
Supported configs are:
- dH = 64/128
- bf16/fp8 non block scaling kv cache, fp32/bf16/fp8 output
- QKVO are all dH major
- Arbitrary seq len and batch size
- CUDA graph support
- Flash decoding, configurable number of splits
- Cluster reduction with configurable number of reduction cta
- Attention sink and sliding window
Unsupported features are:
- Persistent schedule
- MTP
- Paged KV cache
## Kernel Design
Each cluster (of size `1x1xMAX_SPLITS`) works on a single kv head in a batch.
It divides the kvL (kv sequence length) evenly into each CTA in the cluster in flash decoding style.
And the final reduction is performed by `NUM_REDUCTION_CTA` number of CTAs in the cluster in parallel.
![cta](./figures/cta.png)
The example figure above shows a problem size of 1 batch, 2 KV heads, 4 Q heads.
Each cluster processes 1 batch, 1 KV head, 2 Q heads.
And the KV sequence length is divided into 6 128-length tiles, and we evenly distribute the 6 tiles to 4 CTAs in the cluster.
CTA0 and CTA1 will get 2 tiles while CTA2 and CTA3 will get 1 tile each.
In the reduction phase of flash decoding, in this configuration, only 2 CTAs in the cluster will participate in the final reduction.
So each of CTA in the cluster will send their partial results (`Acc2`) to the 2 reduction CTAs in the cluster.
We have 7 warps in total, 1 for DMA_Q, 1 for DMA_KV, 1 for MMA, 4 for EPILOG (softmax + cluster reduction).
The imagine below shows the how the data flows across the warps as well as how the control dependencies are established between the warps.
![tgv_gqa](./figures/tgv_gqa.png)
## Fmax Reduction Mapping
We want to get fmax for each column (q token) of the fmax tensor.
Each thread individually get fmax for all q tokens using credux (and inter warp reduction).
Then T0,32,64,96 stores the local fmax to destination cta's dsmem for whole cluster wide fmax reduction.
Each reduction cta will hold the cluster wide fmax values of 2 q tokens (8 q tokens in total, divided by 4 reduction ctas).
![fmax_mapping](./figures/fmax.png)
## Fsum Reduction Mapping
We want to get fsum for each column (q token) of the fsum tensor.
Do a reswizzle of fsum tensor in rmem through smem such that each thread holds a partial column of the fsum tensor.
Then do intra-thread and intra-warp reduction to get the fsum for each column (q token).
Then T0,32,64,96 stores the local fsum to destination cta's dsmem for whole cluster wide fsum reduction.
Each reduction cta will hold the cluster wide fsum values of 2 q tokens (8 q tokens in total, divided by 4 reduction ctas).
![fsum_mapping](./figures/fsum.png)
## Acc2 Reduction Mapping
Each thread in the reduction cta will be responsible for generating 2 q tokens in the final output tensor (8 q tokens in total, divided by 4 reduction ctas).
The reduction is a scaled accumulation (similar to the correction step in attention mainloop).
![acc2_mapping](./figures/acc2.png)
examples/93_blackwell_low_latency_gqa/tgv_gqa.cu
+750
View File
@@ -0,0 +1,750 @@
/***************************************************************************************************
* 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 Example implementation of low latency GQA for the NVIDIA Blackwell SM100/SM103
architecture using CUTLASS 3.
Input tensor shapes:
K has shape (kvL, dH, kvH, BS)
Q has shape ((qHLocal, qL), dH, kvH, BS)
V has shape (dH, kvL, kvH, BS)
O has shape (dH, (qHLocal, qL), kvH, BS)
kvL is max_seq_len, seq_lens[BS] is the actual seq len for each batch
sinks has shape (qHLocal * kvH), i.e. one sink per q head
Example usage:
$ ./examples/93_blackwell_low_latency_gqa --kvL 8192 --kvH 8 --qH 64 --BS 1
*/
// Standard library includes
#include <cassert>
#include <vector>
#include <memory>
#include <cmath>
#include <iostream>
#include <ctime>
#include <getopt.h>
#include <cuda_runtime.h>
#include <cuda_profiler_api.h>
// Use Thrust to handle host/device allocations
#include <thrust/host_vector.h>
#include <thrust/device_vector.h>
// Cutlass includes
#include <cutlass/numeric_types.h>
#include <cutlass/util/print_error.hpp>
#include <cutlass/numeric_conversion.h>
// CuTe includes
#include <cute/tensor.hpp> // CuTe tensor implementation
#include "tgv_gqa.cuh"
using namespace cute;
// K (kvL, dH, kvH, BS)
// Q ((qHLocal, qL), dH, kvH, BS)
// V (dH, kvL, kvH, BS)
// O (dH, (qHLocal, qL), kvH, BS)
// seq_lens (BS)
// Sinks ((qHLocal, qL), kvH)
template <
class TypeAcc,
int CTA_kvL,
bool NoSink,
class TensorQ,
class TensorK,
class TensorV,
class TensorO,
class TensorSinks>
void
reference_gqa(
TensorK const& tensor_K,
TensorQ const& tensor_Q,
TensorV const& tensor_V,
TensorO const& tensor_O,
int* seq_lens,
TensorSinks const& tensor_Sinks,
float softmax_scale,
int sliding_window_size) {
using TypeQKV = typename TensorQ::element_type;
using TypeO = typename TensorO::element_type;
using namespace cute;
int kvL = size<0>(tensor_K);
int dH = size<1>(tensor_K);
int kvH = size<2>(tensor_K);
int BS = size<3>(tensor_K);
int qHLocal = size<0>(shape<0>(tensor_Q));
int qL = size<1>(shape<0>(tensor_Q));
// reference code don't handle oob either
//assert(kvL % CTA_kvL == 0);
int MaxKVBlocks = cutlass::ceil_div(kvL, CTA_kvL);
// allocate intermediate tensors
thrust::host_vector<TypeAcc> host_Acc1(kvL * qHLocal * qL * kvH * BS);
auto tensor_Acc1 = make_tensor(host_Acc1.data(), make_layout(make_shape(kvL, make_shape(qHLocal, qL), kvH, BS))); // (kvL, (qHLocal, qL), kvH, BS)
// CTA level fmax and fsum
thrust::host_vector<TypeAcc> host_Fmax(MaxKVBlocks * qHLocal * qL * kvH * BS);
Tensor tensor_Fmax = make_tensor(host_Fmax.data(), make_layout(make_shape(MaxKVBlocks, make_shape(qHLocal, qL), kvH, BS))); // (MaxKVBlocks, (qHLocal, qL), kvH, BS)
fill(tensor_Fmax, -cutlass::platform::numeric_limits<TypeAcc>::infinity());
thrust::host_vector<TypeAcc> host_Fsum(MaxKVBlocks * qHLocal * qL * kvH * BS);
Tensor tensor_Fsum = make_tensor(host_Fsum.data(), make_layout(make_shape(MaxKVBlocks, make_shape(qHLocal, qL), kvH, BS))); // (MaxKVBlocks, (qHLocal, qL), kvH, BS)
clear(tensor_Fsum);
thrust::host_vector<TypeAcc> host_P(CTA_kvL * MaxKVBlocks * qHLocal * qL * kvH * BS);
Tensor tensor_P = make_tensor(host_P.data(), make_layout(make_shape(make_shape(CTA_kvL, MaxKVBlocks), make_shape(qHLocal, qL), kvH, BS))); // ((CTA_kvL, MaxKVBlocks), (qHLocal, qL), kvH, BS)
thrust::host_vector<TypeQKV> host_P_converted(CTA_kvL * MaxKVBlocks * qHLocal * qL * kvH * BS);
Tensor tensor_P_converted = make_tensor(host_P_converted.data(), make_layout(make_shape(make_shape(CTA_kvL, MaxKVBlocks), make_shape(qHLocal, qL), kvH, BS))); // ((CTA_kvL, MaxKVBlocks), (qHLocal, qL), kvH, BS)
thrust::host_vector<TypeAcc> host_Acc2(dH * qHLocal * qL * kvH * MaxKVBlocks * BS);
auto tensor_Acc2 = make_tensor(host_Acc2.data(), make_layout(make_shape(dH, make_shape(qHLocal, qL), kvH, BS, MaxKVBlocks))); // (dH, (qHLocal, qL), kvH, BS, MaxKVBlocks)
// cluster level fmax and fsum
thrust::host_vector<TypeAcc> host_Fmax_cluster(qHLocal * qL * kvH * BS);
Tensor tensor_Fmax_cluster = make_tensor(host_Fmax_cluster.data(), make_layout(make_shape(make_shape(qHLocal, qL), kvH, BS))); // ((qHLocal, qL), kvH, BS)
thrust::host_vector<TypeAcc> host_Beta(qHLocal * qL * kvH * MaxKVBlocks * BS);
Tensor tensor_Beta = make_tensor(host_Beta.data(), make_layout(make_shape(make_shape(qHLocal, qL), kvH, BS, MaxKVBlocks))); // ((qHLocal, qL), kvH, BS, MaxKVBlocks)
thrust::host_vector<TypeAcc> host_Fsum_cluster(qHLocal * qL * kvH * BS);
Tensor tensor_Fsum_cluster = make_tensor(host_Fsum_cluster.data(), make_layout(make_shape(make_shape(qHLocal, qL), kvH, BS))); // ((qHLocal, qL), kvH, BS)
float softmax_scale_log2 = softmax_scale * static_cast<float>(M_LOG2E);
for (int _BS = 0; _BS < BS; ++_BS) {
int seq_len = seq_lens[_BS];
int NumKVBlocks = cutlass::ceil_div(seq_len, CTA_kvL);
int kvL_start = (sliding_window_size == 0) ? 0 : std::max(0, seq_len - sliding_window_size);
int kvBlock_start = kvL_start / CTA_kvL;
// bmm1 s = q * k * softmax_scale_log2
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
int start = _kvBlock * CTA_kvL;
int end = std::min(start + CTA_kvL, seq_len);
assert(start < end);
for (int _kvL = start; _kvL < end; ++_kvL) {
TypeAcc acc = TypeAcc(0.f);
for (int _dH = 0; _dH < dH; ++_dH) {
acc += tensor_K(_kvL, _dH, _kvH, _BS) * tensor_Q(make_coord(_qHLocal, _qL), _dH, _kvH, _BS);
}
tensor_Acc1(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS) = (_kvL < kvL_start) ? -INFINITY : (acc * softmax_scale_log2);
}
}
}
}
}
// calculate m_ij
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
int start = _kvBlock * CTA_kvL;
int end = std::min(start + CTA_kvL, seq_len);
assert(start < end);
TypeAcc& fmax = tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS);
fmax = tensor_Acc1(start, make_coord(_qHLocal, _qL), _kvH, _BS);
for (int _kvL = start + 1; _kvL < end; ++_kvL) {
fmax = std::max(fmax, tensor_Acc1(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS));
}
}
}
}
}
// calculate p = exp2f(s - m_ij)
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
int start = _kvBlock * CTA_kvL;
int end = std::min(start + CTA_kvL, seq_len);
assert(start < end);
for (int _kvL = start; _kvL < end; ++_kvL) {
tensor_P(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS) = std::exp2f(tensor_Acc1(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS) - tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS));
}
}
}
}
}
// calculate l_ij = sum(p)
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
int start = _kvBlock * CTA_kvL;
int end = std::min(start + CTA_kvL, seq_len);
assert(start < end);
TypeAcc sum = TypeAcc(0.f);
for (int _kvL = start; _kvL < end; ++_kvL) {
sum += tensor_P(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS);
}
tensor_Fsum(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS) = sum;
}
}
}
}
// convert P from fp32 to bf16/fp8
cutlass::NumericConverter<TypeQKV, TypeAcc> converter_p;
for (int i = 0; i < tensor_P(_,_,_,_BS).size(); i++) {
tensor_P_converted(_,_,_,_BS)[i] = converter_p(tensor_P(_,_,_,_BS)[i]);
}
// bmm2 acc2 = v * p
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _dH = 0; _dH < dH; ++_dH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
int start = _kvBlock * CTA_kvL;
int end = std::min(start + CTA_kvL, seq_len);
assert(start < end);
TypeAcc acc = TypeAcc(0.f);
for (int _kvL = start; _kvL < end; ++_kvL) {
acc += tensor_V(_dH, _kvL, _kvH, _BS) * tensor_P_converted(_kvL, make_coord(_qHLocal, _qL), _kvH, _BS);
}
tensor_Acc2(_dH, make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock) = acc;
}
}
}
}
}
// calculate cluster level fmax and fsum
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
TypeAcc& fmax = tensor_Fmax_cluster(make_coord(_qHLocal, _qL), _kvH, _BS);
fmax = tensor_Fmax(kvBlock_start, make_coord(_qHLocal, _qL), _kvH, _BS);
for (int _kvBlock = kvBlock_start + 1; _kvBlock < NumKVBlocks; ++_kvBlock) {
fmax = std::max(fmax, tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS));
}
// calculate beta
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
tensor_Beta(make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock) = std::exp2f(tensor_Fmax(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS) - fmax);
}
// calculate fsum
TypeAcc sum = TypeAcc(0.f);
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
sum += tensor_Fsum(_kvBlock, make_coord(_qHLocal, _qL), _kvH, _BS) * tensor_Beta(make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock);
}
tensor_Fsum_cluster(make_coord(_qHLocal, _qL), _kvH, _BS) = sum;
if constexpr (!NoSink) {
tensor_Fsum_cluster(make_coord(_qHLocal, _qL), _kvH, _BS) += std::exp2f(tensor_Sinks(make_coord(_qHLocal, _qL), _kvH) * (float)M_LOG2E - fmax);
}
}
}
}
// convert O from fp32 to bf16/fp8
cutlass::NumericConverter<TypeO, TypeAcc> converter_o;
// final reduction
for (int _kvH = 0; _kvH < kvH; ++_kvH) {
for (int _dH = 0; _dH < dH; ++_dH) {
for (int _qHLocal = 0; _qHLocal < qHLocal; ++_qHLocal) {
for (int _qL = 0; _qL < qL; ++_qL) {
TypeAcc acc = TypeAcc(0.f);
for (int _kvBlock = kvBlock_start; _kvBlock < NumKVBlocks; ++_kvBlock) {
acc += tensor_Acc2(_dH, make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock) * tensor_Beta(make_coord(_qHLocal, _qL), _kvH, _BS, _kvBlock);
}
tensor_O(_dH, make_coord(_qHLocal, _qL), _kvH, _BS) = converter_o(acc / tensor_Fsum_cluster(make_coord(_qHLocal, _qL), _kvH, _BS));
}
}
}
}
}
/*int example_row = 0;
int example_kvH = 0;
int example_kvBlock = 0;
int example_BS = 3;
print("tensor_Acc1:\t"); print(tensor_Acc1(_, _, example_kvH, example_BS)); print("\n");
print("tensor_Fmax:\t"); print_tensor(tensor_Fmax(example_kvBlock,_,example_kvH,example_BS)); print("\n");
print("tensor_P:\t"); print_tensor(tensor_P(make_coord(example_row, example_kvBlock),_,example_kvH,example_BS)); print("\n");
print("tensor_Fsum:\t"); print_tensor(tensor_Fsum(example_kvBlock,_,example_kvH,example_BS)); print("\n");
print("tensor_P_converted:\t"); print_tensor(tensor_P_converted(make_coord(example_row, example_kvBlock),_,example_kvH,example_BS)); print("\n");
print("tensor_Acc2:\t"); print_tensor(tensor_Acc2(example_row,_,example_kvH,example_BS,example_kvBlock)); print("\n");
print("tensor_Fmax_cluster:\t"); print_tensor(tensor_Fmax_cluster(_,example_kvH,example_BS)); print("\n");
print("tensor_Beta:\t"); print_tensor(tensor_Beta(_,example_kvH,example_BS,_)); print("\n");
print("tensor_Fsum_cluster:\t"); print_tensor(tensor_Fsum_cluster(_,example_kvH,example_BS)); print("\n");
print("tensor_O:\t"); print_tensor(tensor_O(example_row,_,example_kvH,example_BS)); print("\n");*/
}
template <class Tensor>
void
initialize_tensor(Tensor& tensor, cute::tuple<int, int> value_range = {-4, 4}) {
using DataType = typename Tensor::element_type;
auto [min, max] = value_range;
for (int i = 0; i < cute::size(tensor); i++) {
tensor(i) = DataType(int((max-min)*(rand() / double(RAND_MAX)) + min));
}
}
// Compares two CuTe tensors with torch.allclose semantics
// Returns true if: |input_i - other_i| <= atol + rtol x |other_i| for all elements
template <class TensorInput, class TensorOther>
bool
cute_allclose(
TensorInput const& input,
TensorOther const& other,
float rtol = 1e-05f,
float atol = 1e-08f,
bool equal_nan = false) {
using namespace cute;
// Tensors must have the same size
if (size(input) != size(other)) {
std::cerr << "Error: Tensor sizes don't match. input size: " << size(input)
<< ", other size: " << size(other) << std::endl;
return false;
}
int mismatches = 0;
const int max_print = 10; // Only print first 10 mismatches
for (int i = 0; i < size(input); ++i) {
float input_val = float(input(i));
float other_val = float(other(i));
// Handle NaN comparison
bool input_is_nan = std::isnan(input_val);
bool other_is_nan = std::isnan(other_val);
if (input_is_nan || other_is_nan) {
if (equal_nan && input_is_nan && other_is_nan) {
continue; // Both NaN and equal_nan is true
}
else if (input_is_nan || other_is_nan) {
if (mismatches < max_print) {
std::cerr << "Mismatch at index " << i << ": input=" << input_val
<< ", other=" << other_val << " (NaN detected)" << std::endl;
}
mismatches++;
continue;
}
}
// Check torch.allclose condition: |input - other| <= atol + rtol * |other|
float diff = std::abs(input_val - other_val);
float threshold = atol + rtol * std::abs(other_val);
if (diff > threshold) {
if (mismatches < max_print) {
std::cerr << "Mismatch at index " << i << ": input=" << input_val
<< ", other=" << other_val << ", diff=" << diff
<< ", threshold=" << threshold << std::endl;
}
mismatches++;
}
}
if (mismatches > 0) {
std::cerr << "Total mismatches: " << mismatches << " out of " << size(input) << " elements" << std::endl;
return false;
}
return true;
}
struct ProblemStride {
int stride_Q_kvH;
int stride_Q_qHLocal;
int stride_Q_qL;
int stride_Q_dH;
int stride_Q_BS;
int stride_K_kvH;
int stride_K_kvL;
int stride_K_dH;
int stride_K_BS;
int stride_V_kvH;
int stride_V_kvL;
int stride_V_dH;
int stride_V_BS;
int stride_O_kvH;
int stride_O_qHLocal;
int stride_O_qL;
int stride_O_dH;
int stride_O_BS;
};
ProblemStride make_gqa_stride(int kvH, int qHLocal, int qL, int kvL, int dH, int BS) {
ProblemStride stride;
// Q shape ((qHLocal, qL), dH, kvH, BS), where dH is contiguous
// slowest moving dim->fastest moving dim: BS, qL, kvH, qHLocal, dH
stride.stride_Q_kvH = qHLocal * dH;
stride.stride_Q_qHLocal = dH;
stride.stride_Q_qL = kvH * qHLocal * dH;
stride.stride_Q_dH = 1;
stride.stride_Q_BS = kvH * qHLocal * dH * qL;
// K shape (kvL, dH, kvH, BS), where dH is contiguous
// slowest moving dim->fastest moving dim: BS, kvL, kvH, dH
stride.stride_K_kvH = dH;
stride.stride_K_kvL = kvH * dH;
stride.stride_K_dH = 1;
stride.stride_K_BS = kvL * dH * kvH;
// V shape (dH, kvL, kvH, BS), where dH is contiguous
// slowest moving dim->fastest moving dim: BS, kvL, kvH, dH
stride.stride_V_kvH = dH;
stride.stride_V_kvL = kvH * dH;
stride.stride_V_dH = 1;
stride.stride_V_BS = kvL * dH * kvH;
// O shape (dH, (qHLocal, qL), kvH, BS), where dH is contiguous
// slowest moving dim->fastest moving dim: BS, qL, kvH, qHLocal, dH
stride.stride_O_kvH = qHLocal * dH;
stride.stride_O_qHLocal = dH;
stride.stride_O_qL = kvH * qHLocal * dH;
stride.stride_O_dH = 1;
stride.stride_O_BS = kvH * qHLocal * dH * qL;
return stride;
}
class GQATester {
public:
// Kernel config constants
using TypeQKV = cutlass::bfloat16_t;
using TypeO = cutlass::bfloat16_t;
using TypeAcc = float;
static constexpr int CTA_qHLocal = 8;
static constexpr int CTA_qL = 1;
static constexpr int CTA_kvL = 128;
static constexpr int CTA_dH = 64;
static constexpr int BMM1_DMA_Stage = 3;
static constexpr int BMM2_DMA_Stage = 3;
static constexpr int MaxSplits = 8;
static constexpr int NumReductionCTA = 8;
static constexpr bool NoSink = true;
private:
int kvH_, qHLocal_, qL_, kvL_, dH_, BS_;
float softmax_scale_;
ProblemStride stride_;
int sliding_window_size_;
// Host vectors
thrust::host_vector<TypeQKV> host_Q_;
thrust::host_vector<TypeQKV> host_K_;
thrust::host_vector<TypeQKV> host_V_;
thrust::host_vector<TypeO> host_O_;
thrust::host_vector<TypeO> host_reference_O_;
thrust::host_vector<int> host_seq_lens_;
thrust::host_vector<TypeAcc> host_sinks_;
// Device vectors
thrust::device_vector<TypeQKV> device_Q_;
thrust::device_vector<TypeQKV> device_K_;
thrust::device_vector<TypeQKV> device_V_;
thrust::device_vector<TypeO> device_O_;
thrust::device_vector<int> device_seq_lens_;
thrust::device_vector<TypeAcc> device_sinks_;
public:
GQATester(int kvH, int qH, int qL, int kvL, int dH, int BS, float softmax_scale, int sliding_window_size) :
kvH_(kvH), qHLocal_(qH / kvH), qL_(qL), kvL_(kvL), dH_(dH), BS_(BS), softmax_scale_(softmax_scale), sliding_window_size_(sliding_window_size) {
assert(sliding_window_size_ >= 0);
// Allocate host memory
host_Q_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_);
host_K_.resize(kvH_ * kvL_ * dH_ * BS_);
host_V_.resize(kvH_ * kvL_ * dH_ * BS_);
host_O_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_);
host_reference_O_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_);
host_seq_lens_.resize(BS_);
host_sinks_.resize(qHLocal_ * kvH_); // one sink per q head
stride_ = make_gqa_stride(kvH_, qHLocal_, qL_, kvL_, dH_, BS_);
// Create host CuTe tensors for initialization
auto host_tensor_Q = make_tensor(host_Q_.data(), TGV::gqa::make_layout_Q(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_Q_kvH, stride_.stride_Q_qHLocal, stride_.stride_Q_qL, stride_.stride_Q_dH, stride_.stride_Q_BS));
auto host_tensor_K = make_tensor(host_K_.data(), TGV::gqa::make_layout_K(kvH_, kvL_, dH_, BS_, stride_.stride_K_kvH, stride_.stride_K_kvL, stride_.stride_K_dH, stride_.stride_K_BS));
auto host_tensor_V = make_tensor(host_V_.data(), TGV::gqa::make_layout_V(kvH_, kvL_, dH_, BS_, stride_.stride_V_kvH, stride_.stride_V_kvL, stride_.stride_V_dH, stride_.stride_V_BS));
auto host_tensor_sinks = make_tensor(host_sinks_.data(), TGV::gqa::make_layout_sinks(qHLocal_, qL_, kvH_));
// Initialize Q, K, V tensors with random values
initialize_tensor(host_tensor_Q);
initialize_tensor(host_tensor_K);
initialize_tensor(host_tensor_V);
// have batch size matching kvL (i.e. max seq len) for now
bool test_var_seq_lens = false;
for (int i = 0; i < BS_; ++i) {
if (test_var_seq_lens) {
host_seq_lens_[i] = rand() % kvL_ + 1;
}
else { // all the batch have the same seq len
host_seq_lens_[i] = kvL_;
}
}
for (int i = 0; i < qHLocal_ * kvH_; ++i) {
host_sinks_[i] = rand() / (float)RAND_MAX;
}
// Allocate device memory and copy H2D
device_Q_ = host_Q_;
device_K_ = host_K_;
device_V_ = host_V_;
device_O_.resize(kvH_ * qHLocal_ * qL_ * dH_ * BS_);
device_seq_lens_ = host_seq_lens_;
device_sinks_ = host_sinks_;
gpuErrChk(cudaDeviceSynchronize());
}
void run_kernel(bool pdl, int pdl_count = -1, cudaStream_t stream = 0) {
TGV::gqa::gqa_host<
TypeQKV, TypeO, TypeAcc,
CTA_qHLocal, CTA_qL, CTA_kvL, CTA_dH,
BMM1_DMA_Stage, BMM2_DMA_Stage,
MaxSplits,
NumReductionCTA>(
device_K_.data().get(), device_Q_.data().get(), device_V_.data().get(), device_O_.data().get(),
device_seq_lens_.data().get(),
NoSink ? nullptr : device_sinks_.data().get(),
kvH_, qHLocal_, qL_, kvL_, dH_, BS_,
stride_.stride_K_kvH, stride_.stride_K_kvL, stride_.stride_K_dH, stride_.stride_K_BS,
stride_.stride_Q_kvH, stride_.stride_Q_qHLocal, stride_.stride_Q_qL, stride_.stride_Q_dH, stride_.stride_Q_BS,
stride_.stride_V_kvH, stride_.stride_V_kvL, stride_.stride_V_dH, stride_.stride_V_BS,
stride_.stride_O_kvH, stride_.stride_O_qHLocal, stride_.stride_O_qL, stride_.stride_O_dH, stride_.stride_O_BS,
softmax_scale_,
sliding_window_size_,
pdl, pdl_count, stream);
}
bool verify() {
// Run the GPU kernel
run_kernel(false);
gpuErrChk(cudaDeviceSynchronize());
// Copy D2H
host_O_ = device_O_;
// Create tensors for verification using helper methods
auto host_tensor_Q = make_tensor(host_Q_.data(), TGV::gqa::make_layout_Q(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_Q_kvH, stride_.stride_Q_qHLocal, stride_.stride_Q_qL, stride_.stride_Q_dH, stride_.stride_Q_BS));
auto host_tensor_K = make_tensor(host_K_.data(), TGV::gqa::make_layout_K(kvH_, kvL_, dH_, BS_, stride_.stride_K_kvH, stride_.stride_K_kvL, stride_.stride_K_dH, stride_.stride_K_BS));
auto host_tensor_V = make_tensor(host_V_.data(), TGV::gqa::make_layout_V(kvH_, kvL_, dH_, BS_, stride_.stride_V_kvH, stride_.stride_V_kvL, stride_.stride_V_dH, stride_.stride_V_BS));
auto host_tensor_O = make_tensor(host_O_.data(), TGV::gqa::make_layout_O(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_O_kvH, stride_.stride_O_qHLocal, stride_.stride_O_qL, stride_.stride_O_dH, stride_.stride_O_BS));
auto host_reference_tensor_O = make_tensor(host_reference_O_.data(), TGV::gqa::make_layout_O(kvH_, qHLocal_, qL_, dH_, BS_, stride_.stride_O_kvH, stride_.stride_O_qHLocal, stride_.stride_O_qL, stride_.stride_O_dH, stride_.stride_O_BS));
auto host_seq_lens_tensor = make_tensor(host_seq_lens_.data(), make_layout(make_shape(BS_)));
auto host_tensor_sinks = make_tensor(host_sinks_.data(), TGV::gqa::make_layout_sinks(qHLocal_, qL_, kvH_));
gpuErrChk(cudaDeviceSynchronize());
print("host_seq_lens_tensor:\t"); print_tensor(host_seq_lens_tensor);
print("host_tensor_Q:\t"); print(host_tensor_Q); print("\n");
print("host_tensor_K:\t"); print(host_tensor_K); print("\n");
print("host_tensor_V:\t"); print(host_tensor_V); print("\n");
print("host_tensor_O:\t"); print(host_tensor_O); print("\n");
print("host_reference_tensor_O:\t"); print(host_reference_tensor_O); print("\n");
print("host_tensor_sinks:\t"); print(host_tensor_sinks); print("\n");
// Execute reference GQA kernel
reference_gqa<TypeAcc, CTA_kvL, NoSink>(
host_tensor_K, host_tensor_Q, host_tensor_V, host_reference_tensor_O, host_seq_lens_.data(), host_tensor_sinks, softmax_scale_, sliding_window_size_);
// Compare results using torch.allclose semantics
// For bfloat16, use more relaxed tolerances due to reduced precision
bool success = cute_allclose(host_tensor_O, host_reference_tensor_O,
1e-2f, // rtol (relative tolerance)
1e-3f); // atol (absolute tolerance)
//print("host_tensor_O:\t"); print_tensor(host_tensor_O(_,_,0)); print("\n");
//print("host_reference_tensor_O:\t"); print_tensor(host_reference_tensor_O(_,_,0)); print("\n");
std::cout << "Execution is " << ((success) ? "successful." : "failed.") << std::endl;
return success;
}
};
void benchmark_gqa(int kvH, int qH, int qL, int kvL, int dH, int BS, float softmax_scale, int sliding_window_size, bool pdl, int pdl_count, int num_testers = 4, int bench_iters = 100) {
std::cout << "=== GQA Benchmark ===" << std::endl;
std::cout << "Problem size: kvH=" << kvH << ", qH=" << qH << ", qL=" << qL << ", kvL=" << kvL << ", dH=" << dH << ", BS=" << BS << ", sliding_window_size=" << sliding_window_size << std::endl;
std::cout << "Number of testers (L2 thrashing): " << num_testers << std::endl;
std::cout << "Benchmark iterations: " << bench_iters << std::endl;
// Create multiple tester instances to thrash L2 cache
std::vector<std::unique_ptr<GQATester>> testers;
for (int i = 0; i < num_testers; ++i) {
testers.push_back(std::make_unique<GQATester>(kvH, qH, qL, kvL, dH, BS, softmax_scale, sliding_window_size));
}
std::cout << "Created " << num_testers << " GQATester instances" << std::endl;
// Create CUDA stream for graph capture
cudaStream_t stream;
gpuErrChk(cudaStreamCreate(&stream));
// Capture CUDA graph
std::cout << "Capturing CUDA graph..." << std::endl;
cudaGraph_t graph;
cudaGraphExec_t graph_exec;
gpuErrChk(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal));
// Capture round robin execution pattern
for (int iter = 0; iter < bench_iters; ++iter) {
int tester_idx = iter % num_testers;
// Note: We need to run kernels on the same stream for graph capture
// This requires modifying the kernel launch to accept a stream parameter
// For now, we'll capture a simpler pattern and measure accordingly
testers[tester_idx]->run_kernel(pdl, pdl_count, stream);
}
gpuErrChk(cudaStreamEndCapture(stream, &graph));
gpuErrChk(cudaGraphInstantiate(&graph_exec, graph, NULL, NULL, 0));
std::cout << "CUDA graph captured and instantiated" << std::endl;
// Warmup: run kernels in round robin fashion
std::cout << "Starting warmup..." << std::endl;
gpuErrChk(cudaGraphLaunch(graph_exec, stream));
gpuErrChk(cudaDeviceSynchronize());
std::cout << "Warmup completed" << std::endl;
// Benchmark: replay graph and measure time
std::cout << "Starting benchmark..." << std::endl;
// Create CUDA events for timing
cudaEvent_t start_event, stop_event;
gpuErrChk(cudaEventCreate(&start_event));
gpuErrChk(cudaEventCreate(&stop_event));
gpuErrChk(cudaProfilerStart());
gpuErrChk(cudaEventRecord(start_event, stream));
gpuErrChk(cudaGraphLaunch(graph_exec, stream));
gpuErrChk(cudaEventRecord(stop_event, stream));
gpuErrChk(cudaProfilerStop());
gpuErrChk(cudaDeviceSynchronize());
// Calculate timing
float total_time_ms;
gpuErrChk(cudaEventElapsedTime(&total_time_ms, start_event, stop_event));
float avg_time_ms = total_time_ms / bench_iters;
float avg_time_us = avg_time_ms * 1000.0f;
// Calculate FLOPS
long long ops = 2LL * qH * qL * kvL * dH * 2LL * BS; // 2 ops per multiply-add
double gflops = (ops * bench_iters) / (total_time_ms * 1e6);
double gflops_per_iter = ops / (avg_time_ms * 1e6);
// Calculate DRAM bandwidth
long long bytes_per_iter = ((long long)kvH * kvL * dH * 2 * sizeof(GQATester::TypeQKV) + (long long)qH * qL * dH * (sizeof(GQATester::TypeQKV) + sizeof(GQATester::TypeO))) * BS; // Q(bf16) + K(bf16) + V(bf16) + O(bf16)
float avg_dram_bw_gbps = (bytes_per_iter / (avg_time_ms / 1000.0f)) / (1024.0f * 1024.0f * 1024.0f);
// Report results
std::cout << "\n=== Benchmark Results ===" << std::endl;
std::cout << "Average time per iteration: " << avg_time_ms << " ms (" << avg_time_us << " μs)" << std::endl;
std::cout << "GFLOPS per iteration: " << gflops_per_iter << std::endl;
std::cout << "Average DRAM bandwidth: " << avg_dram_bw_gbps << " GB/s" << std::endl;
// Cleanup
gpuErrChk(cudaGraphExecDestroy(graph_exec));
gpuErrChk(cudaGraphDestroy(graph));
gpuErrChk(cudaEventDestroy(start_event));
gpuErrChk(cudaEventDestroy(stop_event));
gpuErrChk(cudaStreamDestroy(stream));
std::cout << "=== Benchmark Complete ===" << std::endl;
}
int main(int argc, char* argv[]) {
srand(time(NULL));
int kvH = 8; // num KV head
int qH = 64; // num Q head
int qL = 1; // Q sequence length
int kvL = 2048; // KV sequence length
int dH = 64; // hidden dimension
int BS = 1; // batch size
float softmax_scale = 1.0f / (float)sqrt(dH);
int sliding_window_size = 0; // when sliding_window_size = 0, it's disabled
bool pdl = false;
// don't support it yet
int pdl_count = -1;
// arg parsing
while (1) {
static struct option long_options[] = {
{"kvL", required_argument, 0, 0},
{"kvH", required_argument, 0, 0},
{"qH", required_argument, 0, 0},
{"qL", required_argument, 0, 0},
{"BS", required_argument, 0, 0},
{"sliding_window_size", required_argument, 0, 0},
{0, 0, 0, 0} // denote end of array
};
int option_index = 0;
// M no argument
// M: required argument
// M:: optional argument
int c = getopt_long(argc, argv, "", long_options, &option_index);
if (c==-1) break;
switch (c) {
case 0:
// Long option
if (option_index == 0) kvL = atoi(optarg);
else if (option_index == 1) kvH = atoi(optarg);
else if (option_index == 2) qH = atoi(optarg);
else if (option_index == 3) qL = atoi(optarg);
else if (option_index == 4) BS = atoi(optarg);
else if (option_index == 5) sliding_window_size = atoi(optarg);
break;
default: assert(false);
}
}
GQATester tester(kvH, qH, qL, kvL, dH, BS, softmax_scale, sliding_window_size);
bool success = tester.verify();
std::cout << "Correctness test " << (success ? "PASSED" : "FAILED") << std::endl;
benchmark_gqa(kvH, qH, qL, kvL, dH, BS, softmax_scale, sliding_window_size, pdl, pdl_count, 100, 1000);
}
examples/93_blackwell_low_latency_gqa/tgv_gqa.cuh
+2215
View File
File diff suppressed because it is too large Load Diff
examples/94_ada_fp8_blockwise/CMakeLists.txt
+32
View File
@@ -0,0 +1,32 @@
# 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.
cutlass_example_add_executable(
94_ada_fp8_blockwise
ada_fp8_blockwise.cu
)
examples/94_ada_fp8_blockwise/ada_fp8_blockwise.cu
+489
View File
@@ -0,0 +1,489 @@
/***************************************************************************************************
* 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;
}
examples/CMakeLists.txt
+4 -1
View File
@@ -1,4 +1,3 @@
# Copyright (c) 2017 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
#
@@ -171,6 +170,10 @@ foreach(EXAMPLE
90_sm103_fp4_ultra_grouped_gemm
91_fp4_gemv
92_blackwell_moe_gemm
93_blackwell_low_latency_gqa
94_ada_fp8_blockwise
111_hopper_ssd
112_blackwell_ssd
)
add_subdirectory(${EXAMPLE})
examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py
+1832
View File
File diff suppressed because it is too large Load Diff
examples/python/CuTeDSL/cute/export/export_to_c.py
+107
View File
@@ -0,0 +1,107 @@
# 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.
import cutlass
import cutlass.cute as cute
import os
import subprocess
import cuda.bindings.driver as cuda
"""Example demonstrating how to export a CuTe function.
This example shows how to:
1. Compile a CuTe function
2. Export the compiled function to a object file and a C header file
3. Compile the object file to a shared library
To run this example:
.. code-block:: bash
# export the compiled function to a object file and a C header file
python cutlass_ir/compiler/python/examples/cute/export/export_to_c.py
"""
@cute.kernel
def print_tensor_kernel(a: cute.Tensor):
cute.printf("a: {}", a)
@cute.jit
def print_tensor(a: cute.Tensor, stream: cuda.CUstream):
print_tensor_kernel(a).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream)
@cute.kernel
def add_one_kernel(a: cute.Tensor, b: cute.Tensor):
a[0] = b[0] + 1
@cute.jit
def add_one(a: cute.Tensor, b: cute.Tensor, stream: cuda.CUstream):
add_one_kernel(a, b).launch(grid=(1, 1, 1), block=(1, 1, 1), stream=stream)
def run():
from cutlass.cute.runtime import make_fake_compact_tensor
shape = (cute.SymInt(divisibility=16), cute.SymInt())
a = make_fake_compact_tensor(cutlass.Float32, shape, stride_order=(1, 0))
b = make_fake_compact_tensor(cutlass.Float32, shape, stride_order=(1, 0))
stream = cuda.CUstream(0)
compiled_print_tensor = cute.compile(print_tensor, a, stream=stream)
compiled_add_one = cute.compile(add_one, a, b, stream=stream)
os.makedirs("./build", exist_ok=True)
compiled_print_tensor.export_to_c(
file_path="./build",
file_name="print_tensor_example",
function_prefix="print_tensor",
)
compiled_add_one.export_to_c(
file_path="./build", file_name="add_one_example", function_prefix="add_one"
)
cc = os.environ.get("CC", "gcc")
# compile the object file to a shared library
cmd = [
cc,
"-shared",
"-o",
"./build/libexport_example.so",
"./build/print_tensor_example.o",
"./build/add_one_example.o",
]
subprocess.run(cmd, check=True)
if __name__ == "__main__":
run()
examples/python/CuTeDSL/cute/export/load_in_python.py
+79
View File
@@ -0,0 +1,79 @@
# 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.
import cutlass
import cutlass.cute as cute
import cuda.bindings.driver as cuda
"""Example demonstrating how to load a CuTe module/function in Python.
This example shows how to:
1. Load a CuTe module from a object file or a shared library
2. Extract the function from the module and call it in Python
To run this example:
.. code-block:: bash
# prerequesites: export the compiled functions to object files and compile them into a shared library
python examples/cute/export/export_to_c.py
# load the module from a object file or a shared library
python examples/cute/export/load_in_python.py
"""
def run():
import torch
from cutlass.cute.runtime import from_dlpack
a = (
torch.arange(16 * 10, dtype=torch.float32, device="cuda")
.reshape(16, 10)
.permute(1, 0)
)
b = torch.ones(16, 10, dtype=torch.float32, device="cuda").permute(1, 0)
a_cute = from_dlpack(a).mark_layout_dynamic()
b_cute = from_dlpack(b).mark_layout_dynamic()
stream = cuda.CUstream(0)
print_tensor_mod = cute.runtime.load_module("./build/print_tensor_example.o")
print_tensor_mod.print_tensor(a_cute, stream=stream)
add_one_mod = cute.runtime.load_module("./build/add_one_example.o")
add_one_mod.add_one(a_cute, b_cute, stream=stream)
assert a[0, 0] == b[0, 0] + 1
shared_mod = cute.runtime.load_module("./build/libexport_example.so")
shared_mod.print_tensor(a_cute, stream=stream)
shared_mod.add_one(a_cute, b_cute, stream=stream)
assert a[0, 0] == b[0, 0] + 1
if __name__ == "__main__":
run()
examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.cpp
+254
View File
@@ -0,0 +1,254 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES.
* All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary
*
* NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
* property and proprietary rights in and to this material, related
* documentation and any modifications thereto. Any use, reproduction,
* disclosure or distribution of this material and related documentation
* without an express license agreement from NVIDIA CORPORATION or
* its affiliates is strictly prohibited.
*/
/**
* This example demonstrates how to load a CuTe module/function from a object
* file or a shared library and run it in a CUDA kernel.
*
* To run this example:
*
* .. code-block:: bash
*
* # prerequesites: export the compiled functions to object files and compile
* them into a shared library python examples/cute/export/export_to_c.py # run
* the example bash ./examples/cute/export/run_with_dynamic_loading.sh
*/
#include "CuteDSLRuntime.h"
#include <cuda_runtime.h>
#include <fstream>
#include <vector>
std::vector<unsigned char> read_file(const std::string &filename) {
std::ifstream file(filename, std::ios::binary);
std::vector<unsigned char> content((std::istreambuf_iterator<char>(file)),
std::istreambuf_iterator<char>());
return content;
}
void initialize_cuda_context() {
// Initialize cuda context
cudaSetDevice(0);
}
void check_error(CuteDSLRT_Error_t error) {
if (error != CuteDSLRT_Error_Success) {
printf("Got runtime error: %d\n", error);
exit(1);
}
}
// Copy the definition of the tensor from `print_tensor_example.h` to here
typedef struct {
void *data;
int32_t dynamic_shapes[2];
int64_t dynamic_strides[1];
} print_tensor_Tensor_a_t;
// Copy the definition of the tensor from `add_one_example.h` to here
typedef struct {
void *data;
int32_t dynamic_shapes[2];
int64_t dynamic_strides[1];
} add_one_Tensor_a_t;
typedef struct {
void *data;
int32_t dynamic_shapes[2];
int64_t dynamic_strides[1];
} add_one_Tensor_b_t;
print_tensor_Tensor_a_t prepare_print_tensor_tensor() {
print_tensor_Tensor_a_t tensor;
tensor.data = NULL;
tensor.dynamic_shapes[0] = 32;
tensor.dynamic_shapes[1] = 16;
tensor.dynamic_strides[0] = 16;
return tensor;
}
add_one_Tensor_a_t prepare_add_one_tensor_a() {
float a_host_ptr[32 * 16];
for (int i = 0; i < 32 * 16; i++) {
a_host_ptr[i] = i;
}
void *a_device_ptr;
cudaMalloc(&a_device_ptr, sizeof(float) * 32 * 16);
cudaMemcpy(a_device_ptr, a_host_ptr, sizeof(float) * 32 * 16,
cudaMemcpyHostToDevice);
add_one_Tensor_a_t tensor;
tensor.data = (void *)a_device_ptr;
tensor.dynamic_shapes[0] = 32;
tensor.dynamic_shapes[1] = 16;
tensor.dynamic_strides[0] = 16;
return tensor;
}
add_one_Tensor_b_t prepare_add_one_tensor_b() {
float b_host_ptr[32 * 16];
for (int i = 0; i < 32 * 16; i++) {
b_host_ptr[i] = 32 * 16 - i;
}
void *b_device_ptr;
cudaMalloc(&b_device_ptr, sizeof(float) * 32 * 16);
cudaMemcpy(b_device_ptr, b_host_ptr, sizeof(float) * 32 * 16,
cudaMemcpyHostToDevice);
add_one_Tensor_b_t tensor;
tensor.data = (void *)b_device_ptr;
tensor.dynamic_shapes[0] = 32;
tensor.dynamic_shapes[1] = 16;
tensor.dynamic_strides[0] = 16;
return tensor;
}
void run_print_tensor_with_object_file() {
// prepare tensor
print_tensor_Tensor_a_t tensor = prepare_print_tensor_tensor();
// prepare stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// load module and function
CuteDSLRT_Module_t *module = nullptr;
std::vector<unsigned char> print_tensor_example_o_bytes =
read_file("build/print_tensor_example.o");
size_t print_tensor_example_o_size = print_tensor_example_o_bytes.size();
CuteDSLRT_Error_t error = CuteDSLRT_Module_Create_From_Bytes(
&module, print_tensor_example_o_bytes.data(), print_tensor_example_o_size,
nullptr, 0);
check_error(error);
CuteDSLRT_Function_t *function = nullptr;
error = CuteDSLRT_Module_Get_Function(&function, module, "print_tensor");
check_error(error);
// run kernel, refer to the wrapper function in `print_tensor_example.h`
int32_t cuda_result;
void *args[3] = {&tensor, &stream, &cuda_result};
error = CuteDSLRT_Function_Run(function, args, 3);
check_error(error);
// synchronize stream
cudaStreamSynchronize(stream);
// unload module
error = CuteDSLRT_Module_Destroy(module);
check_error(error);
cudaStreamDestroy(stream);
}
void run_add_one_with_object_file() {
// prepare a tensor
add_one_Tensor_a_t tensor_a = prepare_add_one_tensor_a();
// prepare b tensor
add_one_Tensor_b_t tensor_b = prepare_add_one_tensor_b();
// prepare stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// load module
CuteDSLRT_Module_t *module = nullptr;
std::vector<unsigned char> add_one_example_o_bytes =
read_file("build/add_one_example.o");
size_t add_one_example_o_size = add_one_example_o_bytes.size();
CuteDSLRT_Error_t error = CuteDSLRT_Module_Create_From_Bytes(
&module, add_one_example_o_bytes.data(), add_one_example_o_size, nullptr,
0);
check_error(error);
CuteDSLRT_Function_t *function = nullptr;
error = CuteDSLRT_Module_Get_Function(&function, module, "add_one");
check_error(error);
// run kernel, refer to the wrapper function in `add_one_example.h`
int32_t cuda_result;
void *args[4] = {&tensor_a, &tensor_b, &stream, &cuda_result};
error = CuteDSLRT_Function_Run(function, args, 4);
check_error(error);
cudaStreamSynchronize(stream);
// unload module
error = CuteDSLRT_Module_Destroy(module);
check_error(error);
cudaStreamDestroy(stream);
// check result
float a_host_ptr[32 * 16];
cudaMemcpy(a_host_ptr, tensor_a.data, sizeof(float) * 32 * 16,
cudaMemcpyDeviceToHost);
if (a_host_ptr[0] != 32 * 16 + 1) {
printf("Error: a_host_ptr[0] = %f, expected %d\n", a_host_ptr[0], 32 * 16);
exit(1);
}
}
void run_example_with_shared_library() {
// prepare tensor
print_tensor_Tensor_a_t tensor = prepare_print_tensor_tensor();
// prepare a tensor
add_one_Tensor_a_t tensor_a = prepare_add_one_tensor_a();
// prepare b tensor
add_one_Tensor_b_t tensor_b = prepare_add_one_tensor_b();
// prepare stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// load module
CuteDSLRT_Module_t *module = nullptr;
const char *shared_libs[] = {"build/libexport_example.so"};
CuteDSLRT_Error_t error =
CuteDSLRT_Module_Create_From_Bytes(&module, nullptr, 0, shared_libs, 1);
check_error(error);
// get print tensor function
CuteDSLRT_Function_t *print_tensor_function = nullptr;
error = CuteDSLRT_Module_Get_Function(&print_tensor_function, module,
"print_tensor");
check_error(error);
// get add one function
CuteDSLRT_Function_t *add_one_function = nullptr;
error = CuteDSLRT_Module_Get_Function(&add_one_function, module, "add_one");
check_error(error);
// run print tensor kernel
int32_t cuda_result;
void *print_tensor_args[3] = {&tensor, &stream, &cuda_result};
error = CuteDSLRT_Function_Run(print_tensor_function, print_tensor_args, 3);
check_error(error);
cudaStreamSynchronize(stream);
// run add one kernel
void *add_one_args[4] = {&tensor_a, &tensor_b, &stream, &cuda_result};
error = CuteDSLRT_Function_Run(add_one_function, add_one_args, 4);
check_error(error);
cudaStreamSynchronize(stream);
// unload module
error = CuteDSLRT_Module_Destroy(module);
check_error(error);
cudaStreamDestroy(stream);
}
int main() {
initialize_cuda_context();
run_print_tensor_with_object_file();
run_add_one_with_object_file();
run_example_with_shared_library();
return 0;
}
examples/python/CuTeDSL/cute/export/run_with_dynamic_loading.sh
Executable
+80
View File
@@ -0,0 +1,80 @@
# 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.
#!/bin/bash
set -eu
# Try to find the wheel path of nvidia-cutlass-dsl
WHEEL_PATH=$(python3 -c "import cutlass, os; print(os.path.dirname(cutlass.__file__))" 2>/dev/null)/../..
if [[ -z "$WHEEL_PATH" ]]; then
echo "nvidia-cutlass-dsl wheel not found in the current Python environment."
exit 1
else
echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH"
fi
CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/"
export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH}:${CUDA_HOME}/lib64:./build
if [ -z "$CUDA_HOME" ]; then
CUDA_HOME=/usr/local/cuda
echo "CUDA_HOME not set, using default: $CUDA_HOME"
else
echo "CUDA_HOME found: $CUDA_HOME"
fi
SOURCE_FILE="$(dirname "$0")/run_with_dynamic_loading.cpp"
echo "Compiling the executable..."
# Search for a common C++ compiler: g++, clang++, or c++
if [ -n "${CXX-}" ] && command -v "$CXX" &> /dev/null; then
CXX="$CXX"
elif command -v g++ &> /dev/null; then
CXX="g++"
elif command -v clang++ &> /dev/null; then
CXX="clang++"
elif command -v c++ &> /dev/null; then
CXX="c++"
else
echo "Error: No common C++ compiler found (g++, clang++, or c++). Please install a C++ compiler to continue."
exit 1
fi
$CXX -o build/run_with_dynamic_loading \
-I${CUDA_HOME}/include \
-I${WHEEL_PATH}/include \
${SOURCE_FILE} \
-L${CUTE_DSL_LIB_PATH} \
-L${CUDA_HOME}/lib64 \
-L./build \
-lcudart \
-lcute_dsl_runtime \
-lexport_example
echo "Running the executable..."
./build/run_with_dynamic_loading
examples/python/CuTeDSL/cute/export/run_with_static_linking.cpp
+124
View File
@@ -0,0 +1,124 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES.
* All rights reserved. SPDX-License-Identifier: LicenseRef-NvidiaProprietary
*
* NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
* property and proprietary rights in and to this material, related
* documentation and any modifications thereto. Any use, reproduction,
* disclosure or distribution of this material and related documentation
* without an express license agreement from NVIDIA CORPORATION or
* its affiliates is strictly prohibited.
*/
/**
* This example demonstrates how to load a CuTe module/function from compilation
* and static linking and run it in a CUDA kernel.
*
* To run this example:
*
* .. code-block:: bash
*
* # prerequesites: export the compiled functions to object files and compile
* them into a shared library python examples/cute/export/export_to_c.py # run
* the example bash ./examples/cute/export/run_with_static_linking.sh
*/
#include "add_one_example.h"
#include "print_tensor_example.h"
#include <cuda_runtime.h>
void initialize_cuda_context() {
// Initialize cuda context
int device_id = 0;
cudaSetDevice(device_id);
}
void run_print_tensor() {
// prepare tensor
print_tensor_Tensor_a_t tensor;
tensor.data = NULL;
tensor.dynamic_shapes[0] = 32;
tensor.dynamic_shapes[1] = 16;
tensor.dynamic_strides[0] = 16;
// prepare stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// load module
print_tensor_Kernel_Module_t module;
print_tensor_Kernel_Module_Load(&module);
// run kernel
cute_dsl_print_tensor_wrapper(&module, &tensor, stream);
// synchronize stream
cudaStreamSynchronize(stream);
// unload module
print_tensor_Kernel_Module_Unload(&module);
cudaStreamDestroy(stream);
}
void run_add_one() {
// prepare a tensor
add_one_Tensor_a_t a_tensor;
float a_host_ptr[32 * 16];
for (int i = 0; i < 32 * 16; i++) {
a_host_ptr[i] = i;
}
void *a_device_ptr;
cudaMalloc(&a_device_ptr, sizeof(float) * 32 * 16);
cudaMemcpy(a_device_ptr, a_host_ptr, sizeof(float) * 32 * 16,
cudaMemcpyHostToDevice);
a_tensor.data = (void *)a_device_ptr;
a_tensor.dynamic_shapes[0] = 32;
a_tensor.dynamic_shapes[1] = 16;
a_tensor.dynamic_strides[0] = 16;
// prepare b tensor
add_one_Tensor_b_t b_tensor;
float b_host_ptr[32 * 16];
for (int i = 0; i < 32 * 16; i++) {
b_host_ptr[i] = 32 * 16 - i;
}
void *b_device_ptr;
cudaMalloc(&b_device_ptr, sizeof(float) * 32 * 16);
cudaMemcpy(b_device_ptr, b_host_ptr, sizeof(float) * 32 * 16,
cudaMemcpyHostToDevice);
b_tensor.data = (void *)b_device_ptr;
b_tensor.dynamic_shapes[0] = 32;
b_tensor.dynamic_shapes[1] = 16;
b_tensor.dynamic_strides[0] = 16;
// prepare stream
cudaStream_t stream;
cudaStreamCreate(&stream);
// load module
add_one_Kernel_Module_t module;
add_one_Kernel_Module_Load(&module);
// run kernel
cute_dsl_add_one_wrapper(&module, &a_tensor, &b_tensor, stream);
cudaStreamSynchronize(stream);
// unload module
add_one_Kernel_Module_Unload(&module);
cudaStreamDestroy(stream);
// check result
cudaMemcpy(a_host_ptr, a_device_ptr, sizeof(float) * 32 * 16,
cudaMemcpyDeviceToHost);
if (a_host_ptr[0] != 32 * 16 + 1) {
printf("Error: a_host_ptr[0] = %f, expected %d\n", a_host_ptr[0], 32 * 16);
}
}
int main() {
initialize_cuda_context();
run_print_tensor();
run_add_one();
return 0;
}
examples/python/CuTeDSL/cute/export/run_with_static_linking.sh
Executable
+76
View File
@@ -0,0 +1,76 @@
# 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.
#!/bin/bash
set -eu
# Try to find the wheel path of nvidia-cutlass-dsl
WHEEL_PATH=$(python3 -c "import cutlass, os; print(os.path.dirname(cutlass.__file__))" 2>/dev/null)/../..
if [[ -z "$WHEEL_PATH" ]]; then
echo "nvidia-cutlass-dsl wheel not found in the current Python environment."
exit 1
else
echo "nvidia-cutlass-dsl wheel path found at: $WHEEL_PATH"
fi
CUTE_DSL_LIB_PATH="${WHEEL_PATH}/lib/"
export LD_LIBRARY_PATH=${CUTE_DSL_LIB_PATH}
if [ -z "$CUDA_HOME" ]; then
CUDA_HOME=/usr/local/cuda
echo "CUDA_HOME not set, using default: $CUDA_HOME"
else
echo "CUDA_HOME found: $CUDA_HOME"
fi
SOURCE_FILE="$(dirname "$0")/run_with_static_linking.cpp"
echo "Compiling the executable..."
# Search for a common C++ compiler: g++, clang++, or c++
if [ -n "${CXX-}" ] && command -v "$CXX" &> /dev/null; then
CXX="$CXX"
elif command -v g++ &> /dev/null; then
CXX="g++"
elif command -v clang++ &> /dev/null; then
CXX="clang++"
elif command -v c++ &> /dev/null; then
CXX="c++"
else
echo "Error: No common C++ compiler found (g++, clang++, or c++). Please install a C++ compiler to continue."
exit 1
fi
# Note: The options -ldl, -lrt, and -lpthread are optional to satisfy symbol dependencies required by `libcudart_static.a`.
$CXX -o build/run_with_static_linking \
-I${CUDA_HOME}/include \
-I./build \
${SOURCE_FILE} build/add_one_example.o build/print_tensor_example.o \
${CUTE_DSL_LIB_PATH}/libcuda_dialect_runtime_static.a \
${CUDA_HOME}/lib64/libcudart_static.a -ldl -lrt -lpthread
echo "Running the executable..."
./build/run_with_static_linking
examples/python/CuTeDSL/cute/tvm_ffi/aot_export.py
+6 -6
View File
@@ -37,21 +37,19 @@ To run this example:
.. code-block:: bash
python examples/cute/tvm_ffi/aot_export.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_export.py
# run example to use in torch
python examples/cute/tvm_ffi/aot_use_in_torch.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_torch.py
# run example to use in jax
python examples/cute/tvm_ffi/aot_use_in_jax.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_jax.py
# run example to use in c++ bundle
bash examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
bash cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
"""
from pathlib import Path
import torch
import os
import subprocess
import tvm_ffi
import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@@ -69,6 +67,8 @@ def add_one(a: cute.Tensor, b: cute.Tensor):
def main():
import torch
# compile the kernel with "--enable-tvm-ffi" option
a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.cpp
+36 -39
View File
@@ -15,32 +15,31 @@
// This example shows how to interface with an AOT compiled function in a C++
// bundle. to build and run the example, run the following command in project
// root bash
// examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
// cutlass_ir/compiler/python/examples/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
#include <cuda_runtime.h>
#include <iostream>
#include <tvm/ffi/container/tensor.h>
#include <tvm/ffi/error.h>
#include <tvm/ffi/extra/module.h>
#include <iostream>
#include <vector>
namespace ffi = tvm::ffi;
struct CUDANDAlloc {
void AllocData(DLTensor *tensor) {
void AllocData(DLTensor* tensor) {
size_t data_size = ffi::GetDataSize(*tensor);
void *ptr = nullptr;
void* ptr = nullptr;
cudaError_t err = cudaMalloc(&ptr, data_size);
TVM_FFI_ICHECK_EQ(err, cudaSuccess)
<< "cudaMalloc failed: " << cudaGetErrorString(err);
TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaMalloc failed: " << cudaGetErrorString(err);
tensor->data = ptr;
}
void FreeData(DLTensor *tensor) {
void FreeData(DLTensor* tensor) {
if (tensor->data != nullptr) {
cudaError_t err = cudaFree(tensor->data);
TVM_FFI_ICHECK_EQ(err, cudaSuccess)
<< "cudaFree failed: " << cudaGetErrorString(err);
TVM_FFI_ICHECK_EQ(err, cudaSuccess) << "cudaFree failed: " << cudaGetErrorString(err);
tensor->data = nullptr;
}
}
@@ -51,47 +50,45 @@ inline ffi::Tensor Empty(ffi::Shape shape, DLDataType dtype, DLDevice device) {
}
// symbol from the shared library
extern "C" int __tvm_ffi_add_one(void *, const TVMFFIAny *, int32_t,
TVMFFIAny *);
extern "C" int __tvm_ffi_add_one(void*, const TVMFFIAny*, int32_t, TVMFFIAny*);
// Redirects into the exported function in object
void CallAddOne(ffi::TensorView x, ffi::TensorView y) {
tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y);
tvm::ffi::Function::InvokeExternC(nullptr, __tvm_ffi_add_one, x, y);
}
int main() {
DLDataType f32_dtype{kDLFloat, 32, 1};
DLDevice cuda_device{kDLCUDA, 0};
DLDataType f32_dtype{kDLFloat, 32, 1};
DLDevice cuda_device{kDLCUDA, 0};
constexpr int ARRAY_SIZE = 10;
constexpr int ARRAY_SIZE = 10;
ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
ffi::Tensor x = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
ffi::Tensor y = Empty({ARRAY_SIZE}, f32_dtype, cuda_device);
std::vector<float> host_x(ARRAY_SIZE);
for (int i = 0; i < ARRAY_SIZE; ++i) {
host_x[i] = static_cast<float>(i);
}
std::vector<float> host_x(ARRAY_SIZE);
for (int i = 0; i < ARRAY_SIZE; ++i) {
host_x[i] = static_cast<float>(i);
}
size_t nbytes = host_x.size() * sizeof(float);
cudaError_t err =
cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice);
TVM_FFI_ICHECK_EQ(err, cudaSuccess)
<< "cudaMemcpy host to device failed: " << cudaGetErrorString(err);
size_t nbytes = host_x.size() * sizeof(float);
cudaError_t err = cudaMemcpy(x.data_ptr(), host_x.data(), nbytes, cudaMemcpyHostToDevice);
TVM_FFI_ICHECK_EQ(err, cudaSuccess)
<< "cudaMemcpy host to device failed: " << cudaGetErrorString(err);
// Call into the FFI function; tensors remain on device because they carry a
// kDLCUDA device tag.
CallAddOne(x, y);
// Call into the FFI function; tensors remain on device because they carry a
// kDLCUDA device tag.
CallAddOne(x, y);
std::vector<float> host_y(host_x.size());
err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost);
TVM_FFI_ICHECK_EQ(err, cudaSuccess)
<< "cudaMemcpy device to host failed: " << cudaGetErrorString(err);
std::vector<float> host_y(host_x.size());
err = cudaMemcpy(host_y.data(), y.data_ptr(), nbytes, cudaMemcpyDeviceToHost);
TVM_FFI_ICHECK_EQ(err, cudaSuccess)
<< "cudaMemcpy device to host failed: " << cudaGetErrorString(err);
std::cout << "y after add_one_cuda(x, y)" << std::endl;
for (float value : host_y) {
std::cout << value << " ";
}
std::cout << std::endl;
return 0;
std::cout << "y after add_one_cuda(x, y)" << std::endl;
for (float value : host_y) {
std::cout << value << " ";
}
std::cout << std::endl;
return 0;
}
examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_cpp_bundle.sh
+6 -6
View File
@@ -27,8 +27,8 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#!/bin/bash
CUDA_DIALECT_PATH="build/lib/"
export LD_LIBRARY_PATH=${CUDA_DIALECT_PATH}:`tvm-ffi-config --libdir`
# Set up library paths for runtime
export LD_LIBRARY_PATH=$(python3 -m cutlass.cute.export.aot_config --libdir):$(tvm-ffi-config --libdir)
CUDA_HOME=/usr/local/cuda
SOURCE_FILE="$(dirname "$0")/aot_use_in_cpp_bundle.cpp"
@@ -38,11 +38,11 @@ g++ -o build/aot_use_in_cpp_bundle \
-I${CUDA_HOME}/include \
`tvm-ffi-config --cxxflags` \
${SOURCE_FILE} build/add_one.o \
-L${CUDA_DIALECT_PATH} \
$(python3 -m cutlass.cute.export.aot_config --ldflags) \
-L${CUDA_HOME}/lib64 \
-lcuda_dialect_runtime -lcuda -lcudart \
`tvm-ffi-config --ldflags` \
`tvm-ffi-config --libs`
$(python3 -m cutlass.cute.export.aot_config --libs) -lcuda -lcudart \
$(tvm-ffi-config --ldflags) \
$(tvm-ffi-config --libs)
echo "Running the executable..."
./build/aot_use_in_cpp_bundle
examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_jax.py
+1 -1
View File
@@ -37,7 +37,7 @@ def main():
a_jax = jnp.arange(10, dtype=jnp.float32)
b_jax = jnp.zeros(10, dtype=jnp.float32)
lib_path = "./build/add_one.so"
aot_mod = cute.runtime.load_module(lib_path)
aot_mod = cute.runtime.load_module(lib_path, enable_tvm_ffi=True)
jax_tvm_ffi.register_ffi_target("add_one_cute", aot_mod.add_one, platform="gpu")
b_jax = jax.ffi.ffi_call(
"add_one_cute",
examples/python/CuTeDSL/cute/tvm_ffi/aot_use_in_torch.py
+4 -2
View File
@@ -27,14 +27,16 @@
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import cutlass.cute as cute
import torch
# now load it back
def main():
import torch
a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
lib_path = "./build/add_one.so"
aot_mod = cute.runtime.load_module(lib_path)
aot_mod = cute.runtime.load_module(lib_path, enable_tvm_ffi=True)
aot_mod.add_one(a_torch, b_torch)
print("result of b after aot_mod.add_one(a, b)")
print(b_torch)
examples/python/CuTeDSL/cute/tvm_ffi/error_reporting.py
+2 -1
View File
@@ -36,8 +36,9 @@ To run this example:
.. code-block:: bash
python examples/cute/tvm_ffi/error_reporting.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/error_reporting.py
"""
import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_jax.py
+1 -1
View File
@@ -39,7 +39,7 @@ To run this example:
pip install jax-tvm-ffi
pip install jax[cuda13]
python examples/cute/tvm_ffi/jit_and_use_in_jax.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/jit_and_use_in_jax.py
"""
import jax
examples/python/CuTeDSL/cute/tvm_ffi/jit_and_use_in_torch.py
+4 -2
View File
@@ -36,9 +36,9 @@ To run this example:
.. code-block:: bash
python examples/cute/tvm_ffi/jit_and_use_in_torch.py
python cutlass_ir/compiler/python/examples/cute/tvm_ffi/jit_and_use_in_torch.py
"""
import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
@@ -56,6 +56,8 @@ def add_one(a: cute.Tensor, b: cute.Tensor):
def main():
import torch
# compile the kernel with "--enable-tvm-ffi" option
a_torch = torch.arange(10, dtype=torch.float32, device="cuda")
b_torch = torch.zeros(10, dtype=torch.float32, device="cuda")
examples/python/CuTeDSL/cute/tvm_ffi/requirements.txt
+1 -1
View File
@@ -1,2 +1,2 @@
apache-tvm-ffi
torch-c-dlpack-ext
torch-c-dlpack-ext
examples/python/CuTeDSL/jax/cutlass_call_basic.py
+260
View File
@@ -0,0 +1,260 @@
# 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.
from functools import partial
import jax
import jax.numpy as jnp
import cutlass
import cutlass.cute as cute
import cutlass.jax as cjax
import cuda.bindings.driver as cuda
"""
Examples of calling CuTe DSL from jax.jit function using cutlass_call.
cutlass_call is a Jax primitive the enables calling of CuTe DSL kernels within a
a jit-compiled Jax function. During the lowering process cutlass_call will
trigger compilation of the kernel and embed it into the HLO computation. It can
then be efficiently launched by XLA without callback to Python.
This example assumes familiarity with CuTe DSL concepts such as layouts and
dynamic shapes.
To run this example:
.. code-block:: bash
# Run with addition operation
python examples/jax/cutlass_call_basic.py
"""
# This is a typical CuTe DSL kernel function that accepts both tensor and scalar values.
@cute.jit
def launch(
A: cute.Tensor,
B: cute.Tensor,
x: cute.Int32,
y: cute.Int32,
C: cute.Tensor,
D: cute.Tensor,
stream: cuda.CUstream,
):
# Print layouts
print("A layout: ", A.layout)
print("B layout: ", B.layout)
print("C layout: ", C.layout)
print("D layout: ", D.layout)
cute.printf("A layout: {}", A.layout)
cute.printf("B layout: {}", B.layout)
cute.printf("C layout: {}", C.layout)
cute.printf("D layout: {}", D.layout)
cute.printf("")
# Print non-tensor values
print("X is: ", x)
print("Y is: ", y)
cute.printf("X is: {}", x)
cute.printf("Y is: {}", y)
print()
# cutlass_call uses a fixed function signature to pass arguments between Jax and CuTeDSL kernel.
#
# Function Signature Requirement:
# stream, inputs, outputs, *, kwargs...
#
# The first argument must be the CUstream that the kernel is run. This stream is managed by the XLA runtime
# and is necessary to schedule and synchronize launches with the rest of your computation.
#
# The second set of arguments are the Jax arrays for inputs and outputs. Inputs must be passed before
# outputs.
#
# Lastly static arguments (i.e. static_argnums or static_argnames) values are passed as keyword only arguments
# by name.
#
# The the kernel does not match this signature a wrapper functions like the one shown below can be written
# or an inline lambda function can be used to rebind the arguments into the appropriate order.
@cute.jit
def launch_jax_wrapper(
stream: cuda.CUstream,
A: cute.Tensor,
B: cute.Tensor,
C: cute.Tensor,
D: cute.Tensor,
*,
x: cute.Int32,
y: cute.Int32,
):
launch(A, B, x, y, C, D, stream)
@cute.jit
def launch_aliased(
A: cute.Tensor, B: cute.Tensor, x: cute.Int32, y: cute.Int32, stream: cuda.CUstream
):
# Print layouts
print("A layout: ", A.layout)
print("B layout: ", B.layout)
cute.printf("A layout: {}", A.layout)
cute.printf("B layout: {}", B.layout)
cute.printf("")
# Print non-tensor values
print("X is: ", x)
print("Y is: ", y)
cute.printf("X is: {}", x)
cute.printf("Y is: {}", y)
print()
if __name__ == "__main__":
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel(a, b, x, y):
call = cjax.cutlass_call(
launch_jax_wrapper,
# Jax requires output shapes/dtype information for each output
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments
x=x,
y=y,
)
# Returned value is a callable to invoke the kernel passing only jax arrays.
return call(a, b)
print("\nExample: example_basic_call_from_jit")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel(A, B, 0, 1)
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_lambda(a, b, x, y):
call = cjax.cutlass_call(
# A lambda function may be used to wrap and bind arguments passed by jax
# to the kernel. Alternatively you can wrap using another separate cute.jit
# function.
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
# Jax requires output shapes/dtype information for each output
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# Static jit arguments are passed via additional keyword arguments
x=x,
y=y,
)
# Returned value is a callable to invoke the kernel passing only jax arrays.
return call(a, b)
print("\nExample: run_cutlass_kernel_lambda")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel_lambda(A, B, 1, 2)
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_static_shapes(a, b, x, y):
call = cjax.cutlass_call(
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# By default cutlass_call will treat all tensors as dynamic shape.
# Dynamic shapes are often expected for kernels so this default ensures
# the broadest support. If you know that a kernel can accept fully static
# tensors then you can enable this flag to pass all tensors shapes and
# layouts known at compile time.
use_static_tensors=True,
x=x,
y=y,
)
return call(a, b)
print("\nExample: run_cutlass_kernel_static_shapes")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel_static_shapes(A, B, 3, 4)
@partial(jax.jit, static_argnums=[2, 3])
def run_cutlass_kernel_with_modes(a, b, x, y):
call = cjax.cutlass_call(
lambda stream, a, b, c, d, *, x, y: launch(a, b, x, y, c, d, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, a.dtype),
),
# The modes of the layout for each tensor can be specified using the
# TensorSpec. By default modes will align with the physical layout
# but can be mapped to specific index position. If None is passed
# then the default mode is assumed for that tensor.
#
# Individual static/dynamic settings may also be applied. For example
# a specific tensor can be marked to have static shape.
input_spec=(
cjax.TensorSpec(mode=(1, 0, 2), static=True),
cjax.TensorSpec(mode=(3, 1, 2, 0)),
),
output_spec=(None, cjax.TensorSpec(mode=(0, 1, 3, 2))),
x=x,
y=y,
)
return call(a, b)
print("\nExample: run_cutlass_kernel_with_modes")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
C, D = run_cutlass_kernel_with_modes(A, B, 5, 6)
@partial(jax.jit, static_argnums=[2, 3], donate_argnums=[0, 1])
def run_cutlass_kernel_aliased_outputs(a, b, x, y):
call = cjax.cutlass_call(
lambda stream, a, b, *, x, y: launch_aliased(a, b, x, y, stream),
output_shape_dtype=(
jax.ShapeDtypeStruct(a.shape, a.dtype),
jax.ShapeDtypeStruct(b.shape, b.dtype),
),
# Can specify the input tensors that are aliasing outputs of this call.
# To avoid allocating separate output buffers. This is useful for kernels
# that update a tensor.
input_output_aliases={0: 0, 1: 1},
x=x,
y=y,
)
return call(a, b)
print("\nExample: run_cutlass_kernel_aliased_outputs")
A = jnp.zeros((512, 32, 64))
B = jnp.zeros((1, 256, 64, 128))
A, B = run_cutlass_kernel_aliased_outputs(A, B, 7, 8)
examples/python/CuTeDSL/jax/cutlass_call_export.py
+168
View File
@@ -0,0 +1,168 @@
# 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.
import pytest
from functools import partial
import argparse
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import jax
import jax.numpy as jnp
from jax import export
from cutlass.jax import cutlass_call, get_export_disabled_safety_checks
from cutlass.jax.testing import create_tensor
"""
Examples of using jax.export APIs with functions using cutlass_call.
This example demonstrates the use of jax.export with CuTe DSL kernel. It assumes
familiarity with CuTe DSL concepts such as layouts and dynamic shapes as well as
Jax's exporting and serialization features:
https://docs.jax.dev/en/latest/export/index.html#export
To run this example:
.. code-block:: bash
# Run with defaults
python examples/jax/cutlass_call_export.py
# Run with shape (1024, 512)
python examples/jax/cutlass_call_export.py --M 1024 --N 512
# Export with symbolic shapes.
python examples/jax/cutlass_call_export.py --export_symbolic
"""
@cute.kernel
def kernel(gA: cute.Tensor, gB: cute.Tensor, gC: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bdim, _, _ = cute.arch.block_dim()
thread_idx = bidx * bdim + tidx
m, n = gA.shape
ni = thread_idx % n
mi = thread_idx // n
a_val = gA[mi, ni]
b_val = gB[mi, ni]
gC[mi, ni] = a_val + b_val
@cute.jit
def launch(stream: cuda.CUstream, mA: cute.Tensor, mB: cute.Tensor, mC: cute.Tensor):
print("mA: ", mA.layout)
print("mB: ", mB.layout)
print("mC: ", mC.layout)
num_threads_per_block = 256
m, n = mA.shape
kernel(mA, mB, mC).launch(
grid=((m * n) // num_threads_per_block, 1, 1),
block=(num_threads_per_block, 1, 1),
stream=stream,
)
def run_example(M, N, export_symbolic_shapes):
@jax.jit
def f(a, b):
call = cutlass_call(launch, output_shape_dtype=a)
return jax.nn.sigmoid(call(a, b))
@jax.jit
def ref_f(a, b):
return jax.nn.sigmoid(a + b)
# Symbolic or partially shapes are supported by cutlass_call and cute.Tensor
# This allows export of functions calling Cut eDSL kernels w/o having to re-compile
# the kernel for each new shape.
if export_symbolic_shapes:
a, b = export.symbolic_shape("a, b")
export_shape_dtype = jax.ShapeDtypeStruct((a, b), jnp.float32)
else:
export_shape_dtype = jax.ShapeDtypeStruct((M, N), jnp.float32)
print("Exporting with input signature: ")
print(f"({export_shape_dtype}, {export_shape_dtype})")
# jax.export can be used to export a jit function containing cutlass_call.
# The function get_export_disabled_safety_checks() returns a list of custom
# call targets that are used by cutlass_call not part of Jax's built-in
# list of stable custom calls.
exported = jax.export.export(f, disabled_checks=get_export_disabled_safety_checks())
traced = exported(export_shape_dtype, export_shape_dtype)
# Serialize the computation to a byte blob.
blob = traced.serialize()
print(f"Serialized computation is {len(blob)} bytes.")
# Deserialize and run
rehydrated = export.deserialize(blob)
key = jax.random.key(1123)
a_key, b_key = jax.random.split(key, 2)
a = create_tensor((M, N), dtype=jnp.float32, key=a_key)
b = create_tensor((M, N), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
# If the computation was exported with dynamic shapes then we can also
# call it with different shapes. The kernel will not be re-compiled
# even though the shapes are changing.
if export_symbolic_shapes:
a = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=a_key)
b = create_tensor((M * 2, N * 4), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
a = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=a_key)
b = create_tensor((M * 4, N * 4), dtype=jnp.float32, key=b_key)
c = rehydrated.call(a, b)
assert jnp.allclose(c, ref_f(a, b))
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Demonstration of using jax.export with functions with cutlass_call"
)
parser.add_argument("--M", default=512, type=int)
parser.add_argument("--N", default=256, type=int)
parser.add_argument("--export_symbolic", action="store_true")
args = parser.parse_args()
run_example(args.M, args.N, args.export_symbolic)
print("PASS")
examples/python/CuTeDSL/jax/cutlass_call_sharding.py
+140
View File
@@ -0,0 +1,140 @@
# 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.
from functools import partial
import argparse
import jax
import jax.numpy as jnp
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from jax.experimental.custom_partitioning import custom_partitioning
import cutlass
import cutlass.cute as cute
import cutlass.jax as cjax
from cutlass.jax.testing import create_tensor
import cuda.bindings.driver as cuda
"""
Examples of combining jax.jit and jax.shard_map for sharding and executing kernels
across multiple GPU devices.
To run this example:
.. code-block:: bash
# Run with addition operation
python examples/jax/cutlass_call_sharding.py
"""
@cute.kernel
def kernel(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
frgA = cute.make_rmem_tensor(cute.size(a, mode=[0]), a.element_type)
frgB = cute.make_rmem_tensor(cute.size(b, mode=[0]), b.element_type)
frgC = cute.make_rmem_tensor(cute.size(c, mode=[0]), c.element_type)
cute.autovec_copy(a[None, tidx, bidx], frgA)
cute.autovec_copy(b[None, tidx, bidx], frgB)
frgC.store(frgA.load() + frgB.load())
cute.autovec_copy(frgC, c[None, tidx, bidx])
@cute.jit
def launch(
stream: cuda.CUstream,
a: cute.Tensor,
b: cute.Tensor,
c: cute.Tensor,
):
cute.printf("a: {}", a.layout)
cute.printf("b: {}", b.layout)
cute.printf("c: {}", c.layout)
kernel(a, b, c).launch(
grid=[a.shape[-1], 1, 1], block=[a.shape[-2], 1, 1], stream=stream
)
def run_example():
# Create a device mesh with one axis b
ngpu = jax.device_count()
mesh = jax.make_mesh((ngpu,), "b")
if ngpu == 1:
print("Note: only 1 GPU was detected.")
# We will shard our 3D tensors over b
sharding = P("b", None, None)
@partial(jax.jit, static_argnums=[0, 1])
def allocate_sharded_tensors(shape, dtype):
key = jax.random.key(1123)
a_key, b_keys = jax.random.split(key, 2)
a = create_tensor(shape, dtype, a_key)
b = create_tensor(shape, dtype, b_keys)
a = jax.lax.with_sharding_constraint(a, NamedSharding(mesh, sharding))
b = jax.lax.with_sharding_constraint(b, NamedSharding(mesh, sharding))
return a, b
@jax.jit
def compute(a, b):
# This jax.shard_map partitions the cutlass_call over the mesh.
@partial(
jax.shard_map,
mesh=mesh,
in_specs=(sharding, sharding),
out_specs=(sharding, sharding),
)
def sharded_call(a_block, b_block):
call = cjax.cutlass_call(
launch,
use_static_tensors=True,
output_shape_dtype=jax.ShapeDtypeStruct(a_block.shape, a_block.dtype),
)
ref_result = a_block + b_block
return call(a_block, b_block), ref_result
return sharded_call(a, b)
# Allocate (32, 16, 64) on each GPU
shape = (32 * ngpu, 16, 64)
dtype = jnp.float32
a, b = allocate_sharded_tensors(shape, dtype)
c, c_ref = compute(a, b)
assert jnp.allclose(c, c_ref)
if __name__ == "__main__":
run_example()
print("PASS")
examples/python/CuTeDSL/jax/elementwise_apply_example.py
+329
View File
@@ -0,0 +1,329 @@
# 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.
import argparse
import operator
from functools import partial
from typing import List, Type
import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
"""
An Elementwise Apply Example using CuTe DSL with cutlass.jax.cutlass_call
This example is similar to examples/ampere/elementwise_apply.py but demonstrates
how to run the code in a jax specific way using the cutlass_call primitive. It assumes
familiarity with basic CuTe DSL concepts as well as the cutlass_call primitive.
To run this example:
.. code-block:: bash
# Run with addition operation
python examples/jax/elementwise_apply_example.py --M 1024 --N 512 --op add
# Run with multiplication operation
python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op mul
# Run with subtraction operation
python examples/ampere/elementwise_apply_example.py --M 1024 --N 512 --op sub
"""
@cute.kernel
def elementwise_apply_kernel(
op: cutlass.Constexpr,
mInputs: List[cute.Tensor],
mC: cute.Tensor,
cC: cute.Tensor, # coordinate tensor
shape: cute.Shape,
tv_layout: cute.Layout, # (tid, vid) -> logic coord
):
tidx, _, _ = cute.arch.thread_idx()
bidx, bidy, _ = cute.arch.block_idx()
###############################################################################
# Slice to local tile of thread block
###############################################################################
blk_crd = ((None, None), (bidx, bidy))
# Leverage the meta-programming capability of the DSL to slice the tensors for each input
# All for loops below on input tensors would be fully unrolled automatically at compile time
# logical coord -> memory address
gInputs = [t[blk_crd] for t in mInputs] # (TileM, TileN)
gC = mC[blk_crd] # (TileM, TileN)
gCrd = cC[blk_crd] # (TileM, TileN)
print("[DSL INFO] Sliced Tensors per thread block:")
for i in cutlass.range_constexpr(len(gInputs)):
print(f"[DSL INFO] ctaInputs{i} = {gInputs[i].type}")
print(f"[DSL INFO] gC = {gC.type}")
print(f"[DSL INFO] gCrd = {gCrd.type}")
###############################################################################
# Compose with thread block TV layout to map thread & value indices to memory address
###############################################################################
# (tid, vid) -> memory address
tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs]
tidfrgC = cute.composition(gC, tv_layout)
tidfrgCrd = cute.composition(gCrd, tv_layout)
# repeat None like vid to remove hierarchy of layout
thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1]))
###############################################################################
# Slice to local tile of thread
###############################################################################
# vid -> address
thrInputs = [t[thr_crd] for t in tidfrgInputs] # (V)
thrC = tidfrgC[thr_crd] # (V)
thrCrd = tidfrgCrd[thr_crd]
print("[DSL INFO] Sliced Tensors per thread:")
for i in cutlass.range_constexpr(len(thrInputs)):
print(f"[DSL INFO] thrInputs{i} = {thrInputs[i].type}")
print(f"[DSL INFO] thrC = {thrC.type}")
print(f"[DSL INFO] thrCrd = {thrCrd.type}")
###############################################################################
# Compute predicate for out of boundary checks
###############################################################################
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
print(f"[DSL INFO] frgPred = {frgPred.type}")
for i in cutlass.range_constexpr(cute.size(frgPred)):
frgPred[i] = cute.elem_less(thrCrd[i], shape)
# if tidx == 0 and bidx == 0:
# cute.print_tensor(frgPred)
##########################################################
# Load data and compute result
##########################################################
# Load data before use. The compiler will optimize the copy and load
# operations to convert some memory ld/st into register uses.
result = op(*[thrInput.load() for thrInput in thrInputs])
thrC.store(result)
@cute.jit
def elementwise_apply(
op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream
):
"""CUDA kernel applying binary operator on each element of two n-D input tensors in
CuTe Python and store to result tensor.
:param op: Binary operator or lambda function to apply element-wise
:type op: cutlass.Constexpr
:param a: First input tensor
:type a: cute.Tensor
:param b: Second input tensor
:type b: cute.Tensor
:param result: Output tensor to store the results of op(a, b)
:type result: cute.Tensor
:return: None
:rtype: None
"""
# Baseline: naive TV layout
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (512, 4) tile
# * tidx maps to mode-0 but input layout is contiguous on mode-1, performance will be bad
# tv_layout = cute.make_layout((128, (4, 4)), stride=(4, (512, 1)))
# cta_tiler = (512, 4)
# Opt-1: better TV layout with better 1D thread layout (SOL with 1D thread layout)
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (4, 512) tile
# * tidx maps to mode-1 which is leading mode of input tensor for coalesced load
# tv_layout = cute.make_layout((128, (4, 4)), stride=(16, (4, 1)))
# cta_tiler = (4, 512)
# Opt-2: 2D tile but worse
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (128, 16) logical tile
# * V layout is bad as contiguous mode is not on right-most
# * `cute.copy` only supports vectorize when stride-1 of v-layout on right-most )
# tv_layout = cute.make_layout(((32, 4), (4, 4)), stride=((4, 512), (1, 128)))
# cta_tiler = (128, 16)
# Opt-3: SOL with 2D thread tile
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (64, 256) logical tile
# * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store
# Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type
coalesced_ldst_bytes = 16
# Compile time validation: expect same element type for all input tensors
assert all(t.element_type == inputs[0].element_type for t in inputs)
dtype = inputs[0].element_type
thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
val_layout = cute.recast_layout(dtype.width, 8, val_layout)
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
print("[DSL INFO] Input Tensors:")
for i, t in enumerate(inputs):
print(f"[DSL INFO] inputs{i} = {t}")
print(f"[DSL INFO] result = {result}")
print("[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_layout}")
print("[DSL INFO] Tiled Tensors:")
mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs]
# ((TileM, TileN), (RestM, RestN))
mC = cute.zipped_divide(result, tiler_mn)
# (RestM, RestN) -> (RestN, RestM)
remap_block = cute.make_ordered_layout(
cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0)
)
for i, t in enumerate(mInputs):
print(f"[DSL INFO] gInputs{i} = {mInputs[i]}")
mInputs[i] = cute.composition(t, (None, remap_block))
print(f"[DSL INFO] gInputs{i} (remapped) = {mInputs[i]}")
mC = cute.composition(mC, (None, remap_block))
print(f"[DSL INFO] gC = {mC}")
idC = cute.make_identity_tensor(result.shape)
cC = cute.zipped_divide(idC, tiler=tiler_mn)
print(f"[DSL INFO] coord tensor = {cC}")
# Launch the kernel asynchronously
# Group input tensors into a list as a single argument
elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch(
# Compute production at each mode of mC.shape[1] to get multi-dimensional grid size
grid=cute.product_each(mC.shape[1]),
block=[cute.size(tv_layout, mode=[0]), 1, 1],
stream=stream,
)
@cutlass.dsl_user_op
def leaky_relu(x, alpha, *, loc=None, ip=None):
return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip)
def leaky_relu_ref(x, alpha):
import jax.numpy as jnp
return jnp.where(x > 0, x, alpha * x)
def run_and_verify(op, M, N, dtype, skip_ref_check=False):
import jax
import jax.numpy as jnp
import cutlass.jax as cjax
import cutlass.jax.testing as testing
if op == "leaky_relu":
op = partial(leaky_relu, alpha=0.01)
ref_op = partial(leaky_relu_ref, alpha=0.01)
num_inputs = 1
else:
op = getattr(operator, op)
ref_op = op
num_inputs = 2
# This jax function is transformed using jax.jit to compile its contents
# into an efficient HLO executable.
@partial(jax.jit, static_argnums=[1])
def jax_function(inputs, op):
call = cjax.cutlass_call(
# Bind jax arguments to kernel signature
lambda stream, inputs, output, *, op: elementwise_apply(
op, inputs, output, stream
),
# Specify output shape/dtype of result
output_shape_dtype=jax.ShapeDtypeStruct(inputs[0].shape, inputs[0].dtype),
# Pass static/constexpr values as kwargs
op=op,
)
# Call the kernel!
return call(inputs)
@partial(jax.jit, static_argnums=[1])
def jax_ref_function(inputs, op):
return op(*inputs)
print("\nRunning Elementwise Apply test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Input and Output Data type: {dtype}")
jax_dtype = cjax.cutlass_to_jax_dtype(dtype)
keys = jax.random.split(jax.random.key(1435), num_inputs)
inputs = [testing.create_tensor((M, N), jax_dtype, key) for key in keys]
print("Input tensor shapes:")
for i in range(num_inputs):
print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}")
epsilon = 1.2
if op in (operator.truediv, operator.floordiv):
inputs[1] = jnp.where(inputs[1] == 0, epsilon, inputs[1])
# Call the jax.jit function which will compile the kernel
c = jax_function(inputs, op)
if not skip_ref_check:
print("Executing elementwise apply kernel...")
c = jax_function(inputs, op)
print("Verifying results...")
assert jnp.allclose(ref_op(*inputs), c)
print("Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3]}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Demonstration of calling a kernel with cutlass_call"
)
parser.add_argument("--M", default=4096, type=int)
parser.add_argument("--N", default=4096, type=int)
parser.add_argument("--op", default="add", type=str)
parser.add_argument("--skip_ref_check", action="store_true")
args = parser.parse_args()
run_and_verify(
args.op,
args.M,
args.N,
dtype=cutlass.Float32,
skip_ref_check=args.skip_ref_check,
)
print("\nPASS")