leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 25 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