Automated PR - 2026-07-07
This commit is contained in:
github-actions[bot]
@@ -0,0 +1,89 @@
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
#include <c10/cuda/CUDAGuard.h>
|
||||
#include <torch/extension.h>
|
||||
#include <vector>
|
||||
#include <stdio.h>
|
||||
|
||||
#ifdef __SM90__
|
||||
#include "sm90_fp8_gemm_1d2d_bias.hpp"
|
||||
#endif
|
||||
|
||||
#include "sm89_fp8_gemm_1d2d.hpp"
|
||||
|
||||
namespace blockwise{
|
||||
template <int N>
|
||||
static auto get_shape(const torch::Tensor& t) {
|
||||
return [&t] <size_t... Is> (std::index_sequence<Is...>) {
|
||||
return std::make_tuple(static_cast<int>(t.sizes()[Is])...);
|
||||
}(std::make_index_sequence<N>());
|
||||
}
|
||||
|
||||
#ifdef __SM90__
|
||||
static void fp8_gemm_nt_sm90(const std::pair<torch::Tensor, torch::Tensor>& a,
|
||||
const std::pair<torch::Tensor, torch::Tensor>& b,
|
||||
const torch::Tensor& d,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const std::optional<torch::Tensor>& c, const int num_sms) {
|
||||
|
||||
// Type and shape checks
|
||||
const auto& [m , k ] = get_shape<2>(a.first);
|
||||
const auto& [n , k_] = get_shape<2>(b.first);
|
||||
const auto& [m_, n_] = get_shape<2>(d);
|
||||
|
||||
// The SM90 kernel always adds bias; synthesize a zero bias when the layer is
|
||||
// bias-less (e.g. the no-bias video FFN of v3 checkpoints), mirroring SM89 below.
|
||||
torch::Tensor bias_tensor = bias.has_value()
|
||||
? bias.value()
|
||||
: torch::zeros({n}, d.options().dtype(torch::kFloat32));
|
||||
sm90_fp8_gemm_1d2d_bias(a.first, a.second, b.first, b.second, bias_tensor, c, d, m, n, k, num_sms);
|
||||
}
|
||||
#endif
|
||||
|
||||
static void fp8_gemm_nt_sm89(const std::pair<torch::Tensor, torch::Tensor>& a,
|
||||
const std::pair<torch::Tensor, torch::Tensor>& b,
|
||||
const torch::Tensor& d,
|
||||
const std::optional<torch::Tensor>& bias,
|
||||
const bool use_fast_accum = true) {
|
||||
|
||||
const auto& [m, k] = get_shape<2>(a.first);
|
||||
const auto& [n, k_] = get_shape<2>(b.first);
|
||||
const auto& [m_, n_] = get_shape<2>(d);
|
||||
|
||||
// The SM89 kernel always adds bias; synthesize a zero bias when the layer is
|
||||
// bias-less so we add 0 rather than uninitialized memory (mirrors SM90 above).
|
||||
torch::Tensor bias_tensor = bias.has_value()
|
||||
? bias.value()
|
||||
: torch::zeros({n}, d.options().dtype(torch::kFloat32));
|
||||
|
||||
blockwise::sm89_fp8_gemm_1d2d_bias(
|
||||
a.first, a.second, // a data, sfa scales
|
||||
b.first, b.second, // b data, sfb scales
|
||||
bias_tensor, // bias (or empty tensor)
|
||||
d, // output
|
||||
m, n, k,
|
||||
use_fast_accum); // pass through accumulation mode
|
||||
}
|
||||
|
||||
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
||||
// m.def("package_name", &function_name, "function_docstring"")
|
||||
#ifdef __SM90__
|
||||
m.def("fp8_gemm_nt_sm90", &fp8_gemm_nt_sm90,
|
||||
py::arg("a"), py::arg("b"), py::arg("d"),
|
||||
py::arg("bias") = std::nullopt,
|
||||
py::arg("c") = std::nullopt,
|
||||
py::arg("num_sms") = 132
|
||||
);
|
||||
#endif
|
||||
m.def("fp8_gemm_nt_sm89", &fp8_gemm_nt_sm89,
|
||||
py::arg("a"), py::arg("b"), py::arg("d"),
|
||||
py::arg("bias") = std::nullopt,
|
||||
py::arg("use_fast_accum") = true
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,92 @@
|
||||
#pragma once
|
||||
#include <torch/python.h>
|
||||
#include <cute/arch/mma_sm100_umma.hpp>
|
||||
#include "utils.hpp"
|
||||
#include "exceptions.hpp"
|
||||
|
||||
namespace blockwise{
|
||||
struct MulticastConfig {
|
||||
int num_multicast;
|
||||
bool is_multicast_on_a;
|
||||
|
||||
MulticastConfig(const int& num_multicast, const bool& is_multicast_on_a):
|
||||
num_multicast(num_multicast), is_multicast_on_a(is_multicast_on_a) {
|
||||
DG_HOST_ASSERT(1 <= num_multicast and num_multicast <= 2);
|
||||
}
|
||||
};
|
||||
|
||||
struct SharedMemoryConfig {
|
||||
int smem_size;
|
||||
int swizzle_a_mode;
|
||||
int swizzle_b_mode;
|
||||
int swizzle_cd_mode;
|
||||
};
|
||||
|
||||
struct ThreadConfig {
|
||||
int num_threads;
|
||||
|
||||
// SM90
|
||||
int num_tma_threads;
|
||||
int num_math_threads;
|
||||
|
||||
// SM100
|
||||
int num_non_epilogue_threads;
|
||||
int num_epilogue_threads;
|
||||
|
||||
static ThreadConfig sm90(const int& num_tma_threads,
|
||||
const int& num_math_threads) {
|
||||
auto config = ThreadConfig();
|
||||
config.num_threads = num_tma_threads + num_math_threads;
|
||||
config.num_tma_threads = num_tma_threads;
|
||||
config.num_math_threads = num_math_threads;
|
||||
return config;
|
||||
}
|
||||
|
||||
static ThreadConfig sm100(const int& num_non_epilogue_threads,
|
||||
const int& num_epilogue_threads) {
|
||||
auto config = ThreadConfig();
|
||||
config.num_threads = num_non_epilogue_threads + num_epilogue_threads;
|
||||
config.num_non_epilogue_threads = num_non_epilogue_threads;
|
||||
config.num_epilogue_threads = num_epilogue_threads;
|
||||
return config;
|
||||
}
|
||||
};
|
||||
|
||||
template<int SM>
|
||||
struct GemmConfig{};
|
||||
// {
|
||||
// // Templated configs
|
||||
|
||||
// at::ScalarType ab_dtype, cd_dtype;
|
||||
// bool with_accumulation;
|
||||
// int block_m, block_n, block_k;
|
||||
// int num_stages, num_last_stages;
|
||||
|
||||
// // Templated device configs
|
||||
// int num_sms;
|
||||
|
||||
// // Structured configs
|
||||
// MulticastConfig multicast_config;
|
||||
// SharedMemoryConfig smem_config;
|
||||
// ThreadConfig thread_config;
|
||||
// };
|
||||
|
||||
|
||||
template <>
|
||||
struct GemmConfig<90>
|
||||
{
|
||||
at::ScalarType ab_dtype = torch::kFloat8_e4m3fn;
|
||||
at::ScalarType cd_dtype = torch::kBFloat16;
|
||||
bool with_accumulation = false;
|
||||
int block_m = 256;
|
||||
int block_n = 128;
|
||||
int block_k = 128;
|
||||
int num_stages = 3;
|
||||
int num_last_stages = 2;
|
||||
int num_sms = 132;
|
||||
MulticastConfig multicast_config{2, true};
|
||||
SharedMemoryConfig smem_config{216240, 128, 128, 128};
|
||||
ThreadConfig thread_config = ThreadConfig::sm90(128, 256);
|
||||
};
|
||||
|
||||
};
|
||||
@@ -0,0 +1,65 @@
|
||||
#pragma once
|
||||
|
||||
#include <exception>
|
||||
#include <string>
|
||||
#include <sstream>
|
||||
|
||||
namespace blockwise {
|
||||
|
||||
class DGException final : public std::exception {
|
||||
std::string message = {};
|
||||
|
||||
public:
|
||||
explicit DGException(const char *name, const char* file, const int line, const std::string& error) {
|
||||
message = std::string(name) + " error (" + file + ":" + std::to_string(line) + "): " + error;
|
||||
}
|
||||
|
||||
const char *what() const noexcept override {
|
||||
return message.c_str();
|
||||
}
|
||||
};
|
||||
|
||||
#ifndef DG_STATIC_ASSERT
|
||||
#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__)
|
||||
#endif
|
||||
|
||||
#ifndef DG_HOST_ASSERT
|
||||
#define DG_HOST_ASSERT(cond) \
|
||||
do { \
|
||||
if (not (cond)) { \
|
||||
throw DGException("Assertion", __FILE__, __LINE__, #cond); \
|
||||
} \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#ifndef DG_HOST_UNREACHABLE
|
||||
#define DG_HOST_UNREACHABLE(reason) (throw DGException("Assertion", __FILE__, __LINE__, reason))
|
||||
#endif
|
||||
|
||||
// #ifndef DG_CUDA_DRIVER_CHECK
|
||||
// #define DG_CUDA_DRIVER_CHECK(cmd) \
|
||||
// do { \
|
||||
// const auto& e = (cmd); \
|
||||
// if (e != CUDA_SUCCESS) { \
|
||||
// std::stringstream ss; \
|
||||
// const char *name, *info; \
|
||||
// cuGetErrorName(e, &name), cuGetErrorString(e, &info); \
|
||||
// ss << static_cast<int>(e) << " (" << name << ", " << info << ")"; \
|
||||
// throw DGException("CUDA driver", __FILE__, __LINE__, ss.str()); \
|
||||
// } \
|
||||
// } while (0)
|
||||
// #endif
|
||||
|
||||
#ifndef DG_CUDA_RUNTIME_CHECK
|
||||
#define DG_CUDA_RUNTIME_CHECK(cmd) \
|
||||
do { \
|
||||
const auto& e = (cmd); \
|
||||
if (e != cudaSuccess) { \
|
||||
std::stringstream ss; \
|
||||
ss << static_cast<int>(e) << " (" << cudaGetErrorName(e) << ", " << cudaGetErrorString(e) << ")"; \
|
||||
throw DGException("CUDA runtime", __FILE__, __LINE__, ss.str()); \
|
||||
} \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
} // namespace deep_gemm
|
||||
+48
@@ -0,0 +1,48 @@
|
||||
#pragma once
|
||||
|
||||
namespace cute {
|
||||
|
||||
struct ignore_t {
|
||||
template <typename T>
|
||||
constexpr const ignore_t& operator=(T&&) const noexcept {
|
||||
return *this;
|
||||
}
|
||||
};
|
||||
|
||||
inline constexpr ignore_t ignore{};
|
||||
|
||||
} // namespace cute
|
||||
|
||||
#define CUTE_TIE_CONCAT_IMPL(A, B) A##B
|
||||
#define CUTE_TIE_CONCAT(A, B) CUTE_TIE_CONCAT_IMPL(A, B)
|
||||
|
||||
#define CUTE_TIE_GET_NTH_ARG(_1, _2, _3, _4, _5, _6, _7, _8, _9, _10, N, ...) N
|
||||
#define CUTE_TIE_COUNT_ARGS(...) \
|
||||
CUTE_TIE_GET_NTH_ARG(__VA_ARGS__, 10, 9, 8, 7, 6, 5, 4, 3, 2, 1, 0)
|
||||
|
||||
#define CUTE_TIE_OP_DECL(I, TUPLE, VAR) auto VAR = ::cute::get<I>(TUPLE)
|
||||
#define CUTE_TIE_OP_ASSIGN(I, TUPLE, VAR) VAR = ::cute::get<I>(TUPLE)
|
||||
|
||||
#define CUTE_TIE_APPLY_OP_1(OP, T, V1) OP(0, T, V1);
|
||||
#define CUTE_TIE_APPLY_OP_2(OP, T, V1, V2) OP(0, T, V1); OP(1, T, V2);
|
||||
#define CUTE_TIE_APPLY_OP_3(OP, T, V1, V2, V3) OP(0, T, V1); OP(1, T, V2); OP(2, T, V3);
|
||||
#define CUTE_TIE_APPLY_OP_4(OP, T, V1, V2, V3, V4) OP(0, T, V1); OP(1, T, V2); OP(2, T, V3); OP(3, T, V4);
|
||||
#define CUTE_TIE_APPLY_OP_5(OP, T, V1, V2, V3, V4, V5) OP(0, T, V1); OP(1, T, V2); OP(2, T, V3); OP(3, T, V4); OP(4, T, V5);
|
||||
|
||||
#define CUTE_TIE_DECL(TUPLE_EXPR, ...) \
|
||||
auto&& CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__) = (TUPLE_EXPR); \
|
||||
CUTE_TIE_CONCAT(CUTE_TIE_APPLY_OP_, CUTE_TIE_COUNT_ARGS(__VA_ARGS__)) ( \
|
||||
CUTE_TIE_OP_DECL, \
|
||||
CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__), \
|
||||
__VA_ARGS__ \
|
||||
)
|
||||
|
||||
#define CUTE_TIE(TUPLE_EXPR, ...) \
|
||||
do { \
|
||||
auto&& CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__) = (TUPLE_EXPR); \
|
||||
CUTE_TIE_CONCAT(CUTE_TIE_APPLY_OP_, CUTE_TIE_COUNT_ARGS(__VA_ARGS__)) ( \
|
||||
CUTE_TIE_OP_ASSIGN, \
|
||||
CUTE_TIE_CONCAT(cute_tie__temp_tuple_, __LINE__), \
|
||||
__VA_ARGS__ \
|
||||
); \
|
||||
} while (0)
|
||||
+27
@@ -0,0 +1,27 @@
|
||||
#pragma once
|
||||
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
struct EpilogueIdentity {
|
||||
template <uint32_t STORE_BLOCK_N>
|
||||
__device__ __forceinline__ static uint32_t apply_index_n(const uint32_t &n_idx) {
|
||||
return n_idx;
|
||||
}
|
||||
};
|
||||
|
||||
template <uint32_t kLeft, uint32_t kMid, uint32_t kRight>
|
||||
struct EpilogueHeadSplits: EpilogueIdentity {
|
||||
template <uint32_t STORE_BLOCK_N>
|
||||
__device__ __forceinline__ static uint32_t apply_index_n(const uint32_t &n_idx) {
|
||||
DG_STATIC_ASSERT(kLeft % STORE_BLOCK_N == 0 and kMid % STORE_BLOCK_N == 0
|
||||
and kRight % STORE_BLOCK_N == 0, "Invalid head splits config");
|
||||
return n_idx + (n_idx + kRight) / (kLeft + kRight) * kMid;
|
||||
}
|
||||
};
|
||||
|
||||
#pragma clang diagnostic pop
|
||||
|
||||
} // namespace deep_gemm
|
||||
+44
@@ -0,0 +1,44 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/utility>
|
||||
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
|
||||
// Operation functors
|
||||
template <typename T> struct ReduceSum { __device__ T operator()(T a, T b) const { return a + b; } };
|
||||
template <typename T> struct ReduceMax { __device__ T operator()(T a, T b) const { return a > b ? a : b; } };
|
||||
template <typename T> struct ReduceMin { __device__ T operator()(T a, T b) const { return a < b ? a : b; } };
|
||||
template <typename T> struct ReduceAnd { __device__ T operator()(T a, T b) const { return a & b; } };
|
||||
template <typename T> struct ReduceOr { __device__ T operator()(T a, T b) const { return a | b; } };
|
||||
|
||||
// Unified reduction function
|
||||
template <int kNumLanesPerGroup, bool kIntergroupReduce, typename T, typename Op>
|
||||
__forceinline__ __device__ T warp_reduce(T value, Op op) {
|
||||
DG_STATIC_ASSERT(kNumLanesPerGroup == 32 or kNumLanesPerGroup == 16 or kNumLanesPerGroup == 8 or
|
||||
kNumLanesPerGroup == 4 or kNumLanesPerGroup == 2 or kNumLanesPerGroup == 1,
|
||||
"Invalid number of lanes");
|
||||
constexpr uint32_t mask = 0xffffffff;
|
||||
if constexpr (kIntergroupReduce) {
|
||||
if constexpr (kNumLanesPerGroup <= 1) value = op(value, __shfl_xor_sync(mask, value, 1));
|
||||
if constexpr (kNumLanesPerGroup <= 2) value = op(value, __shfl_xor_sync(mask, value, 2));
|
||||
if constexpr (kNumLanesPerGroup <= 4) value = op(value, __shfl_xor_sync(mask, value, 4));
|
||||
if constexpr (kNumLanesPerGroup <= 8) value = op(value, __shfl_xor_sync(mask, value, 8));
|
||||
if constexpr (kNumLanesPerGroup <= 16) value = op(value, __shfl_xor_sync(mask, value, 16));
|
||||
} else {
|
||||
if constexpr (kNumLanesPerGroup >= 32) value = op(value, __shfl_xor_sync(mask, value, 16));
|
||||
if constexpr (kNumLanesPerGroup >= 16) value = op(value, __shfl_xor_sync(mask, value, 8));
|
||||
if constexpr (kNumLanesPerGroup >= 8) value = op(value, __shfl_xor_sync(mask, value, 4));
|
||||
if constexpr (kNumLanesPerGroup >= 4) value = op(value, __shfl_xor_sync(mask, value, 2));
|
||||
if constexpr (kNumLanesPerGroup >= 2) value = op(value, __shfl_xor_sync(mask, value, 1));
|
||||
}
|
||||
return value;
|
||||
}
|
||||
|
||||
// Convenience aliases
|
||||
template <int kNumLanesPerGroup = 32, bool kIntergroupReduce = false, typename T>
|
||||
__forceinline__ __device__ T warp_reduce_sum(T value) {
|
||||
return warp_reduce<kNumLanesPerGroup, kIntergroupReduce, T>(value, ReduceSum<T>{});
|
||||
}
|
||||
+239
@@ -0,0 +1,239 @@
|
||||
#pragma once
|
||||
|
||||
#include <deep_gemm/common/types.hpp>
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
enum class KGroupedIndexType {
|
||||
MN,
|
||||
K,
|
||||
SF_K,
|
||||
};
|
||||
|
||||
template <GemmType kGemmType, uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t kNumSMs, bool kIsMulticastOnA>
|
||||
static constexpr uint32_t get_num_1d_blocks_per_group() {
|
||||
// Select the best from candidates
|
||||
uint32_t num_best_blocks = 0, min_usage = cute::numeric_limits<uint32_t>::max();
|
||||
for (const auto& candidate: {8u, 16u}) {
|
||||
const auto& usage = kIsMulticastOnA ?
|
||||
candidate * BLOCK_N + constexpr_ceil_div(kNumSMs, candidate) * BLOCK_M: // Grouping on N
|
||||
candidate * BLOCK_M + constexpr_ceil_div(kNumSMs, candidate) * BLOCK_N; // Grouping on M
|
||||
if (usage < min_usage)
|
||||
min_usage = usage, num_best_blocks = candidate;
|
||||
}
|
||||
return num_best_blocks;
|
||||
}
|
||||
|
||||
#pragma clang diagnostic push
|
||||
#pragma ide diagnostic ignored "cppcoreguidelines-pro-type-member-init"
|
||||
template <GemmType kGemmType,
|
||||
uint32_t BLOCK_M, uint32_t BLOCK_N,
|
||||
uint32_t kNumGroups,
|
||||
uint32_t kNumMulticast, bool kIsMulticastOnA,
|
||||
uint32_t kNumSMs,
|
||||
uint32_t SF_K_ALIGNMENT = 512u, // for k-grouped GEMM only: 128 (SM90 float SF) or 512 (SM100 UE8M0 SF)
|
||||
uint32_t kNum1DBlocksPerGroup = get_num_1d_blocks_per_group<kGemmType, BLOCK_M, BLOCK_N, kNumSMs, kIsMulticastOnA>()>
|
||||
struct Scheduler {
|
||||
int current_iter = -1;
|
||||
|
||||
// Block configs
|
||||
uint32_t num_blocks;
|
||||
uint32_t num_m_blocks;
|
||||
uint32_t num_n_blocks;
|
||||
|
||||
// For SM90 multicast checks
|
||||
uint32_t num_blocks_in_group;
|
||||
bool is_peer_cta_alive = true;
|
||||
|
||||
// For grouped GEMM
|
||||
int* grouped_layout;
|
||||
uint32_t current_group_idx = 0;
|
||||
// Only used for masked layout
|
||||
uint32_t current_m_cumsum = 0;
|
||||
// Only used for k-grouped layout
|
||||
uint32_t current_shape_k, current_num_valid_groups = 0, current_k_cumsum = 0, current_sf_k_cumsum = 0;
|
||||
uint32_t next_group_idx, next_shape_k;
|
||||
|
||||
// Only used for k-grouped gemm
|
||||
__device__ __forceinline__ void get_next_k_group(uint32_t &group_idx, uint32_t &shape_k) const {
|
||||
for (; group_idx < kNumGroups; ++ group_idx) {
|
||||
shape_k = __ldg(grouped_layout + group_idx);
|
||||
if (shape_k > 0)
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
// ReSharper disable once CppPossiblyUninitializedMember
|
||||
__device__ __forceinline__ explicit Scheduler(const uint32_t& shape_m, const uint32_t& shape_n, const uint32_t& shape_k,
|
||||
int* grouped_layout = nullptr) {
|
||||
num_m_blocks = ceil_div(shape_m, BLOCK_M);
|
||||
num_n_blocks = ceil_div(shape_n, BLOCK_N);
|
||||
current_shape_k = shape_k;
|
||||
if constexpr (kGemmType == GemmType::Normal) {
|
||||
num_blocks = num_m_blocks * num_n_blocks;
|
||||
} else if (kGemmType == GemmType::MGroupedContiguous) {
|
||||
num_blocks = num_m_blocks * num_n_blocks;
|
||||
this->grouped_layout = grouped_layout;
|
||||
} else if (kGemmType == GemmType::MGroupedMasked) {
|
||||
this->grouped_layout = grouped_layout;
|
||||
} else if (kGemmType == GemmType::KGroupedContiguous) {
|
||||
this->grouped_layout = grouped_layout;
|
||||
get_next_k_group(current_group_idx, current_shape_k);
|
||||
next_group_idx = current_group_idx + 1;
|
||||
get_next_k_group(next_group_idx, next_shape_k);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void get_swizzled_block_idx(const uint32_t& block_idx, uint32_t& m_block_idx, uint32_t& n_block_idx) {
|
||||
DG_STATIC_ASSERT(kNum1DBlocksPerGroup % kNumMulticast == 0, "Invalid group size");
|
||||
|
||||
// Swizzle for better L2 usages
|
||||
const auto& primary_num_blocks = kIsMulticastOnA ? num_n_blocks : num_m_blocks;
|
||||
const auto& secondary_num_blocks = kIsMulticastOnA ? num_m_blocks : num_n_blocks;
|
||||
const auto& num_blocks_per_group = secondary_num_blocks * kNum1DBlocksPerGroup;
|
||||
const auto& group_idx = block_idx / num_blocks_per_group;
|
||||
auto first_block_idx = group_idx * kNum1DBlocksPerGroup;
|
||||
auto in_group_idx = block_idx % num_blocks_per_group;
|
||||
num_blocks_in_group = min(kNum1DBlocksPerGroup, primary_num_blocks - first_block_idx);
|
||||
|
||||
// Fix unaligned TMA multicast
|
||||
// NOTES: for SM90 only, as SM90 can dynamically disable TMA multicast
|
||||
// while SM100 uses 2-CTA, which can not be dynamically disabled
|
||||
#if __CUDA_ARCH__ < 1000
|
||||
if (kNumMulticast > 1 and num_blocks_in_group % 2 != 0) {
|
||||
if (in_group_idx < (num_blocks_in_group ^ 1) * secondary_num_blocks) {
|
||||
num_blocks_in_group = num_blocks_in_group ^ 1;
|
||||
} else {
|
||||
in_group_idx = in_group_idx - (num_blocks_in_group ^ 1) * secondary_num_blocks;
|
||||
first_block_idx += num_blocks_in_group ^ 1;
|
||||
num_blocks_in_group = 1;
|
||||
}
|
||||
}
|
||||
#endif
|
||||
|
||||
// Convert to final M/N block indices
|
||||
// `kIsMulticastOnA == true` leads to groups on N
|
||||
if constexpr (kIsMulticastOnA) {
|
||||
m_block_idx = in_group_idx / num_blocks_in_group;
|
||||
n_block_idx = first_block_idx + in_group_idx % num_blocks_in_group;
|
||||
} else {
|
||||
m_block_idx = first_block_idx + in_group_idx % num_blocks_in_group;
|
||||
n_block_idx = in_group_idx / num_blocks_in_group;
|
||||
}
|
||||
}
|
||||
|
||||
template <bool kWithGroupOffset, KGroupedIndexType kIndexType = KGroupedIndexType::MN>
|
||||
__device__ __forceinline__ uint32_t get_global_idx(const uint32_t shape_dim, const uint32_t block_size,
|
||||
const uint32_t& block_idx, const uint32_t& m_block_idx = 0) {
|
||||
if constexpr (kGemmType == GemmType::Normal) {
|
||||
return block_idx * block_size;
|
||||
} else if constexpr (kGemmType == GemmType::MGroupedContiguous) {
|
||||
const auto offset = kWithGroupOffset ? cute::max(0, __ldg(grouped_layout + m_block_idx * BLOCK_M)) : 0;
|
||||
return offset * shape_dim + block_idx * block_size;
|
||||
} else if constexpr (kGemmType == GemmType::MGroupedMasked) {
|
||||
const auto offset = kWithGroupOffset ? current_group_idx : 0;
|
||||
return offset * shape_dim + block_idx * block_size;
|
||||
} else if constexpr (kGemmType == GemmType::KGroupedContiguous) {
|
||||
auto offset = 0;
|
||||
if constexpr (kWithGroupOffset) {
|
||||
if constexpr (kIndexType == KGroupedIndexType::MN)
|
||||
offset = current_group_idx * shape_dim;
|
||||
else if constexpr (kIndexType == KGroupedIndexType::K)
|
||||
offset = current_k_cumsum;
|
||||
else if constexpr (kIndexType == KGroupedIndexType::SF_K)
|
||||
offset = current_sf_k_cumsum;
|
||||
}
|
||||
return offset + block_idx * block_size;
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ bool get_next_block(uint32_t& m_block_idx, uint32_t& n_block_idx) {
|
||||
const auto next_block_idx = (++ current_iter) * kNumSMs + blockIdx.x;
|
||||
|
||||
if constexpr (kGemmType == GemmType::MGroupedMasked) {
|
||||
while (true) {
|
||||
// End of the task
|
||||
if (current_group_idx == kNumGroups)
|
||||
return false;
|
||||
|
||||
// Within current group
|
||||
num_m_blocks = ceil_div(static_cast<uint32_t>(__ldg(grouped_layout + current_group_idx)), BLOCK_M);
|
||||
const auto current_m_block_cumsum = current_m_cumsum + num_m_blocks;
|
||||
if (next_block_idx < current_m_block_cumsum * num_n_blocks)
|
||||
break;
|
||||
|
||||
// Move to check the next group
|
||||
current_group_idx ++, current_m_cumsum = current_m_block_cumsum;
|
||||
}
|
||||
|
||||
get_swizzled_block_idx(next_block_idx - current_m_cumsum * num_n_blocks, m_block_idx, n_block_idx);
|
||||
} else if (kGemmType == GemmType::KGroupedContiguous) {
|
||||
while (true) {
|
||||
// End of the task
|
||||
if (current_group_idx == kNumGroups)
|
||||
return false;
|
||||
|
||||
// Within current group
|
||||
if (next_block_idx < (current_num_valid_groups + 1) * num_m_blocks * num_n_blocks)
|
||||
break;
|
||||
|
||||
// Move to check the next group
|
||||
current_k_cumsum += current_shape_k;
|
||||
current_sf_k_cumsum += ceil_div(current_shape_k, SF_K_ALIGNMENT);
|
||||
current_num_valid_groups ++;
|
||||
|
||||
current_group_idx = next_group_idx ++;
|
||||
current_shape_k = next_shape_k;
|
||||
get_next_k_group(next_group_idx, next_shape_k);
|
||||
}
|
||||
|
||||
get_swizzled_block_idx(next_block_idx - current_num_valid_groups * num_m_blocks * num_n_blocks, m_block_idx, n_block_idx);
|
||||
} else {
|
||||
if (next_block_idx >= num_blocks)
|
||||
return false;
|
||||
|
||||
// For SM90 only
|
||||
// NOTES: we don't have to set `is_peer_cta_alive` for masked grouped GEMM, as it must be aligned
|
||||
is_peer_cta_alive = num_n_blocks % kNumMulticast == 0 or // Always aligned on N (constant bypass)
|
||||
num_m_blocks % kNumMulticast == 0 or // Always aligned on M (constant bypass)
|
||||
(next_block_idx ^ 1) < num_blocks; // Peer CTA in bound
|
||||
get_swizzled_block_idx(next_block_idx, m_block_idx, n_block_idx);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
// For SM90 only
|
||||
__device__ __forceinline__ bool is_tma_multicast_valid(const uint32_t& m_block_idx) const {
|
||||
if (num_blocks_in_group == 1)
|
||||
return false;
|
||||
if constexpr (kGemmType == GemmType::Normal or kGemmType == GemmType::MGroupedMasked or kGemmType == GemmType::KGroupedContiguous) {
|
||||
return true;
|
||||
} else {
|
||||
DG_STATIC_ASSERT(kGemmType == GemmType::MGroupedContiguous, "Invalid Gemm type");
|
||||
if constexpr (kIsMulticastOnA) {
|
||||
return true;
|
||||
} else {
|
||||
const auto& group_idx = __ldg(grouped_layout + m_block_idx * BLOCK_M);
|
||||
const auto& peer_group_idx = __ldg(grouped_layout + (m_block_idx ^ 1) * BLOCK_M);
|
||||
return group_idx == peer_group_idx;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// For SM90 only
|
||||
// ReSharper disable once CppNotAllPathsReturnValue
|
||||
__device__ __forceinline__ bool is_computation_valid(const uint32_t& m_block_idx, const uint32_t& m_offset) const {
|
||||
if constexpr (kGemmType == GemmType::Normal) {
|
||||
return true;
|
||||
} else if constexpr (kGemmType == GemmType::MGroupedContiguous) {
|
||||
return __ldg(grouped_layout + m_offset + m_block_idx * BLOCK_M) >= 0;
|
||||
} else if constexpr (kGemmType == GemmType::MGroupedMasked) {
|
||||
return m_offset + m_block_idx * BLOCK_M < __ldg(grouped_layout + current_group_idx);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
#pragma clang diagnostic pop
|
||||
|
||||
} // namespace deep_gemm
|
||||
+260
@@ -0,0 +1,260 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/atom/mma_traits_sm100.hpp>
|
||||
#include <cute/arch/mma_sm100_umma.hpp>
|
||||
#include <cute/arch/tmem_allocator_sm100.hpp>
|
||||
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
|
||||
namespace deep_gemm::sm100 {
|
||||
|
||||
template <uint32_t BLOCK_INNER, uint32_t kSwizzleMode, typename dtype_t>
|
||||
constexpr uint32_t get_inner_block_atom_size() {
|
||||
return kSwizzleMode == 0 ? BLOCK_INNER : kSwizzleMode / sizeof(dtype_t);
|
||||
}
|
||||
|
||||
template <uint32_t BLOCK_INNER, uint32_t BLOCK_OUTER,
|
||||
uint32_t kSwizzleMode, uint32_t kNumMulticast,
|
||||
typename dtype_t>
|
||||
__device__ __forceinline__ void
|
||||
tma_copy(void const* desc_ptr, cutlass::arch::ClusterTransactionBarrier* barrier_ptr,
|
||||
dtype_t* smem_ptr, const uint32_t& inner_idx, const int32_t& outer_idx) {
|
||||
DG_STATIC_ASSERT(1 <= kNumMulticast and kNumMulticast <= 2, "Invalid multicast config");
|
||||
DG_STATIC_ASSERT(static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL) ==
|
||||
static_cast<uint64_t>(cute::TMA::CacheHintSm100::EVICT_NORMAL), "Invalid cache hint");
|
||||
|
||||
// 2-CTA function will send signals to the leader CTA only
|
||||
const auto copy_func = kNumMulticast == 1 ? cute::SM90_TMA_LOAD_2D::copy : cute::SM100_TMA_2SM_LOAD_2D::copy;
|
||||
|
||||
// Issue multiple TMAs
|
||||
constexpr uint32_t BLOCK_INNER_ATOM = get_inner_block_atom_size<BLOCK_INNER, kSwizzleMode, dtype_t>();
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < BLOCK_INNER / BLOCK_INNER_ATOM; ++ i) {
|
||||
copy_func(desc_ptr, reinterpret_cast<uint64_t*>(barrier_ptr),
|
||||
static_cast<uint64_t>(cute::TMA::CacheHintSm100::EVICT_NORMAL),
|
||||
smem_ptr + i * BLOCK_OUTER * BLOCK_INNER_ATOM, inner_idx + i * BLOCK_INNER_ATOM, outer_idx);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
cute::UMMA::SmemDescriptor make_smem_desc(cute::UMMA::LayoutType layout, void* smem_ptr,
|
||||
uint32_t stride_byte_offset, uint32_t leading_byte_offset) {
|
||||
cute::UMMA::SmemDescriptor desc;
|
||||
|
||||
// Set the version for SM100
|
||||
desc.version_ = 1;
|
||||
|
||||
// Legacy mode
|
||||
desc.lbo_mode_ = 0;
|
||||
|
||||
// Layout
|
||||
desc.layout_type_ = static_cast<uint8_t>(layout);
|
||||
|
||||
// Start address
|
||||
const auto uint_ptr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
desc.start_address_ = static_cast<uint16_t>(uint_ptr >> 4);
|
||||
|
||||
// Base offset
|
||||
desc.base_offset_ = 0;
|
||||
|
||||
// SBO and LBO
|
||||
desc.stride_byte_offset_ = stride_byte_offset >> 4;
|
||||
desc.leading_byte_offset_ = leading_byte_offset >> 4;
|
||||
|
||||
return desc;
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
cute::UMMA::SmemDescriptor make_sf_desc(void* smem_ptr) {
|
||||
// NOTES: the UTCCP layout is K-major by default
|
||||
// Atom size: 8 x 128 bits
|
||||
// {SBO, LBO} means the byte stride between atoms on {MN, K}
|
||||
// Since the UTCCP we used is 128b-wide (only 1 atom on K), so LBO can be zero
|
||||
return make_smem_desc(cute::UMMA::LayoutType::SWIZZLE_NONE, smem_ptr, 8 * 16, 0);
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
void replace_smem_desc_addr(cute::UMMA::SmemDescriptor& desc, const void* smem_ptr) {
|
||||
const auto uint_ptr = cute::cast_smem_ptr_to_uint(smem_ptr);
|
||||
desc.start_address_ = static_cast<uint16_t>(uint_ptr >> 4);
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
static uint32_t get_atom_base(const cute::UMMA::LayoutType& layout_type) {
|
||||
return layout_type == cute::UMMA::LayoutType::SWIZZLE_128B_BASE32B ? 32 : 16;
|
||||
}
|
||||
|
||||
// ReSharper disable once CppNotAllPathsReturnValue
|
||||
template <cute::UMMA::Major kMajorMode, uint32_t kSwizzleMode, bool kUseBase32, typename dtype_t>
|
||||
constexpr static cute::UMMA::LayoutType to_umma_layout_type() {
|
||||
DG_STATIC_ASSERT(kSwizzleMode == 0 or kSwizzleMode == 16 or
|
||||
kSwizzleMode == 32 or kSwizzleMode == 64 or
|
||||
kSwizzleMode == 128, "Invalid swizzling mode");
|
||||
// A special case
|
||||
if constexpr ((cute::is_same_v<dtype_t, float> and kMajorMode == cute::UMMA::Major::MN) or kUseBase32) {
|
||||
DG_STATIC_ASSERT(kUseBase32, "Invalid swizzling base");
|
||||
return cute::UMMA::LayoutType::SWIZZLE_128B_BASE32B;
|
||||
}
|
||||
|
||||
// Normal cases
|
||||
if constexpr (kSwizzleMode == 0) return cute::UMMA::LayoutType::SWIZZLE_NONE;
|
||||
if constexpr (kSwizzleMode == 16) return cute::UMMA::LayoutType::SWIZZLE_NONE;
|
||||
if constexpr (kSwizzleMode == 32) return cute::UMMA::LayoutType::SWIZZLE_32B;
|
||||
if constexpr (kSwizzleMode == 64) return cute::UMMA::LayoutType::SWIZZLE_64B;
|
||||
if constexpr (kSwizzleMode == 128) return cute::UMMA::LayoutType::SWIZZLE_128B;
|
||||
}
|
||||
|
||||
template <cute::UMMA::Major kMajorMode, uint32_t BLOCK_MN, uint32_t kSwizzleMode, typename dtype_t>
|
||||
__device__ __forceinline__
|
||||
constexpr uint32_t get_umma_desc_stride_k() {
|
||||
return kMajorMode == cute::UMMA::Major::K ? 1 : get_inner_block_atom_size<BLOCK_MN, kSwizzleMode, dtype_t>();
|
||||
}
|
||||
|
||||
template <cute::UMMA::Major kMajorMode, uint32_t BLOCK_MN, uint32_t kSwizzleMode, typename dtype_t>
|
||||
__device__ __forceinline__
|
||||
uint32_t advance_umma_desc_lo(const uint32_t& base, const uint32_t& offset, const uint32_t& k_idx) {
|
||||
return base + (((offset + k_idx * get_umma_desc_stride_k<kMajorMode, BLOCK_MN, kSwizzleMode, dtype_t>()) * static_cast<uint32_t>(sizeof(dtype_t))) >> 4u);
|
||||
}
|
||||
|
||||
template <cute::UMMA::Major kMajorMode, uint32_t BLOCK_MN, uint32_t BLOCK_K, uint32_t kSwizzleMode, bool kUseBase32 = false, typename dtype_t>
|
||||
__device__ __forceinline__
|
||||
cute::UMMA::SmemDescriptor make_umma_desc(dtype_t* base_smem_ptr, uint32_t mn_idx, uint32_t k_idx) {
|
||||
const uint32_t stride_k = get_umma_desc_stride_k<kMajorMode, BLOCK_MN, kSwizzleMode, dtype_t>();
|
||||
const auto& layout_type = to_umma_layout_type<kMajorMode, kSwizzleMode, kUseBase32, dtype_t>();
|
||||
const auto& num_non_contiguous = 128 / get_atom_base(layout_type);
|
||||
if constexpr (kMajorMode == cute::UMMA::Major::K) {
|
||||
// NOTES: for K-major layout, the swizzle must be 128B (also, atom index must be 0), as `BLOCK_K` is always 128
|
||||
DG_STATIC_ASSERT(kSwizzleMode == BLOCK_K * sizeof(dtype_t), "Unexpected value");
|
||||
|
||||
// Atom size: 8 x `kSwizzleMode` (in bytes, on K)
|
||||
// {SBO, LBO} means the byte stride between atoms on {MN, K}
|
||||
// NOTES: on K, there is only 1 atom as asserted previously, so LBO can be 0
|
||||
const uint32_t stride_byte_offset = num_non_contiguous * BLOCK_K * sizeof(dtype_t);
|
||||
const uint32_t leading_byte_offset = 0;
|
||||
return make_smem_desc(layout_type,
|
||||
base_smem_ptr + mn_idx * BLOCK_K + k_idx * stride_k,
|
||||
stride_byte_offset, leading_byte_offset);
|
||||
} else {
|
||||
constexpr uint32_t BLOCK_MN_ATOM = get_inner_block_atom_size<BLOCK_MN, kSwizzleMode, dtype_t>();
|
||||
|
||||
// Must have no in-atom MN-idx
|
||||
// NOTES: no worries for the runtime assert, the `mn_idx` are constants at compilation time
|
||||
DG_DEVICE_ASSERT(mn_idx % BLOCK_MN_ATOM == 0);
|
||||
DG_STATIC_ASSERT(kSwizzleMode > 0, "Invalid swizzling");
|
||||
|
||||
// Atom size: `kSwizzleMode` (in bytes, on MN) x 8
|
||||
// NOTES: `kSwizzleMode == 16` mean non-swizzling but interleaving
|
||||
// {SBO, LBO} means the byte stride between atoms on {K, MN} for swizzling
|
||||
// {SBO, LBO} means the byte stride between atoms on {MN, K} for non-swizzling
|
||||
uint32_t stride_byte_offset = num_non_contiguous * BLOCK_MN_ATOM * sizeof(dtype_t);
|
||||
uint32_t leading_byte_offset = BLOCK_K * BLOCK_MN_ATOM * sizeof(dtype_t);
|
||||
if constexpr (kSwizzleMode == 16)
|
||||
swap(stride_byte_offset, leading_byte_offset);
|
||||
return make_smem_desc(layout_type,
|
||||
base_smem_ptr + mn_idx * BLOCK_K + k_idx * stride_k,
|
||||
stride_byte_offset, leading_byte_offset);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__
|
||||
uint64_t make_runtime_instr_desc_with_sf_id(cute::UMMA::InstrDescriptorBlockScaled desc, const uint32_t& sf_id) {
|
||||
desc.a_sf_id_ = sf_id, desc.b_sf_id_ = sf_id;
|
||||
return static_cast<uint64_t>(static_cast<uint32_t>(desc)) << 32;
|
||||
}
|
||||
|
||||
template <uint32_t kNumCols>
|
||||
__device__ constexpr uint32_t get_num_aligned_tmem_cols() {
|
||||
DG_STATIC_ASSERT(kNumCols <= 512, "Too many tensor memory columns");
|
||||
if (kNumCols <= 32) return 32;
|
||||
if (kNumCols <= 64) return 64;
|
||||
if (kNumCols <= 128) return 128;
|
||||
if (kNumCols <= 256) return 256;
|
||||
return 512;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void tcgen05_before_thread_sync() {
|
||||
asm volatile("tcgen05.fence::before_thread_sync;");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void tcgen05_after_thread_sync() {
|
||||
asm volatile("tcgen05.fence::after_thread_sync;");
|
||||
}
|
||||
|
||||
// UMMA versions with relaxed assertions
|
||||
struct SM100_MMA_F16BF16_SS {
|
||||
__device__ static void
|
||||
fma(uint64_t const& desc_a,
|
||||
uint64_t const& desc_b,
|
||||
uint32_t const& tmem_c,
|
||||
uint32_t const& scale_c,
|
||||
uint64_t const& desc) {
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
"setp.ne.b32 p, %4, 0;\n\t"
|
||||
"tcgen05.mma.cta_group::1.kind::f16 [%0], %1, %2, %3, p; \n\t"
|
||||
"}\n"
|
||||
:: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c));
|
||||
}
|
||||
};
|
||||
|
||||
struct SM100_MMA_F16BF16_2x1SM_SS {
|
||||
__device__ static void
|
||||
fma(uint64_t const& desc_a,
|
||||
uint64_t const& desc_b,
|
||||
uint32_t const& tmem_c,
|
||||
uint32_t const& scale_c,
|
||||
uint64_t const& desc) {
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
"setp.ne.b32 p, %4, 0;\n\t"
|
||||
"tcgen05.mma.cta_group::2.kind::f16 [%0], %1, %2, %3, p; \n\t"
|
||||
"}\n"
|
||||
:: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c));
|
||||
}
|
||||
};
|
||||
|
||||
struct SM100_MMA_MXF8F6F4_SS {
|
||||
__device__ static void
|
||||
fma(uint64_t const& desc_a,
|
||||
uint64_t const& desc_b,
|
||||
uint32_t const& tmem_c,
|
||||
uint32_t const& scale_c,
|
||||
uint64_t const& desc,
|
||||
uint32_t const& tmem_sfa,
|
||||
uint32_t const& tmem_sfb) {
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
"setp.ne.b32 p, %4, 0;\n\t"
|
||||
"tcgen05.mma.cta_group::1.kind::mxf8f6f4.block_scale [%0], %1, %2, %3, [%5], [%6], p; \n\t"
|
||||
"}\n"
|
||||
:
|
||||
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c),
|
||||
"r"(tmem_sfa), "r"(tmem_sfb));
|
||||
}
|
||||
};
|
||||
|
||||
struct SM100_MMA_MXF8F6F4_2x1SM_SS {
|
||||
__device__ static void
|
||||
fma(uint64_t const& desc_a,
|
||||
uint64_t const& desc_b,
|
||||
uint32_t const& tmem_c,
|
||||
uint32_t const& scale_c,
|
||||
uint64_t const& desc,
|
||||
uint32_t const& tmem_sfa,
|
||||
uint32_t const& tmem_sfb) {
|
||||
asm volatile(
|
||||
"{\n\t"
|
||||
".reg .pred p;\n\t"
|
||||
"setp.ne.b32 p, %4, 0;\n\t"
|
||||
"tcgen05.mma.cta_group::2.kind::mxf8f6f4.block_scale [%0], %1, %2, %3, [%5], [%6], p; \n\t"
|
||||
"}\n"
|
||||
:
|
||||
: "r"(tmem_c), "l"(desc_a), "l"(desc_b), "r"(static_cast<uint32_t>(desc >> 32)), "r"(scale_c),
|
||||
"r"(tmem_sfa), "r"(tmem_sfb));
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace `deep_gemm::sm100`
|
||||
+283
@@ -0,0 +1,283 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/arch/copy_sm90_tma.hpp>
|
||||
#include <cute/arch/cluster_sm90.hpp>
|
||||
#include <cute/arch/mma_sm90_gmma.hpp>
|
||||
#include <cute/arch/mma_sm90_gmma_ext.hpp>
|
||||
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
|
||||
namespace deep_gemm::sm90 {
|
||||
|
||||
template <int N_, typename MMA>
|
||||
struct FP8MMA {
|
||||
|
||||
template <size_t ...Idx>
|
||||
__forceinline__ __device__ static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d, cute::index_sequence<Idx...>) {
|
||||
using namespace cute::SM90::GMMA;
|
||||
MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero));
|
||||
}
|
||||
|
||||
__forceinline__ __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
|
||||
call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence<N_/2>{});
|
||||
}
|
||||
|
||||
static constexpr int M = 64;
|
||||
static constexpr int N = N_;
|
||||
static constexpr int K = 32;
|
||||
static constexpr int kNumAccum = M * N / 128;
|
||||
};
|
||||
|
||||
template <int N>
|
||||
struct FP8MMASelector {
|
||||
|
||||
static constexpr auto select_mma() {
|
||||
using namespace cute::SM90::GMMA;
|
||||
if constexpr (N == 8) return MMA_64x8x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 16) return MMA_64x16x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 24) return MMA_64x24x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 32) return MMA_64x32x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 40) return MMA_64x40x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 48) return MMA_64x48x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 56) return MMA_64x56x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 64) return MMA_64x64x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 72) return MMA_64x72x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 80) return MMA_64x80x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 88) return MMA_64x88x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 96) return MMA_64x96x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 104) return MMA_64x104x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 112) return MMA_64x112x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 120) return MMA_64x120x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 128) return MMA_64x128x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 136) return MMA_64x136x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 144) return MMA_64x144x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 152) return MMA_64x152x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 160) return MMA_64x160x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 168) return MMA_64x168x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 176) return MMA_64x176x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 184) return MMA_64x184x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 192) return MMA_64x192x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 200) return MMA_64x200x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 208) return MMA_64x208x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 216) return MMA_64x216x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 224) return MMA_64x224x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 232) return MMA_64x232x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 240) return MMA_64x240x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 248) return MMA_64x248x32_F32E4M3E4M3_SS_TN();
|
||||
if constexpr (N == 256) return MMA_64x256x32_F32E4M3E4M3_SS_TN();
|
||||
}
|
||||
|
||||
static constexpr auto select_type() {
|
||||
return FP8MMA<N, decltype(select_mma())>();
|
||||
}
|
||||
|
||||
using type = decltype(select_type());
|
||||
};
|
||||
|
||||
template <int N_, typename MMA>
|
||||
struct BF16MMA {
|
||||
|
||||
template <size_t ...Idx>
|
||||
__forceinline__ __device__ static void call_fma_impl(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d, cute::index_sequence<Idx...>) {
|
||||
using namespace cute::SM90::GMMA;
|
||||
MMA::fma(desc_a, desc_b, d[Idx]..., (scale_d ? ScaleOut::One : ScaleOut::Zero));
|
||||
}
|
||||
|
||||
__forceinline__ __device__ static void wgmma(uint64_t const& desc_a, uint64_t const& desc_b, float* d, bool scale_d) {
|
||||
call_fma_impl(desc_a, desc_b, d, scale_d, cute::make_index_sequence<N_/2>{});
|
||||
}
|
||||
|
||||
static constexpr int M = 64;
|
||||
static constexpr int N = N_;
|
||||
static constexpr int K = 16;
|
||||
static constexpr int kNumAccum = M * N / 128;
|
||||
};
|
||||
|
||||
template <int N>
|
||||
struct BF16MMASelector {
|
||||
|
||||
static constexpr auto select_mma() {
|
||||
using namespace cute::SM90::GMMA;
|
||||
if constexpr (N == 8) return MMA_64x8x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 16) return MMA_64x16x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 24) return MMA_64x24x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 32) return MMA_64x32x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 40) return MMA_64x40x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 48) return MMA_64x48x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 56) return MMA_64x56x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 64) return MMA_64x64x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 72) return MMA_64x72x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 80) return MMA_64x80x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 88) return MMA_64x88x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 96) return MMA_64x96x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 104) return MMA_64x104x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 112) return MMA_64x112x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 120) return MMA_64x120x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 128) return MMA_64x128x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 136) return MMA_64x136x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 144) return MMA_64x144x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 152) return MMA_64x152x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 160) return MMA_64x160x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 168) return MMA_64x168x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 176) return MMA_64x176x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 184) return MMA_64x184x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 192) return MMA_64x192x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 200) return MMA_64x200x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 208) return MMA_64x208x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 216) return MMA_64x216x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 224) return MMA_64x224x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 232) return MMA_64x232x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 240) return MMA_64x240x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 248) return MMA_64x248x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
if constexpr (N == 256) return MMA_64x256x16_F32BF16BF16_SS<Major::K, Major::K>();
|
||||
}
|
||||
|
||||
static constexpr auto select_type() {
|
||||
return BF16MMA<N, decltype(select_mma())>();
|
||||
}
|
||||
|
||||
using type = decltype(select_type());
|
||||
};
|
||||
|
||||
|
||||
template <typename dtype_t>
|
||||
struct SM90_U32x2_STSM_N {
|
||||
__device__ __forceinline__ static void
|
||||
copy(dtype_t src_0, dtype_t src_1, void* smem_dst) {
|
||||
const uint32_t src[2] = {*reinterpret_cast<uint32_t*>(&src_0), *reinterpret_cast<uint32_t*>(&src_1)};
|
||||
asm volatile("stmatrix.sync.aligned.x2.m8n8.shared.b16 [%0], {%1, %2};\n"
|
||||
:: "l"(smem_dst), "r"(src[0]), "r"(src[1]));
|
||||
}
|
||||
};
|
||||
|
||||
struct SM90_U32x2_LDSM_N {
|
||||
__device__ __forceinline__ static void
|
||||
copy(uint32_t& dst_0, uint32_t& dst_1, void* smem_src) {
|
||||
asm volatile("ldmatrix.sync.aligned.x2.m8n8.shared.b16 {%0, %1}, [%2];\n"
|
||||
: "=r"(dst_0), "=r"(dst_1)
|
||||
: "l"(smem_src));
|
||||
}
|
||||
};
|
||||
|
||||
struct SM90_U32x4_LDSM_N {
|
||||
__device__ __forceinline__ static void
|
||||
copy(uint32_t& dst_0, uint32_t& dst_1, uint32_t& dst_2, uint32_t& dst_3, void* smem_src) {
|
||||
asm volatile("ldmatrix.sync.aligned.x4.m8n8.shared.b16 {%0, %1, %2, %3}, [%4];\n"
|
||||
: "=r"(dst_0), "=r"(dst_1), "=r"(dst_2), "=r"(dst_3)
|
||||
: "l"(smem_src));
|
||||
}
|
||||
};
|
||||
|
||||
__forceinline__ __device__ void warpgroup_arrive() {
|
||||
asm volatile("wgmma.fence.sync.aligned;\n" ::: "memory");
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void warpgroup_commit_batch() {
|
||||
asm volatile("wgmma.commit_group.sync.aligned;\n" ::: "memory");
|
||||
}
|
||||
|
||||
__forceinline__ __device__ void warpgroup_fence_operand(float& reg) {
|
||||
asm volatile("" : "+f"(reg) :: "memory");
|
||||
}
|
||||
|
||||
template <int N>
|
||||
__forceinline__ __device__ void warpgroup_wait() {
|
||||
DG_STATIC_ASSERT(N >= 0 and N <= 7, "WGMMA wait: N must be in range [0, 7]");
|
||||
asm volatile("wgmma.wait_group.sync.aligned %0;\n" :: "n"(N) : "memory");
|
||||
}
|
||||
|
||||
// TODO: replace with CUTLASS solution
|
||||
union GmmaDescriptor {
|
||||
__host__ __device__ constexpr GmmaDescriptor() noexcept: desc_(0) {}
|
||||
|
||||
__host__ __device__ constexpr GmmaDescriptor(uint64_t desc) noexcept: desc_(desc) {}
|
||||
|
||||
__host__ __device__ constexpr GmmaDescriptor(GmmaDescriptor const &t) noexcept: desc_(t.desc_) {}
|
||||
|
||||
__host__ __device__ constexpr GmmaDescriptor(GmmaDescriptor &&t) noexcept: desc_(t.desc_) {}
|
||||
|
||||
__host__ __device__ constexpr GmmaDescriptor &operator=(GmmaDescriptor const &t) noexcept {
|
||||
desc_ = t.desc_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
__host__ __device__ constexpr GmmaDescriptor &operator=(GmmaDescriptor &&t) noexcept {
|
||||
desc_ = t.desc_;
|
||||
return *this;
|
||||
}
|
||||
|
||||
uint64_t desc_;
|
||||
uint32_t reg32_[2];
|
||||
uint16_t reg16_[4];
|
||||
|
||||
struct {
|
||||
uint16_t start_address_: 14, : 2;
|
||||
uint16_t leading_byte_offset_: 14, : 2;
|
||||
uint16_t stride_byte_offset_: 14, : 2;
|
||||
uint8_t : 1, base_offset_: 3, : 4;
|
||||
uint8_t : 6, layout_type_: 2;
|
||||
} bitfield;
|
||||
|
||||
// Decay to an `uint64_t`
|
||||
__host__ __device__ constexpr operator uint64_t() const noexcept { return desc_; }
|
||||
};
|
||||
|
||||
template <class PointerType>
|
||||
__device__ GmmaDescriptor make_smem_desc(PointerType smem_ptr, const int& layout_type,
|
||||
const int& leading_byte_offset = 0,
|
||||
const int& stride_byte_offset = 1024) {
|
||||
GmmaDescriptor desc;
|
||||
const auto& uint_ptr = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));
|
||||
desc.bitfield.start_address_ = uint_ptr >> 4;
|
||||
desc.bitfield.layout_type_ = layout_type;
|
||||
desc.bitfield.leading_byte_offset_ = leading_byte_offset >> 4;
|
||||
desc.bitfield.stride_byte_offset_ = stride_byte_offset >> 4;
|
||||
desc.bitfield.base_offset_ = 0;
|
||||
return desc;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void
|
||||
tma_copy(void const* desc_ptr, uint64_t* barrier_ptr, void* smem_ptr,
|
||||
const uint32_t& crd_0, const uint32_t& crd_1, const uint32_t& num_tma_multicast = 1) {
|
||||
constexpr auto cache_hint = static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL);
|
||||
if (num_tma_multicast == 1) {
|
||||
cute::SM90_TMA_LOAD_2D::copy(desc_ptr, barrier_ptr, cache_hint, smem_ptr, crd_0, crd_1);
|
||||
} else if (cute::block_rank_in_cluster() == 0) {
|
||||
cute::SM90_TMA_LOAD_MULTICAST_2D::copy(desc_ptr, barrier_ptr, (1 << num_tma_multicast) - 1, cache_hint, smem_ptr, crd_0, crd_1);
|
||||
}
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void
|
||||
tma_3d_copy(void const* desc_ptr, uint64_t* barrier_ptr, void* smem_ptr,
|
||||
const uint32_t& crd_0, const uint32_t& crd_1, const uint32_t& crd_2) {
|
||||
constexpr auto cache_hint = static_cast<uint64_t>(cute::TMA::CacheHintSm90::EVICT_NORMAL);
|
||||
cute::SM90_TMA_LOAD_3D::copy(desc_ptr, barrier_ptr, cache_hint, smem_ptr, crd_0, crd_1, crd_2);
|
||||
}
|
||||
|
||||
// Tensormap related
|
||||
__device__ __forceinline__ void tensor_map_release_cta() {
|
||||
asm volatile ("fence.proxy.tensormap::generic.release.cta;");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void tensor_map_acquire_cta(const cute::TmaDescriptor* gmem_desc_ptr) {
|
||||
auto gmem_int_desc = reinterpret_cast<uint64_t>(gmem_desc_ptr);
|
||||
asm volatile ("fence.proxy.tensormap::generic.acquire.cta [%0], 128;" :: "l"(gmem_int_desc) : "memory");
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void tensor_map_replace_global_addr_in_smem(cute::TmaDescriptor* smem_desc, const void* new_addr) {
|
||||
auto smem_int_desc = static_cast<uint32_t>(__cvta_generic_to_shared(smem_desc));
|
||||
const auto new_int64_addr = reinterpret_cast<uint64_t>(new_addr);
|
||||
asm volatile ("tensormap.replace.tile.global_address.shared::cta.b1024.b64 [%0], %1;" :: "r"(smem_int_desc), "l"(new_int64_addr));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void tensor_map_replace_global_inner_dim_stride_in_smem(cute::TmaDescriptor* smem_desc, const uint32_t& new_dim, const uint64_t& new_stride) {
|
||||
auto smem_int_desc = __cvta_generic_to_shared(smem_desc);
|
||||
asm volatile ("tensormap.replace.tile.global_dim.shared::cta.b1024.b32 [%0], 0, %1;" :: "l"(smem_int_desc), "r"(new_dim));
|
||||
#if ((__CUDACC_VER_MAJOR__ > 12) or ((__CUDACC_VER_MAJOR__ == 12) and (__CUDACC_VER_MINOR__ >= 3)))
|
||||
asm volatile("tensormap.replace.tile.global_stride.shared::cta.b1024.b64 [%0], 0, %1;" :: "l"(smem_int_desc), "l"(new_stride));
|
||||
#else
|
||||
DG_STATIC_ASSERT(false, "Invalid CUDA version");
|
||||
#endif
|
||||
}
|
||||
|
||||
} // namespace `deep_gemm::sm90`
|
||||
+18
@@ -0,0 +1,18 @@
|
||||
#pragma once
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
enum class GemmType {
|
||||
Normal = 0,
|
||||
MGroupedContiguous = 1,
|
||||
MGroupedMasked = 2,
|
||||
KGroupedContiguous = 3,
|
||||
};
|
||||
|
||||
enum class KernelType {
|
||||
Kernel1D1D = 0,
|
||||
Kernel1D2D = 1,
|
||||
KernelNoSF = 2
|
||||
};
|
||||
|
||||
} // namespace deep_gemm
|
||||
+179
@@ -0,0 +1,179 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/utility>
|
||||
#include <cute/container/tuple.hpp>
|
||||
|
||||
#include "cute_tie.cuh"
|
||||
|
||||
#ifdef __CLION_IDE__
|
||||
|
||||
__host__ __device__ __forceinline__ void host_device_printf(const char* format, ...) {
|
||||
asm volatile("trap;");
|
||||
}
|
||||
|
||||
#define printf host_device_printf
|
||||
#endif
|
||||
|
||||
#ifndef DG_DEVICE_ASSERT
|
||||
#define DG_DEVICE_ASSERT(cond) \
|
||||
do { \
|
||||
if (not (cond)) { \
|
||||
printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \
|
||||
asm("trap;"); \
|
||||
} \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#ifndef DG_TRAP_ONLY_DEVICE_ASSERT
|
||||
#define DG_TRAP_ONLY_DEVICE_ASSERT(cond) \
|
||||
do { \
|
||||
if (not (cond)) \
|
||||
asm("trap;"); \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#ifndef DG_STATIC_ASSERT
|
||||
#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__)
|
||||
#endif
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
template <typename FuncT>
|
||||
struct PatternVisitor {
|
||||
FuncT func;
|
||||
|
||||
__device__ __host__
|
||||
explicit PatternVisitor(FuncT&& func): func(std::forward<FuncT>(func)) {}
|
||||
|
||||
__device__ __host__
|
||||
auto operator [](const uint32_t& i) {
|
||||
return func(i);
|
||||
}
|
||||
};
|
||||
|
||||
template <typename T>
|
||||
__device__ __host__ T ceil_div(T a, T b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __host__ constexpr T constexpr_ceil_div(T a, T b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __host__ T align(T a, T b) {
|
||||
return ceil_div(a, b) * b;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __host__ constexpr T constexpr_align(T a, T b) {
|
||||
return constexpr_ceil_div(a, b) * b;
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__device__ __host__ constexpr T constexpr_gcd(T a, T b) {
|
||||
return b == 0 ? a : constexpr_gcd(b, a % b);
|
||||
}
|
||||
|
||||
template<typename T>
|
||||
__forceinline__ __device__ void swap(T& a, T& b) {
|
||||
T temp = a;
|
||||
a = b;
|
||||
b = temp;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ uint32_t get_sm_idx() {
|
||||
uint32_t sm_idx;
|
||||
asm ("mov.u32 %0, %%smid;" : "=r"(sm_idx));
|
||||
return sm_idx;
|
||||
}
|
||||
|
||||
__forceinline__ __device__ uint32_t get_lane_idx() {
|
||||
uint32_t lane_id;
|
||||
asm ("mov.u32 %0, %laneid;" : "=r"(lane_id));
|
||||
return lane_id;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint32_t ld_shared(const uint32_t* ptr) {
|
||||
uint32_t ret;
|
||||
asm volatile("ld.shared.u32 %0, [%1];" : "=r"(ret) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float2 ld_shared(const float2* ptr) {
|
||||
float2 ret;
|
||||
asm volatile("ld.shared.v2.f32 {%0, %1}, [%2];" : "=f"(ret.x), "=f"(ret.y) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float4 ld_shared(const float4* ptr) {
|
||||
float4 ret;
|
||||
asm volatile("ld.shared.v4.f32 {%0, %1, %2, %3}, [%4];" : "=f"(ret.x), "=f"(ret.y), "=f"(ret.z), "=f"(ret.w) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ uint4 ld_shared(const uint4* ptr) {
|
||||
uint4 ret;
|
||||
asm volatile("ld.shared.v4.u32 {%0, %1, %2, %3}, [%4];" : "=r"(ret.x), "=r"(ret.y), "=r"(ret.z), "=r"(ret.w) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ float ld_shared(const float* ptr) {
|
||||
float ret;
|
||||
asm volatile("ld.shared.f32 %0, [%1];" : "=f"(ret) : "l"(ptr));
|
||||
return ret;
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st_shared(const float* ptr, float val) {
|
||||
asm volatile("st.shared.f32 [%0], %1;" :: "l"(ptr), "f"(val));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st_shared(const float2* ptr, float2 val) {
|
||||
asm volatile("st.shared.v2.f32 [%0], {%1, %2};" :: "l"(ptr), "f"(val.x), "f"(val.y));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st_shared(const uint32_t* ptr, uint32_t val) {
|
||||
asm volatile("st.shared.u32 [%0], %1;" :: "l"(ptr), "r"(val));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st_shared(const void* ptr, uint32_t x, uint32_t y) {
|
||||
asm volatile("st.shared.v2.u32 [%0], {%1, %2};" :: "l"(ptr), "r"(x), "r"(y));
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void st_shared(const void* ptr, uint32_t x, uint32_t y, uint32_t z, uint32_t w) {
|
||||
asm volatile("st.shared.v4.u32 [%0], {%1, %2, %3, %4};" :: "l"(ptr), "r"(x), "r"(y), "r"(z), "r"(w));
|
||||
}
|
||||
|
||||
template <typename old_t>
|
||||
__device__ __forceinline__ int cast_into_bf16_and_pack(old_t& x, old_t& y) {
|
||||
auto bf16x2 = __float22bfloat162_rn({*reinterpret_cast<float*>(&x), *reinterpret_cast<float*>(&y)});
|
||||
return *reinterpret_cast<int*>(&bf16x2);
|
||||
}
|
||||
|
||||
__device__ __forceinline__ void prefetch_l1(void *ptr) {
|
||||
asm volatile("prefetch.global.L1 [%0];" :: "l"(ptr));
|
||||
}
|
||||
|
||||
template <uint32_t kNumBytes>
|
||||
struct Vectorized {
|
||||
static auto zeros() {
|
||||
// TODO: add `ulonglong4` for SM100 once `__ldg` support this
|
||||
if constexpr (kNumBytes > 0 and kNumBytes % 16 == 0) {
|
||||
return make_uint4(0, 0, 0, 0);
|
||||
} else if constexpr (kNumBytes > 0 and kNumBytes % 8 == 0) {
|
||||
return make_uint2(0, 0);
|
||||
} else if constexpr (kNumBytes > 0 and kNumBytes % 4 == 0) {
|
||||
return 0;
|
||||
} else {
|
||||
DG_STATIC_ASSERT(kNumBytes > 0 and kNumBytes % 4 == 0, "Invalid vectorization");
|
||||
}
|
||||
}
|
||||
|
||||
using vec_t = decltype(zeros());
|
||||
};
|
||||
|
||||
} // namespace `deep_gemm`
|
||||
+408
@@ -0,0 +1,408 @@
|
||||
#pragma once
|
||||
|
||||
#pragma clang diagnostic push
|
||||
#pragma clang diagnostic ignored "-Wunknown-attributes"
|
||||
|
||||
#include <cutlass/arch/barrier.h>
|
||||
#include <cutlass/arch/reg_reconfig.h>
|
||||
|
||||
#include <cute/arch/cluster_sm90.hpp>
|
||||
#include <cute/arch/copy_sm90_desc.hpp>
|
||||
#include <cute/arch/copy_sm90_tma.hpp>
|
||||
|
||||
#include <deep_gemm/common/epilogue_utils.cuh>
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
#include <deep_gemm/common/scheduler.cuh>
|
||||
#include <deep_gemm/common/sm90_utils.cuh>
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
using namespace deep_gemm::sm90;
|
||||
|
||||
template <uint32_t kNumFormerIters, uint32_t kGap, uint32_t kEnd, typename func_t>
|
||||
__device__ void dispatch_num_former_iters(uint32_t num_former_iters, const func_t& func) {
|
||||
if (num_former_iters == kNumFormerIters) {
|
||||
func(cute::Int<kNumFormerIters>{});
|
||||
return;
|
||||
}
|
||||
|
||||
if constexpr (kNumFormerIters + kGap <= kEnd)
|
||||
dispatch_num_former_iters<kNumFormerIters + kGap, kGap, kEnd>(num_former_iters, func);
|
||||
}
|
||||
|
||||
template <uint32_t SHAPE_M, uint32_t SHAPE_N, uint32_t SHAPE_K,
|
||||
uint32_t kNumGroups,
|
||||
uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
|
||||
uint32_t kSwizzleDMode,
|
||||
uint32_t kNumStages, uint32_t kNumLastStages,
|
||||
uint32_t kNumTMAThreads, uint32_t kNumMathThreads,
|
||||
uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
|
||||
uint32_t kNumSMs, GemmType kGemmType,
|
||||
typename epilogue_type_t>
|
||||
__global__ __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void
|
||||
sm90_fp8_gemm_1d2d_impl(float* sfb, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_a,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_b,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_d,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) {
|
||||
#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900)) or defined(__CLION_IDE__)
|
||||
// Scaling checks
|
||||
DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
|
||||
DG_STATIC_ASSERT(constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or (constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block");
|
||||
|
||||
// Types
|
||||
using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
|
||||
using Barrier = cutlass::arch::ClusterTransactionBarrier;
|
||||
DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
|
||||
|
||||
// Overwrite shape constants if the compiler gives
|
||||
shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m;
|
||||
shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n;
|
||||
shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k;
|
||||
|
||||
// Shared memory
|
||||
static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
|
||||
static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(__nv_bfloat16);
|
||||
static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
|
||||
static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
|
||||
static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
|
||||
const uint32_t& shape_k_scales = ceil_div(shape_k, BLOCK_K);
|
||||
const uint32_t& smem_sfb_size = align<uint32_t>(shape_k_scales * (kMustUseUniformedScaleB ? 1 : 2) * sizeof(float), sizeof(Barrier));
|
||||
|
||||
// Configs
|
||||
const uint32_t num_total_k_blocks = ceil_div(shape_k, BLOCK_K);
|
||||
const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
|
||||
const uint32_t lane_idx = get_lane_idx();
|
||||
|
||||
// Prefetch TMA descriptors at the very beginning
|
||||
if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
|
||||
cute::prefetch_tma_descriptor(&tensor_map_a);
|
||||
cute::prefetch_tma_descriptor(&tensor_map_b);
|
||||
cute::prefetch_tma_descriptor(&tensor_map_sfa);
|
||||
cute::prefetch_tma_descriptor(&tensor_map_d);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
// Align to 1024 bytes for swizzle-128B
|
||||
extern __shared__ __align__(1024) uint8_t smem_buffer[];
|
||||
DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes");
|
||||
|
||||
// Data on shared memory
|
||||
auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer);
|
||||
auto smem_a = PatternVisitor([&](const uint32_t& i) {
|
||||
return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);
|
||||
});
|
||||
auto smem_b = PatternVisitor([&](const uint32_t& i) {
|
||||
return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
|
||||
});
|
||||
constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
|
||||
auto smem_sfa = PatternVisitor([&](const uint32_t& i) {
|
||||
return reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + i * SMEM_SFA_SIZE_PER_STAGE);
|
||||
});
|
||||
auto smem_sfb = reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE);
|
||||
|
||||
// Fill barriers
|
||||
auto barrier_start_ptr = reinterpret_cast<Barrier*>(reinterpret_cast<uint8_t*>(smem_sfb) + smem_sfb_size);
|
||||
auto full_barriers = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; });
|
||||
auto empty_barriers = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; });
|
||||
|
||||
// Initialize barriers
|
||||
DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
|
||||
if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) {
|
||||
// NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
|
||||
// even with TMA multicast disabled, we want to make the behavior aligned
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumStages; ++ i) {
|
||||
full_barriers[i]->init(1);
|
||||
empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
|
||||
}
|
||||
|
||||
// Make initialized barrier visible in async proxy
|
||||
cutlass::arch::fence_barrier_init();
|
||||
}
|
||||
|
||||
// Synchronize all threads to make barrier visible in normal memory model
|
||||
(kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads();
|
||||
|
||||
// Register reconfigurations
|
||||
constexpr uint32_t kNumTMARegisters = 40;
|
||||
constexpr uint32_t kNumMathRegisters = 232;
|
||||
|
||||
// Block scheduler
|
||||
uint32_t m_block_idx, n_block_idx;
|
||||
auto scheduler = Scheduler<kGemmType, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA, kNumSMs>(shape_m, shape_n, shape_k, grouped_layout);
|
||||
|
||||
// Pipeline and TMA phases
|
||||
uint32_t stage_idx = 0, phase = 0;
|
||||
auto advance_pipeline = [&](uint32_t& k_block_idx) {
|
||||
++ k_block_idx;
|
||||
|
||||
// Flip phases only if reach the next first stage
|
||||
stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
|
||||
phase ^= stage_idx == 0;
|
||||
};
|
||||
|
||||
if (warp_idx >= kNumMathThreads / 32) {
|
||||
// TMA warp-group for loading data
|
||||
cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();
|
||||
|
||||
// NOTES: only one thread (or warp) will be used
|
||||
if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
|
||||
// Persistently schedule over blocks
|
||||
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
|
||||
// Assign TMA multicast number into A and B
|
||||
// NOTES: there may be additional odd rows/columns or cases where multicast is not possible.
|
||||
const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx);
|
||||
const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
|
||||
const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
|
||||
DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast");
|
||||
|
||||
for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
|
||||
// Wait consumer release
|
||||
empty_barriers[stage_idx]->wait(phase ^ 1);
|
||||
|
||||
// Issue TMA A
|
||||
constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked;
|
||||
auto& full_barrier = *full_barriers[stage_idx];
|
||||
const uint32_t k_idx = k_block_idx * BLOCK_K;
|
||||
tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
|
||||
smem_a[stage_idx], k_idx, scheduler.get_global_idx<kWithGroupOffsetA>(shape_m, BLOCK_M, m_block_idx),
|
||||
num_tma_multicast_a);
|
||||
tma_copy(&tensor_map_sfa, reinterpret_cast<uint64_t*>(&full_barrier),
|
||||
smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.get_global_idx<kWithGroupOffsetA>(shape_k_scales, 1, k_block_idx),
|
||||
num_tma_multicast_a);
|
||||
|
||||
// Issue TMA B
|
||||
tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
|
||||
smem_b[stage_idx], k_idx, scheduler.get_global_idx<true>(shape_n, BLOCK_N, n_block_idx, m_block_idx),
|
||||
num_tma_multicast_b);
|
||||
full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
|
||||
}
|
||||
}
|
||||
|
||||
// To safely deconstruct distributed shared barriers, we need another round of empty waits
|
||||
if constexpr (kNumTMAMulticast > 1) {
|
||||
for (uint32_t i = 0; i < kNumStages; advance_pipeline(i))
|
||||
empty_barriers[stage_idx]->wait(phase ^ 1);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Math warp-groups for WGMMA
|
||||
cutlass::arch::warpgroup_reg_alloc<kNumMathRegisters>();
|
||||
|
||||
// NOTES: use `__shfl_sync` to encourage NVCC to use unified registers
|
||||
const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
|
||||
const auto r_0 = warp_idx * 16 + lane_idx / 4, r_1 = r_0 + 8;
|
||||
|
||||
auto a_desc = make_smem_desc(smem_a[0] + math_wg_idx * WGMMA::M * BLOCK_K, 1);
|
||||
auto b_desc = make_smem_desc(smem_b[0], 1);
|
||||
const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0);
|
||||
const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0);
|
||||
|
||||
// Persistently schedule over blocks
|
||||
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
|
||||
// Decide the number of scales B to load
|
||||
DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0);
|
||||
uint32_t num_former_iters = BLOCK_N / 8, num_full_iters = num_former_iters;
|
||||
if constexpr (not kMustUseUniformedScaleB) {
|
||||
num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8;
|
||||
num_full_iters = min(shape_n - n_block_idx * BLOCK_N, BLOCK_N) / 8;
|
||||
}
|
||||
uint32_t num_sfb = shape_k_scales * (num_former_iters >= num_full_iters ? 1 : 2);
|
||||
|
||||
// Load B scales with math warp-groups
|
||||
// NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks
|
||||
if (threadIdx.x >= 32) {
|
||||
auto num_previous_lines = scheduler.get_global_idx<true>(ceil_div(shape_n, BLOCK_K), 0, 0, m_block_idx);
|
||||
auto local_sfb = sfb + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * shape_k_scales;
|
||||
#pragma unroll
|
||||
for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32)
|
||||
st_shared(smem_sfb + i, __ldg(local_sfb + i));
|
||||
}
|
||||
cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
|
||||
|
||||
// Accumulation for WGMMA or CUDA promotion
|
||||
constexpr uint32_t WAVE_BLOCK_M = WGMMA::M * (BLOCK_M <= 64 ? 1 : 2);
|
||||
DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");
|
||||
float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
|
||||
|
||||
// Empty barrier arrival
|
||||
auto empty_barrier_arrive = [&]() {
|
||||
if constexpr (kNumTMAMulticast == 1) {
|
||||
lane_idx == 0 ? empty_barriers[stage_idx]->arrive() : void();
|
||||
} else {
|
||||
auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster();
|
||||
lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void();
|
||||
}
|
||||
};
|
||||
|
||||
// Skip useless computations
|
||||
if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) {
|
||||
// The compiler must know the dynamic variable `num_former_iters`'s real value
|
||||
constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;
|
||||
constexpr uint32_t kGap = constexpr_gcd(BLOCK_K, BLOCK_N) / 8;
|
||||
constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0;
|
||||
|
||||
// Dispatch `num_former_iters` and launch MMAs
|
||||
dispatch_num_former_iters<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) {
|
||||
#pragma unroll 8
|
||||
for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
|
||||
const auto& a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16);
|
||||
const auto& b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16);
|
||||
|
||||
// Read B scales
|
||||
float scale_b_0 = ld_shared(smem_sfb + k_block_idx), scale_b_1;
|
||||
// NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks
|
||||
if constexpr (not kMustUseUniformedScaleB)
|
||||
scale_b_1 = ld_shared(smem_sfb + k_block_idx + shape_k_scales);
|
||||
|
||||
// Wait TMA arrivals
|
||||
full_barriers[stage_idx]->wait(phase);
|
||||
|
||||
// TODO: remove some useless computation for unaligned Ms
|
||||
#pragma unroll
|
||||
for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
|
||||
auto m_offset = local_idx * WAVE_BLOCK_M;
|
||||
|
||||
// Read A scales
|
||||
// NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
|
||||
auto scale_a_0 = ld_shared(smem_sfa[stage_idx] + r_0 + m_offset);
|
||||
auto scale_a_1 = ld_shared(smem_sfa[stage_idx] + r_1 + m_offset);
|
||||
|
||||
// Commit WGMMA instructions
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
|
||||
warpgroup_fence_operand(accum[i]);
|
||||
warpgroup_arrive();
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
|
||||
a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16;
|
||||
b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16;
|
||||
WGMMA::wgmma(a_desc, b_desc, accum, k);
|
||||
}
|
||||
warpgroup_commit_batch();
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
|
||||
warpgroup_fence_operand(accum[i]);
|
||||
warpgroup_wait<0>();
|
||||
|
||||
// Notify barrier arrival at the last warpgroup wave
|
||||
if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
|
||||
empty_barrier_arrive();
|
||||
|
||||
// Promote with scales
|
||||
// NOTES: making it as predicates is very important for performance, comparing to two loops
|
||||
float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
|
||||
float scale_0_1, scale_1_1;
|
||||
if constexpr (not kMustUseUniformedScaleB)
|
||||
scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
|
||||
|
||||
auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
|
||||
// NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant
|
||||
bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
|
||||
shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
|
||||
shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
|
||||
shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
|
||||
shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
|
||||
full_barriers[stage_idx]->wait(phase);
|
||||
empty_barrier_arrive();
|
||||
}
|
||||
}
|
||||
|
||||
// TMA checks
|
||||
constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);
|
||||
constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes);
|
||||
constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4;
|
||||
DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom");
|
||||
DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,
|
||||
"Unaligned TMA store or too many TMA store instructions");
|
||||
DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");
|
||||
|
||||
// Wait last TMA store to be finished
|
||||
if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N)
|
||||
cute::tma_store_wait<0>();
|
||||
cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
|
||||
|
||||
// Write back to shared memory using STSM and issue TMA stores
|
||||
DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
|
||||
#pragma unroll
|
||||
for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
|
||||
auto m_offset = local_idx * WAVE_BLOCK_M;
|
||||
auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
|
||||
#pragma unroll
|
||||
for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
|
||||
// Swizzle or padding into the correct address
|
||||
uint8_t* smem_ptr = nullptr;
|
||||
if constexpr (kSwizzleDMode > 0) {
|
||||
// Calculate the swizzling atom offset and in-atom offset
|
||||
constexpr uint32_t kNumBankGroupBytes = 16;
|
||||
auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);
|
||||
|
||||
// Calculate the index of the bank group to be written in the atom
|
||||
auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);
|
||||
|
||||
// Reshape the atom in another view and swizzle
|
||||
// - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
|
||||
// - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
|
||||
constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
|
||||
auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
|
||||
auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
|
||||
col ^= row % (kSwizzleDMode / 16);
|
||||
|
||||
// Add back into the base pointer
|
||||
// NOTES: think twice before modifying this, as changes may affect the number of instructions
|
||||
smem_ptr = reinterpret_cast<uint8_t*>(smem_d) + // Base pointer
|
||||
warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + // Warp offset
|
||||
m_offset * kSwizzleDMode + // Wave offset
|
||||
atom_offset * BLOCK_M * kSwizzleDMode + // Swizzle atom offset (constants)
|
||||
row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
|
||||
} else {
|
||||
// No swizzling, just padding
|
||||
smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8);
|
||||
}
|
||||
|
||||
// NOTES: only 16 lanes' addresses are used
|
||||
SM90_U32x2_STSM_N<nv_bfloat162>::copy(
|
||||
__float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}),
|
||||
__float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}),
|
||||
smem_ptr
|
||||
);
|
||||
}
|
||||
}
|
||||
cute::tma_store_fence();
|
||||
cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
|
||||
|
||||
// Use TMA store to write back to global memory
|
||||
// TODO: compatible with FP32 output
|
||||
constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked;
|
||||
DG_STATIC_ASSERT(kNumMathThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks");
|
||||
if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) {
|
||||
auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N;
|
||||
auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M;
|
||||
cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr,
|
||||
epilogue_type_t::apply_index_n<TMA_D_BLOCK_N>(n_block_idx * BLOCK_N + in_block_n_offset),
|
||||
scheduler.get_global_idx<kWithGroupOffsetD>(shape_m, BLOCK_M, m_block_idx));
|
||||
cute::tma_store_arrive();
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
}
|
||||
#else
|
||||
if (blockIdx.x == 0 and threadIdx.x == 0)
|
||||
DG_DEVICE_ASSERT(false and "This kernel only support sm_90a");
|
||||
#endif
|
||||
}
|
||||
|
||||
}; // namespace deep_gemm
|
||||
|
||||
#pragma clang diagnostic pop
|
||||
+590
@@ -0,0 +1,590 @@
|
||||
#include <cutlass/arch/barrier.h>
|
||||
#include <cutlass/arch/reg_reconfig.h>
|
||||
|
||||
#include <cute/arch/cluster_sm90.hpp>
|
||||
#include <cute/arch/copy_sm90_desc.hpp>
|
||||
#include <cute/arch/copy_sm90_tma.hpp>
|
||||
|
||||
#include <deep_gemm/common/epilogue_utils.cuh>
|
||||
#include <deep_gemm/common/utils.cuh>
|
||||
#include <deep_gemm/common/scheduler.cuh>
|
||||
#include <deep_gemm/common/sm90_utils.cuh>
|
||||
|
||||
#include <stdexcept>
|
||||
#include <string>
|
||||
|
||||
// LT-PATCH: upstream hard-#defines `__CUDA_ARCH__ 900` here, which forces the wgmma
|
||||
// kernel body on every compile pass and makes this source impossible to place in a
|
||||
// multi-arch fat binary (it emits sm_90-only instructions during e.g. the sm_89 pass
|
||||
// -> ptxas error). Removed so the existing `#if __CUDA_ARCH__ >= 900 ... #else assert
|
||||
// #endif` guard takes effect per-arch: the real body is built only into the sm_90a
|
||||
// cubin, other arches get a host-visible assert stub. The sm_90a pass is unchanged
|
||||
// (nvcc defines __CUDA_ARCH__=900 there regardless).
|
||||
|
||||
namespace deep_gemm {
|
||||
|
||||
using namespace deep_gemm::sm90;
|
||||
|
||||
template <uint32_t kNumFormerIters, uint32_t kGap, uint32_t kEnd, typename func_t>
|
||||
__device__ void dispatch_num_former_iters(uint32_t num_former_iters, const func_t& func) {
|
||||
if (num_former_iters == kNumFormerIters) {
|
||||
func(cute::Int<kNumFormerIters>{});
|
||||
return;
|
||||
}
|
||||
|
||||
if constexpr (kNumFormerIters + kGap <= kEnd)
|
||||
dispatch_num_former_iters<kNumFormerIters + kGap, kGap, kEnd>(num_former_iters, func);
|
||||
}
|
||||
|
||||
template <uint32_t SHAPE_M, uint32_t SHAPE_N, uint32_t SHAPE_K,
|
||||
uint32_t kNumGroups,
|
||||
uint32_t BLOCK_M, uint32_t BLOCK_N, uint32_t BLOCK_K,
|
||||
uint32_t kSwizzleDMode,
|
||||
uint32_t kNumStages, uint32_t kNumLastStages,
|
||||
uint32_t kNumTMAThreads, uint32_t kNumMathThreads,
|
||||
uint32_t kNumTMAMulticast, bool kIsTMAMulticastOnA,
|
||||
uint32_t kNumSMs, GemmType kGemmType,
|
||||
typename epilogue_type_t>
|
||||
__global__ __launch_bounds__(kNumTMAThreads + kNumMathThreads, 1) void
|
||||
sm90_fp8_gemm_1d2d_bias_impl(float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_a,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_b,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_d,
|
||||
const __grid_constant__ cute::TmaDescriptor tensor_map_sfa) {
|
||||
// LT-PATCH: was `__CUDA_ARCH__ >= 900`. Tightened to Hopper-only (< 1000) so that in a
|
||||
// multi-arch fat binary that also targets Blackwell (sm_100/sm_120), this wgmma body is
|
||||
// NOT emitted for those passes (wgmma is sm_90a-only) -- they get the `#else` assert stub
|
||||
// instead. Blackwell dispatches to the SM89 kernel at runtime, so the stub is never run.
|
||||
#if (defined(__CUDA_ARCH__) and (__CUDA_ARCH__ >= 900) and (__CUDA_ARCH__ < 1000)) or defined(__CLION_IDE__)
|
||||
// Scaling checks
|
||||
DG_STATIC_ASSERT(BLOCK_K == 128, "Only support per-128-channel FP8 scaling");
|
||||
DG_STATIC_ASSERT(constexpr_ceil_div(BLOCK_N, BLOCK_K) == 1 or (constexpr_gcd(BLOCK_N, BLOCK_K) == BLOCK_N - BLOCK_K), "Too much B scales in a single block");
|
||||
|
||||
// Types
|
||||
using WGMMA = typename FP8MMASelector<BLOCK_N>::type;
|
||||
using Barrier = cutlass::arch::ClusterTransactionBarrier;
|
||||
DG_STATIC_ASSERT(BLOCK_M % WGMMA::M == 0, "Invalid block size");
|
||||
|
||||
// Overwrite shape constants if the compiler gives
|
||||
shape_m = SHAPE_M != 0 ? SHAPE_M : shape_m;
|
||||
shape_n = SHAPE_N != 0 ? SHAPE_N : shape_n;
|
||||
shape_k = SHAPE_K != 0 ? SHAPE_K : shape_k;
|
||||
|
||||
// Shared memory
|
||||
static constexpr bool kMustUseUniformedScaleB = (BLOCK_K % BLOCK_N == 0);
|
||||
static constexpr uint32_t SMEM_D_SIZE = BLOCK_M * BLOCK_N * sizeof(__nv_bfloat16);
|
||||
static constexpr uint32_t SMEM_A_SIZE_PER_STAGE = BLOCK_M * BLOCK_K * sizeof(__nv_fp8_e4m3);
|
||||
static constexpr uint32_t SMEM_B_SIZE_PER_STAGE = BLOCK_N * BLOCK_K * sizeof(__nv_fp8_e4m3);
|
||||
static constexpr uint32_t SMEM_SFA_SIZE_PER_STAGE = BLOCK_M * sizeof(float);
|
||||
const uint32_t& shape_k_scales = ceil_div(shape_k, BLOCK_K);
|
||||
const uint32_t& smem_sfb_size = align<uint32_t>(shape_k_scales * (kMustUseUniformedScaleB ? 1 : 2) * sizeof(float), sizeof(Barrier));
|
||||
|
||||
// Configs
|
||||
const uint32_t num_total_k_blocks = ceil_div(shape_k, BLOCK_K);
|
||||
const uint32_t warp_idx = __shfl_sync(0xffffffff, threadIdx.x / 32, 0);
|
||||
const uint32_t lane_idx = get_lane_idx();
|
||||
|
||||
// Prefetch TMA descriptors at the very beginning
|
||||
if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
|
||||
cute::prefetch_tma_descriptor(&tensor_map_a);
|
||||
cute::prefetch_tma_descriptor(&tensor_map_b);
|
||||
cute::prefetch_tma_descriptor(&tensor_map_sfa);
|
||||
cute::prefetch_tma_descriptor(&tensor_map_d);
|
||||
}
|
||||
__syncwarp();
|
||||
|
||||
// Align to 1024 bytes for swizzle-128B
|
||||
extern __shared__ __align__(1024) uint8_t smem_buffer[];
|
||||
DG_STATIC_ASSERT(SMEM_D_SIZE % 1024 == 0, "Shared memory of A/B must be aligned to 1024 bytes");
|
||||
|
||||
// Data on shared memory
|
||||
auto smem_d = reinterpret_cast<__nv_bfloat16*>(smem_buffer);
|
||||
auto smem_a = PatternVisitor([&](const uint32_t& i) {
|
||||
return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + i * SMEM_A_SIZE_PER_STAGE);
|
||||
});
|
||||
auto smem_b = PatternVisitor([&](const uint32_t& i) {
|
||||
return reinterpret_cast<__nv_fp8_e4m3*>(smem_buffer + SMEM_D_SIZE + kNumStages * SMEM_A_SIZE_PER_STAGE + i * SMEM_B_SIZE_PER_STAGE);
|
||||
});
|
||||
constexpr uint32_t SMEM_SF_OFFSET = SMEM_D_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
|
||||
auto smem_sfa = PatternVisitor([&](const uint32_t& i) {
|
||||
return reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + i * SMEM_SFA_SIZE_PER_STAGE);
|
||||
});
|
||||
auto smem_sfb = reinterpret_cast<float*>(smem_buffer + SMEM_SF_OFFSET + kNumStages * SMEM_SFA_SIZE_PER_STAGE);
|
||||
|
||||
// Fill barriers
|
||||
auto barrier_start_ptr = reinterpret_cast<Barrier*>(reinterpret_cast<uint8_t*>(smem_sfb) + smem_sfb_size);
|
||||
auto full_barriers = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + i; });
|
||||
auto empty_barriers = PatternVisitor([&](const uint32_t& i) { return barrier_start_ptr + kNumStages + i; });
|
||||
|
||||
// Initialize barriers
|
||||
DG_STATIC_ASSERT(kNumTMAMulticast <= 32, "Too many TMA multicast");
|
||||
if (warp_idx == kNumMathThreads / 32 + 1 and cute::elect_one_sync()) {
|
||||
// NOTES: we always use `lane_idx` to arrive for the `lane_idx`-th CTA in the cluster,
|
||||
// even with TMA multicast disabled, we want to make the behavior aligned
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < kNumStages; ++ i) {
|
||||
full_barriers[i]->init(1);
|
||||
empty_barriers[i]->init(kNumTMAMulticast * kNumMathThreads / 32);
|
||||
}
|
||||
|
||||
// Make initialized barrier visible in async proxy
|
||||
cutlass::arch::fence_barrier_init();
|
||||
}
|
||||
|
||||
// Synchronize all threads to make barrier visible in normal memory model
|
||||
(kNumTMAMulticast > 1) ? cute::cluster_sync() : __syncthreads();
|
||||
|
||||
// Register reconfigurations
|
||||
constexpr uint32_t kNumTMARegisters = 40;
|
||||
constexpr uint32_t kNumMathRegisters = 232;
|
||||
|
||||
// Block scheduler
|
||||
uint32_t m_block_idx, n_block_idx;
|
||||
auto scheduler = Scheduler<kGemmType, BLOCK_M, BLOCK_N, kNumGroups, kNumTMAMulticast, kIsTMAMulticastOnA, kNumSMs>(shape_m, shape_n, shape_k, grouped_layout);
|
||||
|
||||
// Pipeline and TMA phases
|
||||
uint32_t stage_idx = 0, phase = 0;
|
||||
auto advance_pipeline = [&](uint32_t& k_block_idx) {
|
||||
++ k_block_idx;
|
||||
|
||||
// Flip phases only if reach the next first stage
|
||||
stage_idx = stage_idx == kNumStages - 1 ? 0 : stage_idx + 1;
|
||||
phase ^= stage_idx == 0;
|
||||
};
|
||||
|
||||
if (warp_idx >= kNumMathThreads / 32) {
|
||||
// TMA warp-group for loading data
|
||||
cutlass::arch::warpgroup_reg_dealloc<kNumTMARegisters>();
|
||||
|
||||
// NOTES: only one thread (or warp) will be used
|
||||
if (warp_idx == kNumMathThreads / 32 and cute::elect_one_sync()) {
|
||||
// Persistently schedule over blocks
|
||||
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
|
||||
// Assign TMA multicast number into A and B
|
||||
// NOTES: there may be additional odd rows/columns or cases where multicast is not possible.
|
||||
const bool is_tma_multicast_valid = scheduler.is_tma_multicast_valid(m_block_idx);
|
||||
const uint32_t num_tma_multicast_a = (kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
|
||||
const uint32_t num_tma_multicast_b = (not kIsTMAMulticastOnA and is_tma_multicast_valid) ? kNumTMAMulticast : 1;
|
||||
DG_STATIC_ASSERT(kNumTMAMulticast <= 2, "Scheduler does not support > 2 TMA multicast");
|
||||
|
||||
for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
|
||||
// Wait consumer release
|
||||
empty_barriers[stage_idx]->wait(phase ^ 1);
|
||||
|
||||
// Issue TMA A
|
||||
constexpr bool kWithGroupOffsetA = kGemmType == GemmType::MGroupedMasked;
|
||||
auto& full_barrier = *full_barriers[stage_idx];
|
||||
const uint32_t k_idx = k_block_idx * BLOCK_K;
|
||||
tma_copy(&tensor_map_a, reinterpret_cast<uint64_t*>(&full_barrier),
|
||||
smem_a[stage_idx], k_idx, scheduler.get_global_idx<kWithGroupOffsetA>(shape_m, BLOCK_M, m_block_idx),
|
||||
num_tma_multicast_a);
|
||||
tma_copy(&tensor_map_sfa, reinterpret_cast<uint64_t*>(&full_barrier),
|
||||
smem_sfa[stage_idx], m_block_idx * BLOCK_M, scheduler.get_global_idx<kWithGroupOffsetA>(shape_k_scales, 1, k_block_idx),
|
||||
num_tma_multicast_a);
|
||||
|
||||
// Issue TMA B
|
||||
tma_copy(&tensor_map_b, reinterpret_cast<uint64_t*>(&full_barrier),
|
||||
smem_b[stage_idx], k_idx, scheduler.get_global_idx<true>(shape_n, BLOCK_N, n_block_idx, m_block_idx),
|
||||
num_tma_multicast_b);
|
||||
full_barrier.arrive_and_expect_tx(SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE);
|
||||
}
|
||||
}
|
||||
|
||||
// To safely deconstruct distributed shared barriers, we need another round of empty waits
|
||||
if constexpr (kNumTMAMulticast > 1) {
|
||||
for (uint32_t i = 0; i < kNumStages; advance_pipeline(i))
|
||||
empty_barriers[stage_idx]->wait(phase ^ 1);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
// Math warp-groups for WGMMA
|
||||
cutlass::arch::warpgroup_reg_alloc<kNumMathRegisters>();
|
||||
|
||||
// NOTES: use `__shfl_sync` to encourage NVCC to use unified registers
|
||||
const auto math_wg_idx = __shfl_sync(0xffffffff, threadIdx.x / 128, 0);
|
||||
const auto r_0 = warp_idx * 16 + lane_idx / 4, r_1 = r_0 + 8;
|
||||
|
||||
auto a_desc = make_smem_desc(smem_a[0] + math_wg_idx * WGMMA::M * BLOCK_K, 1);
|
||||
auto b_desc = make_smem_desc(smem_b[0], 1);
|
||||
const uint32_t a_desc_lo = __shfl_sync(0xffffffff, a_desc.reg32_[0], 0);
|
||||
const uint32_t b_desc_lo = __shfl_sync(0xffffffff, b_desc.reg32_[0], 0);
|
||||
|
||||
// Persistently schedule over blocks
|
||||
while (scheduler.get_next_block(m_block_idx, n_block_idx)) {
|
||||
// Decide the number of scales B to load
|
||||
DG_TRAP_ONLY_DEVICE_ASSERT(shape_n % 8 == 0);
|
||||
uint32_t num_former_iters = BLOCK_N / 8, num_full_iters = num_former_iters;
|
||||
if constexpr (not kMustUseUniformedScaleB) {
|
||||
num_former_iters = min(BLOCK_N, BLOCK_K - n_block_idx * BLOCK_N % BLOCK_K) / 8;
|
||||
num_full_iters = min(shape_n - n_block_idx * BLOCK_N, BLOCK_N) / 8;
|
||||
}
|
||||
uint32_t num_sfb = shape_k_scales * (num_former_iters >= num_full_iters ? 1 : 2);
|
||||
|
||||
// Load B scales with math warp-groups
|
||||
// NOTES: except the first warp, we want to overlap loading B scales with TMA stores between tasks
|
||||
if (threadIdx.x >= 32) {
|
||||
auto num_previous_lines = scheduler.get_global_idx<true>(ceil_div(shape_n, BLOCK_K), 0, 0, m_block_idx);
|
||||
auto local_sfb = sfb + (num_previous_lines + ((n_block_idx * BLOCK_N) / BLOCK_K)) * shape_k_scales;
|
||||
#pragma unroll
|
||||
for (uint32_t i = threadIdx.x - 32; i < num_sfb; i += kNumMathThreads - 32)
|
||||
st_shared(smem_sfb + i, __ldg(local_sfb + i));
|
||||
}
|
||||
cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
|
||||
|
||||
// Accumulation for WGMMA or CUDA promotion
|
||||
constexpr uint32_t WAVE_BLOCK_M = WGMMA::M * (BLOCK_M <= 64 ? 1 : 2);
|
||||
DG_STATIC_ASSERT(BLOCK_M % WAVE_BLOCK_M == 0, "Invalid block sizes");
|
||||
float accum[WGMMA::kNumAccum], final_accum[WGMMA::kNumAccum * (BLOCK_M / WAVE_BLOCK_M)] = {0};
|
||||
|
||||
// Empty barrier arrival
|
||||
auto empty_barrier_arrive = [&]() {
|
||||
if constexpr (kNumTMAMulticast == 1) {
|
||||
lane_idx == 0 ? empty_barriers[stage_idx]->arrive() : void();
|
||||
} else {
|
||||
auto target_cta = scheduler.is_peer_cta_alive ? lane_idx : cute::block_rank_in_cluster();
|
||||
lane_idx < kNumTMAMulticast ? empty_barriers[stage_idx]->arrive(target_cta) : void();
|
||||
}
|
||||
};
|
||||
|
||||
// Skip useless computations
|
||||
if (scheduler.is_computation_valid(m_block_idx, math_wg_idx * WGMMA::M)) {
|
||||
// The compiler must know the dynamic variable `num_former_iters`'s real value
|
||||
constexpr bool kShouldOptimize = BLOCK_K / constexpr_gcd(BLOCK_K, BLOCK_N) <= 4 and not kMustUseUniformedScaleB;
|
||||
constexpr uint32_t kGap = constexpr_gcd(BLOCK_K, BLOCK_N) / 8;
|
||||
constexpr uint32_t kEnd = kShouldOptimize ? BLOCK_K / 8 : 0;
|
||||
|
||||
// Dispatch `num_former_iters` and launch MMAs
|
||||
dispatch_num_former_iters<0, kGap, kEnd>(kShouldOptimize ? num_former_iters : 0, [&](auto _) {
|
||||
#pragma unroll 8
|
||||
for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
|
||||
const auto& a_desc_base_lo = a_desc_lo + stage_idx * (SMEM_A_SIZE_PER_STAGE / 16);
|
||||
const auto& b_desc_base_lo = b_desc_lo + stage_idx * (SMEM_B_SIZE_PER_STAGE / 16);
|
||||
|
||||
// Read B scales
|
||||
float scale_b_0 = ld_shared(smem_sfb + k_block_idx), scale_b_1;
|
||||
// NOTES: even some blocks do not need to read the second row, but we still load one to align with other blocks
|
||||
if constexpr (not kMustUseUniformedScaleB)
|
||||
scale_b_1 = ld_shared(smem_sfb + k_block_idx + shape_k_scales);
|
||||
|
||||
// Wait TMA arrivals
|
||||
full_barriers[stage_idx]->wait(phase);
|
||||
|
||||
// TODO: remove some useless computation for unaligned Ms
|
||||
#pragma unroll
|
||||
for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
|
||||
auto m_offset = local_idx * WAVE_BLOCK_M;
|
||||
|
||||
// Read A scales
|
||||
// NOTES: all shared memory read must be prior to `warpgroup_arrive` to avoid next scheduled block polluting the results
|
||||
auto scale_a_0 = ld_shared(smem_sfa[stage_idx] + r_0 + m_offset);
|
||||
auto scale_a_1 = ld_shared(smem_sfa[stage_idx] + r_1 + m_offset);
|
||||
|
||||
// Commit WGMMA instructions
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
|
||||
warpgroup_fence_operand(accum[i]);
|
||||
warpgroup_arrive();
|
||||
#pragma unroll
|
||||
for (uint32_t k = 0; k < BLOCK_K / WGMMA::K; ++ k) {
|
||||
a_desc.reg32_[0] = a_desc_base_lo + (m_offset * BLOCK_K + k * WGMMA::K) / 16;
|
||||
b_desc.reg32_[0] = b_desc_base_lo + k * WGMMA::K / 16;
|
||||
WGMMA::wgmma(a_desc, b_desc, accum, k);
|
||||
}
|
||||
warpgroup_commit_batch();
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < WGMMA::kNumAccum; ++ i)
|
||||
warpgroup_fence_operand(accum[i]);
|
||||
warpgroup_wait<0>();
|
||||
|
||||
// Notify barrier arrival at the last warpgroup wave
|
||||
if (local_idx == BLOCK_M / WAVE_BLOCK_M - 1)
|
||||
empty_barrier_arrive();
|
||||
|
||||
// Promote with scales
|
||||
// NOTES: making it as predicates is very important for performance, comparing to two loops
|
||||
float scale_0_0 = scale_a_0 * scale_b_0, scale_1_0 = scale_a_1 * scale_b_0;
|
||||
float scale_0_1, scale_1_1;
|
||||
if constexpr (not kMustUseUniformedScaleB)
|
||||
scale_0_1 = scale_a_0 * scale_b_1, scale_1_1 = scale_a_1 * scale_b_1;
|
||||
|
||||
auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
|
||||
#pragma unroll
|
||||
for (uint32_t i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
|
||||
// NOTES: for unrolled `num_former_iters` cases, we expect the compiler to automatically make it a constant
|
||||
bool predicate = kMustUseUniformedScaleB or i < num_former_iters;
|
||||
shifted_accum[i * 4 + 0] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 0];
|
||||
shifted_accum[i * 4 + 1] += (predicate ? scale_0_0 : scale_0_1) * accum[i * 4 + 1];
|
||||
shifted_accum[i * 4 + 2] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 2];
|
||||
shifted_accum[i * 4 + 3] += (predicate ? scale_1_0 : scale_1_1) * accum[i * 4 + 3];
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
} else {
|
||||
#pragma unroll
|
||||
for (uint32_t k_block_idx = 0; k_block_idx < num_total_k_blocks; advance_pipeline(k_block_idx)) {
|
||||
full_barriers[stage_idx]->wait(phase);
|
||||
empty_barrier_arrive();
|
||||
}
|
||||
}
|
||||
|
||||
// TMA checks
|
||||
constexpr uint32_t kNumElemBytes = sizeof(nv_bfloat16);
|
||||
constexpr uint32_t TMA_D_BLOCK_N = kSwizzleDMode == 0 ? BLOCK_N : (kSwizzleDMode / kNumElemBytes);
|
||||
constexpr uint32_t WGMMA_M_PER_WARP = WGMMA::M / 4;
|
||||
DG_STATIC_ASSERT(BLOCK_M % 8 == 0, "Invalid swizzling atom");
|
||||
DG_STATIC_ASSERT(BLOCK_N % TMA_D_BLOCK_N == 0 and BLOCK_N / TMA_D_BLOCK_N <= 32,
|
||||
"Unaligned TMA store or too many TMA store instructions");
|
||||
DG_STATIC_ASSERT(TMA_D_BLOCK_N % 8 == 0, "Invalid TMA block N");
|
||||
// Wait last TMA store to be finished
|
||||
float* bias_ptr = bias + n_block_idx*BLOCK_N + (lane_idx % 4) * 2;
|
||||
#pragma unroll
|
||||
for(uint32_t local_idx=0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx){
|
||||
auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
|
||||
#pragma unroll
|
||||
for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
|
||||
shifted_accum[4*i + 0] += bias_ptr[8*i + 0];
|
||||
shifted_accum[4*i + 1] += bias_ptr[8*i + 1];
|
||||
shifted_accum[4*i + 2] += bias_ptr[8*i + 0];
|
||||
shifted_accum[4*i + 3] += bias_ptr[8*i + 1];
|
||||
}
|
||||
}
|
||||
|
||||
if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N)
|
||||
cute::tma_store_wait<0>();
|
||||
cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
|
||||
|
||||
// Write back to shared memory using STSM and issue TMA stores
|
||||
DG_STATIC_ASSERT(WGMMA::kNumAccum % 4 == 0, "Invalid STSM x2 vectorization");
|
||||
#pragma unroll
|
||||
for (uint32_t local_idx = 0; local_idx < BLOCK_M / WAVE_BLOCK_M; ++ local_idx) {
|
||||
auto m_offset = local_idx * WAVE_BLOCK_M;
|
||||
auto shifted_accum = final_accum + WGMMA::kNumAccum * local_idx;
|
||||
#pragma unroll
|
||||
for (auto i = 0; i < WGMMA::kNumAccum / 4; ++ i) {
|
||||
// Swizzle or padding into the correct address
|
||||
uint8_t* smem_ptr = nullptr;
|
||||
if constexpr (kSwizzleDMode > 0) {
|
||||
// Calculate the swizzling atom offset and in-atom offset
|
||||
constexpr uint32_t kNumBankGroupBytes = 16;
|
||||
auto atom_offset = i / (TMA_D_BLOCK_N / 8), in_atom_offset = i % (TMA_D_BLOCK_N / 8);
|
||||
|
||||
// Calculate the index of the bank group to be written in the atom
|
||||
auto bank_group_index = in_atom_offset + lane_idx * (kSwizzleDMode / kNumBankGroupBytes);
|
||||
|
||||
// Reshape the atom in another view and swizzle
|
||||
// - original: `(BLOCK_M, kSwizzleDMode / kNumBankGroupBytes)`
|
||||
// - new: `(BLOCK_M * kSwizzleDMode / kNumBankGroupBytes / 8, 8)`
|
||||
constexpr bool kHasShortcut = (kSwizzleDMode / kNumBankGroupBytes) == 8;
|
||||
auto row = kHasShortcut ? (in_atom_offset / 8 + lane_idx) : (bank_group_index / 8);
|
||||
auto col = kHasShortcut ? (in_atom_offset) : (bank_group_index % 8);
|
||||
col ^= row % (kSwizzleDMode / 16);
|
||||
|
||||
// Add back into the base pointer
|
||||
// NOTES: think twice before modifying this, as changes may affect the number of instructions
|
||||
smem_ptr = reinterpret_cast<uint8_t*>(smem_d) + // Base pointer
|
||||
warp_idx * (WGMMA_M_PER_WARP * kSwizzleDMode) + // Warp offset
|
||||
m_offset * kSwizzleDMode + // Wave offset
|
||||
atom_offset * BLOCK_M * kSwizzleDMode + // Swizzle atom offset (constants)
|
||||
row * (kNumBankGroupBytes * 8) + col * kNumBankGroupBytes; // In-atom offset
|
||||
} else {
|
||||
// No swizzling, just padding
|
||||
smem_ptr = reinterpret_cast<uint8_t*>(smem_d + (m_offset + warp_idx * WGMMA_M_PER_WARP + lane_idx) * BLOCK_N + i * 8);
|
||||
}
|
||||
|
||||
// NOTES: only 16 lanes' addresses are used
|
||||
SM90_U32x2_STSM_N<nv_bfloat162>::copy(
|
||||
__float22bfloat162_rn({shifted_accum[i * 4 + 0], shifted_accum[i * 4 + 1]}),
|
||||
__float22bfloat162_rn({shifted_accum[i * 4 + 2], shifted_accum[i * 4 + 3]}),
|
||||
smem_ptr
|
||||
);
|
||||
}
|
||||
}
|
||||
cute::tma_store_fence();
|
||||
cutlass::arch::NamedBarrier::sync(kNumMathThreads, 0);
|
||||
|
||||
// Use TMA store to write back to global memory
|
||||
// TODO: compatible with FP32 output
|
||||
constexpr bool kWithGroupOffsetD = kGemmType == GemmType::MGroupedMasked;
|
||||
DG_STATIC_ASSERT(kNumMathThreads >= BLOCK_N / TMA_D_BLOCK_N, "Too many TMA blocks");
|
||||
if (threadIdx.x < BLOCK_N / TMA_D_BLOCK_N) {
|
||||
auto in_block_n_offset = threadIdx.x * TMA_D_BLOCK_N;
|
||||
auto smem_ptr = smem_d + in_block_n_offset * BLOCK_M;
|
||||
cute::SM90_TMA_STORE_2D::copy(&tensor_map_d, smem_ptr,
|
||||
epilogue_type_t::apply_index_n<TMA_D_BLOCK_N>(n_block_idx * BLOCK_N + in_block_n_offset),
|
||||
scheduler.get_global_idx<kWithGroupOffsetD>(shape_m, BLOCK_M, m_block_idx));
|
||||
cute::tma_store_arrive();
|
||||
}
|
||||
__syncwarp();
|
||||
}
|
||||
}
|
||||
#else
|
||||
if (blockIdx.x == 0 and threadIdx.x == 0)
|
||||
DG_DEVICE_ASSERT(false and "This kernel only support sm_90a");
|
||||
#endif
|
||||
}
|
||||
|
||||
static cudaLaunchConfig_t construct_launch_config(const cudaStream_t& stream, const int& smem_size,
|
||||
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
|
||||
|
||||
cudaLaunchConfig_t config;
|
||||
config.gridDim = grid_dim;
|
||||
config.blockDim = block_dim;
|
||||
config.dynamicSmemBytes = smem_size;
|
||||
config.stream = stream;
|
||||
config.numAttrs = 0;
|
||||
config.attrs = nullptr;
|
||||
|
||||
// NOTES: must use `static` or the `attr` will be deconstructed
|
||||
static cudaLaunchAttribute attr;
|
||||
if (cluster_dim > 1) {
|
||||
attr.id = cudaLaunchAttributeClusterDimension;
|
||||
attr.val.clusterDim = {static_cast<unsigned>(cluster_dim), 1, 1};
|
||||
config.attrs = &attr;
|
||||
config.numAttrs = 1;
|
||||
}
|
||||
return config;
|
||||
}
|
||||
|
||||
|
||||
// static auto launch_kernel(auto kernel, const cudaLaunchConfig_t& config, float* sfb, float* bias, int* grouped_layout,
|
||||
// uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
// const CUtensorMap tensor_map_a,
|
||||
// const CUtensorMap tensor_map_b,
|
||||
// const CUtensorMap tensor_map_d,
|
||||
// const CUtensorMap tensor_map_sfa) {
|
||||
// // void* ptr_args[] = {&sfb, &bias, &grouped_layout, &shape_m, &shape_n, &shape_k, &tensor_map_a, &tensor_map_b, &tensor_map_d, &tensor_map_sfa};
|
||||
// return
|
||||
// }
|
||||
|
||||
|
||||
template<int N, int K>
|
||||
void sm90_fp8_gemm_1d2d_bias_launch(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa){
|
||||
dim3 grid{num_sms, 1, 1};
|
||||
dim3 block{num_threads, 1, 1};
|
||||
const auto config = construct_launch_config(stream, smem_size, grid, block, cluster_dim);
|
||||
if(num_sms == 132){
|
||||
auto kernel = &sm90_fp8_gemm_1d2d_bias_impl<0, N, K, 1, 256, 128, 128, 128, 3, (K / 128) % 3, 128, 256, 2, true, 132, GemmType::Normal, EpilogueIdentity>;
|
||||
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
|
||||
cudaLaunchKernelEx(&config, kernel, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
|
||||
} else if(num_sms == 116) {
|
||||
auto kernel = &sm90_fp8_gemm_1d2d_bias_impl<0, N, K, 1, 256, 128, 128, 128, 3, (K / 128) % 3, 128, 256, 2, true, 116, GemmType::Normal, EpilogueIdentity>;
|
||||
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
|
||||
cudaLaunchKernelEx(&config, kernel, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
|
||||
} else if (num_sms == 100) {
|
||||
auto kernel = &sm90_fp8_gemm_1d2d_bias_impl<0, N, K, 1, 256, 128, 128, 128, 3, (K / 128) % 3, 128, 256, 2, true, 100, GemmType::Normal, EpilogueIdentity>;
|
||||
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size);
|
||||
cudaLaunchKernelEx(&config, kernel, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
|
||||
} else {
|
||||
// The supported SM counts are exactly the branches above (the only kernels
|
||||
// instantiated). Fail loudly instead of falling through with no launch,
|
||||
// which would leave the output buffer uninitialized.
|
||||
throw std::runtime_error("Unsupported num_sms=" + std::to_string(num_sms)
|
||||
+ " (blockwise SM90 GEMM is built for 132, 116, and 100 SMs)");
|
||||
}
|
||||
|
||||
// launch_kernel(kernel, config, sfb, bias, grouped_layout, shape_m, shape_n, shape_k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);
|
||||
}
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<2048, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<4096, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<8192, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<16384, 2048>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<2048, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<4096, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<8192, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<16384, 4096>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<2048, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<4096, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<8192, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<16384, 8192>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<2048, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<4096, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<8192, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
template void sm90_fp8_gemm_1d2d_bias_launch<16384, 16384>(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
}; // namespace deep_gemm
|
||||
@@ -0,0 +1,287 @@
|
||||
#include "cutlass/cutlass.h"
|
||||
#include "cutlass/layout/layout.h"
|
||||
#include <cute/tensor.hpp>
|
||||
#include <c10/cuda/CUDAException.h>
|
||||
#include <torch/extension.h>
|
||||
#include <torch/python.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <iostream>
|
||||
|
||||
#include "kernel_traits.cuh"
|
||||
#include "static_switch.h"
|
||||
|
||||
namespace sm89{
|
||||
using namespace cute;
|
||||
|
||||
__device__ static void copy_1d(float* gmem_src, float* smem_dst)
|
||||
{
|
||||
|
||||
uint32_t smem_int_ptr = cast_smem_ptr_to_uint((void*)smem_dst);
|
||||
asm volatile("cp.async.ca.shared.global.L2::128B [%0], [%1], %2;\n"
|
||||
:: "r"(smem_int_ptr),
|
||||
"l"(gmem_src),
|
||||
"n"(sizeof(float)));
|
||||
}
|
||||
|
||||
template <typename KernelTraits=gemm_traits<128, 256, 2, 4096, 2, 4, true, half_t, bfloat16_t>>
|
||||
__global__ void gemm_fp8_kernel(float_e4m3_t* Aptr, float* sfa, float_e4m3_t* Bptr, float* sfb, float* bias_ptr, void* out, int M, int N, int K, int TMA_ALIGNED_M){
|
||||
using output_t = typename KernelTraits::out_t;
|
||||
using SmemLayoutA = typename KernelTraits::SmemLayoutA;
|
||||
using SmemLayoutB = typename KernelTraits::SmemLayoutB;
|
||||
using SmemLayoutC = typename KernelTraits::SmemLayoutC;
|
||||
|
||||
constexpr int BM = KernelTraits::BM;
|
||||
constexpr int BN = KernelTraits::BN;
|
||||
constexpr int BK = KernelTraits::BK;
|
||||
constexpr int Ksfa = KernelTraits::KSF;
|
||||
constexpr bool has_bias = KernelTraits::HasBias;
|
||||
extern __shared__ float smem_[];
|
||||
float *bias_shm = smem_;
|
||||
float *sfa_shm = reinterpret_cast<float*>(bias_shm + cosize(typename KernelTraits::SmemLayoutBias{}));
|
||||
|
||||
output_t* C_shm = reinterpret_cast<output_t*>(sfa_shm + cosize(typename KernelTraits::SmemLayoutSFA{}));
|
||||
float_e4m3_t* A_shm = reinterpret_cast<float_e4m3_t*>(sfa_shm + cosize(typename KernelTraits::SmemLayoutSFA{}));
|
||||
float_e4m3_t* B_shm = reinterpret_cast<float_e4m3_t*>(A_shm + cosize(SmemLayoutA{}));
|
||||
|
||||
int idx = threadIdx.x;
|
||||
int ix = blockIdx.x;
|
||||
int iy = blockIdx.y;
|
||||
// sfa += BM * iy;
|
||||
sfb += KernelTraits::NUM_SFB_PER_STEP * ix * Ksfa;
|
||||
|
||||
|
||||
output_t* Cptr = reinterpret_cast<output_t*>(out);
|
||||
|
||||
Tensor A = make_tensor(make_gmem_ptr(Aptr), make_shape(M, K), make_stride(K, Int<1>{}));
|
||||
Tensor B = make_tensor(make_gmem_ptr(Bptr), make_shape(N, K), make_stride(K, Int<1>{}));
|
||||
Tensor D = make_tensor(make_gmem_ptr(Cptr), make_shape(M, N), make_stride(N, Int<1>{}));
|
||||
Tensor SFA = make_tensor(make_gmem_ptr(sfa), make_shape(M, Ksfa), make_stride(Int<1>{}, TMA_ALIGNED_M));
|
||||
|
||||
Tensor gA = local_tile(A, make_tile(Int<BM>{}, Int<BK>{}), make_coord(iy, _));
|
||||
Tensor gB = local_tile(B, make_tile(Int<BN>{}, Int<BK>{}), make_coord(ix, _));
|
||||
Tensor gD = local_tile(D, make_tile(Int<BM>{}, Int<BN>{}), make_coord(iy, ix));
|
||||
Tensor gSFA = local_tile(SFA, make_tile(Int<BM>{}, Int<1>{}), make_coord(iy, _));
|
||||
|
||||
auto sBias = make_tensor(make_smem_ptr(bias_shm), typename KernelTraits::SmemLayoutBias{});
|
||||
if constexpr (has_bias){
|
||||
Tensor Bias = make_tensor(make_gmem_ptr(bias_ptr), make_shape(_1{}, N), make_stride(N, Int<1>{}));
|
||||
Tensor gBias = local_tile(Bias, make_tile(Int<1>{}, Int<BN>{}), make_coord(_, ix));
|
||||
typename KernelTraits::G2SBiasCopy g2s_bias_copy;
|
||||
auto g2s_bias_thr_copy = g2s_bias_copy.get_slice(idx);
|
||||
auto tCBiasgBias = g2s_bias_thr_copy.partition_S(gBias);
|
||||
auto tCBiassBias = g2s_bias_thr_copy.partition_D(sBias);
|
||||
if(idx < BN){
|
||||
copy_1d((float*)&gBias(0) + idx, (float*)&sBias(0) + idx);
|
||||
}
|
||||
}
|
||||
|
||||
auto sSFA = make_tensor(make_smem_ptr(sfa_shm), typename KernelTraits::SmemLayoutSFA{});
|
||||
auto sA = make_tensor(make_smem_ptr(A_shm), SmemLayoutA{});
|
||||
auto sB = make_tensor(make_smem_ptr(B_shm), SmemLayoutB{});
|
||||
|
||||
typename KernelTraits::MMATile tiled_mma;
|
||||
auto thr_mma = tiled_mma.get_slice(threadIdx.x);
|
||||
|
||||
auto tCrA = thr_mma.partition_fragment_A(gA(_, _, 0));
|
||||
auto tCrB = thr_mma.partition_fragment_B(gB(_, _, 0));
|
||||
auto tCrD = thr_mma.partition_fragment_C(gD);
|
||||
clear(tCrD);
|
||||
auto tCrD_fp32 = make_tensor_like<float>(tCrD);
|
||||
clear(tCrD_fp32);
|
||||
|
||||
typename KernelTraits::G2STiledCopy g2s_tiled_copy;
|
||||
auto g2s_thr_copy = g2s_tiled_copy.get_slice(idx);
|
||||
auto tAgA_copy = g2s_thr_copy.partition_S(gA);
|
||||
auto tAsA_copy = g2s_thr_copy.partition_D(sA);
|
||||
auto tBgB_copy = g2s_thr_copy.partition_S(gB);
|
||||
auto tBsB_copy = g2s_thr_copy.partition_D(sB);
|
||||
|
||||
auto s2r_tiled_copy_a = make_tiled_copy_A(typename KernelTraits::S2RCopyAtomA{}, tiled_mma);
|
||||
auto s2r_thr_copy_a = s2r_tiled_copy_a.get_slice(idx);
|
||||
auto tAsA = s2r_thr_copy_a.partition_S(sA);
|
||||
auto tCrA_view = s2r_thr_copy_a.retile_D(tCrA);
|
||||
|
||||
|
||||
auto s2r_tiled_copy_b = make_tiled_copy_B(typename KernelTraits::S2RCopyAtomB{}, tiled_mma);
|
||||
auto s2r_thr_copy_b = s2r_tiled_copy_b.get_slice(idx);
|
||||
auto tBsB = s2r_thr_copy_b.partition_S(sB);
|
||||
auto tCrB_view = s2r_thr_copy_b.retile_D(tCrB);
|
||||
|
||||
auto cA = make_identity_tensor(make_shape(size<0>(sA), size<1>(sA)));
|
||||
auto tAcA = g2s_thr_copy.partition_S(cA);
|
||||
int residual = M - iy*BM;
|
||||
|
||||
int itile_to_read = 0;
|
||||
int ismem_read = 0;
|
||||
int ismem_write = 0;
|
||||
int ismem_read_sfa = 0;
|
||||
constexpr int kStages = KernelTraits::KStages;
|
||||
|
||||
#pragma unroll
|
||||
for(int istage=0; istage<kStages - 1; ++istage){
|
||||
for (size_t m = 0; m < size<1>(tAsA_copy); m++)
|
||||
{
|
||||
for (size_t k = 0; k < size<2>(tAsA_copy); k++)
|
||||
{
|
||||
if(get<0>(tAcA(0, m, k)) < residual){
|
||||
cute::copy(g2s_tiled_copy, tAgA_copy(_, m, k, istage), tAsA_copy(_, m, k, istage));
|
||||
}
|
||||
}
|
||||
}
|
||||
if(idx < KernelTraits::THREADS_SFA_COPY && (BM * iy + idx * KernelTraits::SFA_ELEMS_PER_COPY < M)) {
|
||||
copy_1d((float*)&gSFA(0, 0, istage) + idx*KernelTraits::SFA_ELEMS_PER_COPY, (float*)&sSFA(0, istage) + idx*KernelTraits::SFA_ELEMS_PER_COPY);
|
||||
}
|
||||
cute::copy(g2s_tiled_copy, tBgB_copy(_, _, _, istage), tBsB_copy(_, _, _, istage));
|
||||
cp_async_fence();
|
||||
++itile_to_read;
|
||||
++ismem_write;
|
||||
}
|
||||
|
||||
cp_async_wait<kStages - 2>();
|
||||
__syncthreads();
|
||||
|
||||
cute::copy(s2r_tiled_copy_a, tAsA(_, _, 0, ismem_read), tCrA_view(_, _, 0));
|
||||
cute::copy(s2r_tiled_copy_b, tBsB(_, _, 0, ismem_read), tCrB_view(_, _, 0));
|
||||
|
||||
static constexpr int nk = size<2>(tCrA);
|
||||
auto sfa_tv = typename KernelTraits::SFAThreadLayout{};
|
||||
static constexpr int NTILES = KernelTraits::NTiles;
|
||||
#pragma unroll
|
||||
for(int itile = 0; itile < NTILES; itile++){
|
||||
clear(tCrD);
|
||||
#pragma unroll
|
||||
for(int ik = 0; ik < nk; ik++){
|
||||
int ik_next = (ik + 1) % nk;
|
||||
if(ik == nk - 1) {
|
||||
cp_async_wait<kStages - 2>();
|
||||
__syncthreads();
|
||||
ismem_read = (ismem_read + 1) % kStages;
|
||||
}
|
||||
cute::copy(s2r_tiled_copy_a, tAsA(_, _, ik_next, ismem_read), tCrA_view(_, _, ik_next));
|
||||
cute::copy(s2r_tiled_copy_b, tBsB(_, _, ik_next, ismem_read), tCrB_view(_, _, ik_next));
|
||||
if(ik == 0){
|
||||
if(itile_to_read < NTILES){
|
||||
for (size_t m = 0; m < size<1>(tAsA_copy); m++)
|
||||
{
|
||||
for (size_t k = 0; k < size<2>(tAsA_copy); k++)
|
||||
{
|
||||
if(get<0>(tAcA(0, m, k)) < residual){
|
||||
cute::copy(g2s_tiled_copy, tAgA_copy(_, m, k, itile_to_read), tAsA_copy(_, m, k, ismem_write));
|
||||
}
|
||||
}
|
||||
}
|
||||
cute::copy(g2s_tiled_copy, tBgB_copy(_, _, _, itile_to_read), tBsB_copy(_, _, _, ismem_write));
|
||||
if(idx < KernelTraits::THREADS_SFA_COPY && (BM * iy + idx * KernelTraits::SFA_ELEMS_PER_COPY < M)) {
|
||||
copy_1d((float*)&gSFA(0, 0, itile_to_read) + idx * KernelTraits::SFA_ELEMS_PER_COPY, (float*)&sSFA(0, ismem_write) + idx*KernelTraits::SFA_ELEMS_PER_COPY);
|
||||
}
|
||||
++itile_to_read;
|
||||
ismem_write = (ismem_write + 1) % kStages;
|
||||
}
|
||||
cp_async_fence();
|
||||
}
|
||||
cute::gemm(tiled_mma, tCrD, tCrA(_, _, ik), tCrB(_, _, ik), tCrD);
|
||||
}
|
||||
|
||||
int sf_ind = itile / KernelTraits::TILES_PER_BLOCK;
|
||||
float sfb_val = sfb[sf_ind];
|
||||
#pragma unroll
|
||||
for(int i = 0; i < size<1>(tCrD); i++){ // (MMA, MMA_M, MMA_N) = (4, 4, 4)
|
||||
float sfa_val_1 = sSFA(sfa_tv(idx) + i * KernelTraits::MMA_WARP_M, ismem_read_sfa);
|
||||
float sfa_val_2 = sSFA(sfa_tv(idx) + 8 + i * KernelTraits::MMA_WARP_M, ismem_read_sfa);
|
||||
#pragma unroll
|
||||
for(int j = 0; j < size<2>(tCrD); j++){
|
||||
tCrD_fp32(0, i, j) += sfa_val_1 * sfb_val * float(tCrD(0, i, j));
|
||||
tCrD_fp32(1, i, j) += sfa_val_1 * sfb_val * float(tCrD(1, i, j));
|
||||
tCrD_fp32(2, i, j) += sfa_val_2 * sfb_val * float(tCrD(2, i, j));
|
||||
tCrD_fp32(3, i, j) += sfa_val_2 * sfb_val * float(tCrD(3, i, j));
|
||||
}
|
||||
}
|
||||
ismem_read_sfa = (ismem_read_sfa + 1) % kStages;
|
||||
__syncthreads();
|
||||
}
|
||||
|
||||
auto tCrBias = make_tensor<float>(Layout<Shape<_2, Int<size<2>(tCrD_fp32)>>>{});
|
||||
auto bias_threads = typename KernelTraits::BiasThreadLayout{};
|
||||
if constexpr (has_bias){
|
||||
#pragma unroll
|
||||
for(int i = 0; i<size<2>(tCrD_fp32); i++){
|
||||
tCrBias(0, i) = sBias(bias_threads(idx) + i * KernelTraits::MMA_WARP_N);
|
||||
tCrBias(1, i) = sBias(1 + bias_threads(idx) + i * KernelTraits::MMA_WARP_N);
|
||||
}
|
||||
#pragma unroll
|
||||
for(int i = 0; i<size<1>(tCrD_fp32); i++){
|
||||
#pragma unroll
|
||||
for (int j = 0; j < size<2>(tCrD_fp32) ; j++)
|
||||
{
|
||||
tCrD_fp32(0, i, j) += tCrBias(0, j);
|
||||
tCrD_fp32(1, i, j) += tCrBias(1, j);
|
||||
tCrD_fp32(2, i, j) += tCrBias(0, j);
|
||||
tCrD_fp32(3, i, j) += tCrBias(1, j);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
auto sC = make_tensor(make_smem_ptr(C_shm), SmemLayoutC{});
|
||||
auto r2s_tiled_copy_c = make_tiled_copy_C(typename KernelTraits::R2SCopyAtomC{}, tiled_mma);
|
||||
auto r2s_thr_copy_c = r2s_tiled_copy_c.get_slice(idx);
|
||||
auto tCrC_r2s = r2s_thr_copy_c.retile_S(tCrD_fp32);
|
||||
auto tCsC_r2s = r2s_thr_copy_c.partition_D(sC);
|
||||
|
||||
typename KernelTraits::S2GCopyC s2g_tiled_copy_c;
|
||||
auto s2g_thr_copy_c = s2g_tiled_copy_c.get_thread_slice(idx);
|
||||
auto tCsC_s2g = s2g_thr_copy_c.partition_S(sC);
|
||||
auto tCgC_s2g = s2g_thr_copy_c.partition_D(gD);
|
||||
|
||||
int pipe = size<2>(tCsC_r2s);
|
||||
|
||||
auto cC = make_identity_tensor(make_shape(size<0>(gD), size<1>(gD)));
|
||||
auto tCcC = s2g_thr_copy_c.partition_D(cC);
|
||||
|
||||
for(int i = 0; i< size<1>(tCrC_r2s); i++){
|
||||
for(int j = 0; j < size<2>(tCrC_r2s); j+=pipe){
|
||||
for(int step = 0; step < pipe; ++step){
|
||||
auto fragment = make_tensor_like<output_t>(tCrC_r2s(_, i, j + step));
|
||||
cute::copy(tCrC_r2s(_, i, j + step), fragment);
|
||||
cute::copy(r2s_tiled_copy_c, fragment, tCsC_r2s(_, 0, step));
|
||||
}
|
||||
__syncthreads();
|
||||
if (get<0>(tCcC(0, i, j / pipe)) < residual){
|
||||
cute::copy(s2g_tiled_copy_c, tCsC_s2g(_, 0, 0), tCgC_s2g(_, i, j / pipe));
|
||||
}
|
||||
__syncthreads();
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
template <bool has_bias, typename accum_type>
|
||||
void fp8_kernel_launch(void* Aptr, void* sfa, void* Bptr, void* sfb, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream) {
|
||||
int TMA_ALIGNED_M = ((M + sizeof(float) - 1) / sizeof(float)) * sizeof(float); // SIZEOF(float) = 4
|
||||
BLOCK_K_SWITCH(K_, M_SWITCH(
|
||||
using KernelTraits = gemm_traits<BM, BN, 3, K_, WARP_ROW, WARP_COL, has_bias, accum_type, bfloat16_t>;
|
||||
auto kernel = &gemm_fp8_kernel<KernelTraits>;
|
||||
int BX = (N + KernelTraits::BN - 1) / KernelTraits::BN;
|
||||
int BY = (M + KernelTraits::BM - 1) / KernelTraits::BM;
|
||||
dim3 block(KernelTraits::NUM_THREADS);
|
||||
dim3 gridDim(BX, BY);
|
||||
cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, KernelTraits::SmemSize);
|
||||
kernel<<<gridDim, KernelTraits::NUM_THREADS, KernelTraits::SmemSize, stream>>>((float_e4m3_t*)Aptr, (float*)sfa, (float_e4m3_t*)Bptr, (float*)sfb, (float*)bias_ptr, out, M, N, K, TMA_ALIGNED_M);
|
||||
C10_CUDA_KERNEL_LAUNCH_CHECK();))
|
||||
}
|
||||
|
||||
template<bool use_fast_accum>
|
||||
void fp8_bias_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream){
|
||||
using accum_type = std::conditional_t<use_fast_accum, half_t, float>;
|
||||
fp8_kernel_launch<true, accum_type>(Aptr, SFA, Bptr, SFB, bias_ptr, out, M, N, K, stream);
|
||||
}
|
||||
// template<bool use_fast_accum>
|
||||
// void fp8_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream){
|
||||
// using accum_type = std::conditional_t<use_fast_accum, half_t, float>;
|
||||
// BLOCK_K_SWITCH(num_acc_upcast_steps, fp8_kernel_launch<false, num_acc_upcast_steps, accum_type>(Aptr, SFA, Bptr, SFB, nullptr, out, M, N, K, stream);)
|
||||
// }
|
||||
|
||||
// template void fp8_gemm_cuda<true>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream);
|
||||
// template void fp8_gemm_cuda<false>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* out, int M, int N, int K, cudaStream_t stream);
|
||||
|
||||
template void fp8_bias_gemm_cuda<true>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
|
||||
template void fp8_bias_gemm_cuda<false>(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
|
||||
}; // namespace sm89
|
||||
@@ -0,0 +1,121 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/tensor.hpp>
|
||||
|
||||
#include <cutlass/cutlass.h>
|
||||
#include <cutlass/layout/layout.h>
|
||||
#include <cutlass/numeric_types.h>
|
||||
|
||||
#include "mma_sm89_fp16.hpp"
|
||||
#include "mma_traits_sm89_fp16.hpp"
|
||||
|
||||
using namespace cute;
|
||||
|
||||
template<int BYTES> struct BytesToType {};
|
||||
template<> struct BytesToType<4> {
|
||||
using Type = uint32_t;
|
||||
static_assert(sizeof(Type) == 4);
|
||||
};
|
||||
template<> struct BytesToType<2> {
|
||||
using Type = uint16_t;
|
||||
static_assert(sizeof(Type) == 2);
|
||||
};
|
||||
|
||||
template<int BM_, int BN_, int KStages_, int K_, int WARP_ROW_=2, int WARP_COL_=2, bool HasBias_=false, typename accum_t_=cutlass::half_t, typename out_t_=cutlass::bfloat16_t>
|
||||
struct gemm_traits {
|
||||
static constexpr int BLOCK_SIZE = 128;
|
||||
static constexpr int K = K_;
|
||||
static constexpr int BM = BM_;
|
||||
static constexpr int BN = BN_;
|
||||
static constexpr int BK = 128;
|
||||
static constexpr int TILES_PER_BLOCK = BLOCK_SIZE / BK;
|
||||
static constexpr int NUM_SFB_PER_STEP = BN / BLOCK_SIZE;
|
||||
static constexpr int NTiles = K / BK;
|
||||
static constexpr int KSF = K / BLOCK_SIZE;
|
||||
static constexpr int KStages = KStages_;
|
||||
static constexpr int WARP_ROW = WARP_ROW_;
|
||||
static constexpr int WARP_COL = WARP_COL_;
|
||||
static constexpr int NUM_WARPS = WARP_ROW * WARP_COL;
|
||||
static constexpr int NUM_THREADS = NUM_WARPS * 32;
|
||||
static constexpr int MMA_WARP_M = WARP_ROW * 16;
|
||||
static constexpr int MMA_WARP_N = WARP_COL * 8;
|
||||
static constexpr int MMA_WARP_K = 32;
|
||||
using accum_t = accum_t_;
|
||||
using out_t = out_t_;
|
||||
using SwizzleLayoutO = std::conditional_t<
|
||||
std::is_same_v<out_t_, cutlass::bfloat16_t>,
|
||||
Swizzle<3, 3, 3>,
|
||||
Swizzle<2, 4, 3>
|
||||
>;
|
||||
using SwizzleLayoutAB = Swizzle<2, 4, 3>;
|
||||
using MMA_Atom_SM89 = std::conditional_t<
|
||||
std::is_same_v<accum_t, cutlass::half_t>,
|
||||
MMA_Atom<SM89_16x8x32_F16E4M3E4M3F16_TN>,
|
||||
MMA_Atom<SM89_16x8x32_F32E4M3E4M3F32_TN>
|
||||
>;
|
||||
static constexpr int INPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(float_e4m3_t);
|
||||
static constexpr int OUTPUT_ELEMS_PER_COPY = sizeof(uint128_t) / sizeof(out_t_);
|
||||
static constexpr int THREADS_PER_ROW = BK / INPUT_ELEMS_PER_COPY;
|
||||
using GMEMLayout = Layout< Shape <Int<NUM_THREADS / THREADS_PER_ROW>, Int<THREADS_PER_ROW>>, Stride<Int<THREADS_PER_ROW>, _1>>;
|
||||
using G2SCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEGLOBAL<cute::uint128_t>, float_e4m3_t>;
|
||||
using G2STiledCopy = decltype(
|
||||
make_tiled_copy(
|
||||
G2SCopyAtom{},
|
||||
GMEMLayout{},
|
||||
Layout<Shape<_1, Int<INPUT_ELEMS_PER_COPY>>>{}
|
||||
)
|
||||
);
|
||||
using S2RCopyAtomA = Copy_Atom<SM75_U32x4_LDSM_N, float_e4m3_t>;
|
||||
using S2RCopyAtomB = Copy_Atom<SM75_U32x2_LDSM_N, float_e4m3_t>;
|
||||
using SmemLayoutAtom = decltype(composition(
|
||||
Swizzle<2, 4, 3>{},
|
||||
make_layout(make_shape(Int<8>{}, Int<BK>{}),
|
||||
make_stride(Int<BK>{}, Int<1>{}))));
|
||||
using SmemLayoutA = decltype(
|
||||
tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BM>{}, Int<BK>{}, Int<KStages>{}))
|
||||
);
|
||||
using SmemLayoutB = decltype(
|
||||
tile_to_shape(SmemLayoutAtom{}, make_shape(Int<BN>{}, Int<BK>{}, Int<KStages>{}))
|
||||
);
|
||||
using MMATile = decltype(
|
||||
make_tiled_mma(
|
||||
MMA_Atom_SM89{},
|
||||
Layout<Shape<Int<WARP_ROW>, Int<WARP_COL>, _1>>{},
|
||||
Tile<Int<MMA_WARP_M>, Int<MMA_WARP_N>, Int<MMA_WARP_K>>{}
|
||||
)
|
||||
);
|
||||
|
||||
static constexpr int ELEMS_PER_TILE = MMA_WARP_M * MMA_WARP_N;
|
||||
static constexpr int NUM_ELEMS_PER_WRITE = NUM_THREADS * sizeof(cute::uint128_t) / sizeof(out_t_);
|
||||
static constexpr int OUT_PIPE = NUM_ELEMS_PER_WRITE / ELEMS_PER_TILE;
|
||||
// using SmemLayoutC = Layout<Shape<Int<BM>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
|
||||
|
||||
using SmemLayoutC = decltype(
|
||||
make_layout(
|
||||
make_shape(Int<MMA_WARP_M>{}, Int<MMA_WARP_N*OUT_PIPE>{}),
|
||||
make_stride(Int<MMA_WARP_N*OUT_PIPE>{}, Int<1>{})
|
||||
)
|
||||
);
|
||||
static constexpr int THREADS_PER_ROW_WRITE = MMA_WARP_N * OUT_PIPE / OUTPUT_ELEMS_PER_COPY;
|
||||
using R2SCopyAtomC = Copy_Atom<UniversalCopy<typename BytesToType<2*sizeof(out_t)>::Type>, out_t>;
|
||||
using S2GCopyAtomC = Copy_Atom<UniversalCopy<cute::uint128_t>, out_t>;
|
||||
using S2GCopyC = decltype(make_tiled_copy(S2GCopyAtomC{},
|
||||
make_layout(make_shape(Int<NUM_THREADS / THREADS_PER_ROW_WRITE>{}, Int<THREADS_PER_ROW_WRITE>{}),
|
||||
make_stride(Int<THREADS_PER_ROW_WRITE>{}, Int<1>{})),
|
||||
make_layout(make_shape(Int<1>{}, Int<OUTPUT_ELEMS_PER_COPY>{}))));
|
||||
|
||||
using G2SBiasCopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<float>, float>;
|
||||
using G2SBiasCopy = decltype(make_tiled_copy(G2SBiasCopyAtom{}, make_layout(
|
||||
make_shape(Int<1>{},Int<BN>{}), make_stride(Int<BN>{}, Int<1>{})),
|
||||
make_layout(make_shape(Int<1>{},Int<1>{}), make_stride(Int<1>{}, Int<1>{}))));
|
||||
using sfa_copy_vtype = float;
|
||||
static constexpr int SFA_ELEMS_PER_COPY = sizeof(sfa_copy_vtype)/sizeof(float);
|
||||
static constexpr int THREADS_SFA_COPY = BM * sizeof(float) / sizeof(sfa_copy_vtype);
|
||||
// using G2SSFACopyAtom = Copy_Atom<SM80_CP_ASYNC_CACHEALWAYS<cute::uint128_t>, float>;
|
||||
static constexpr bool HasBias = HasBias_;
|
||||
using SmemLayoutBias = Layout<Shape<Int<1>, Int<BN>>, Stride<Int<BN>, Int<1>>>;
|
||||
using SmemLayoutSFA = Layout<Shape<Int<BM>, Int<KStages>>, Stride<Int<1>, Int<BM>>>;
|
||||
using BiasThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_2, _0>, Stride<_0, _8>>>;
|
||||
using SFAThreadLayout = Layout<Shape<Shape<_4, _8>, Shape<Int<WARP_ROW>, Int<WARP_COL>>>, Stride<Stride<_0, _1>, Stride<_16, _0>>>;
|
||||
static constexpr int SmemSize = cute::max(cute::cosize(SmemLayoutA{})+cute::cosize(SmemLayoutB{}), cute::cosize(SmemLayoutC{})*sizeof(out_t)) + cute::cosize(SmemLayoutBias{}) * sizeof(float) + cute::cosize(SmemLayoutSFA{})*sizeof(float);
|
||||
};
|
||||
@@ -0,0 +1,84 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/config.hpp>
|
||||
#include <cute/arch/mma.hpp>
|
||||
////////////////////////////////////////////////////////////////////////////////
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 4)
|
||||
# define CUTE_ARCH_MMA_F32_SM89_SUPPORTED
|
||||
#endif
|
||||
|
||||
#if (__CUDACC_VER_MAJOR__ > 12) || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ >= 8)
|
||||
# define CUTE_ARCH_MMA_F16_SM89_SUPPORTED
|
||||
#endif
|
||||
|
||||
#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 890)
|
||||
# if defined(CUTE_ARCH_MMA_F32_SM89_SUPPORTED)
|
||||
# define CUTE_ARCH_MMA_F32_SM89_ENABLED
|
||||
# endif
|
||||
|
||||
# if defined(CUTE_ARCH_MMA_F16_SM89_SUPPORTED)
|
||||
# define CUTE_ARCH_MMA_F16_SM89_ENABLED
|
||||
# endif
|
||||
#endif
|
||||
|
||||
namespace cute {
|
||||
struct SM89_16x8x32_F32E4M3E4M3F32_TN
|
||||
{
|
||||
using DRegisters = float[4];
|
||||
using ARegisters = uint32_t[4];
|
||||
using BRegisters = uint32_t[2];
|
||||
using CRegisters = float[4];
|
||||
|
||||
CUTE_HOST_DEVICE static void
|
||||
fma(float & d0, float & d1, float & d2, float & d3,
|
||||
uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
|
||||
uint32_t const& b0, uint32_t const& b1,
|
||||
float const& c0, float const& c1, float const& c2, float const& c3)
|
||||
{
|
||||
#if defined(CUTE_ARCH_MMA_F32_SM89_ENABLED)
|
||||
asm(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 "
|
||||
"{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"
|
||||
: "=f"(d0), "=f"(d1), "=f"(d2), "=f"(d3)
|
||||
:
|
||||
"r"(a0), "r"(a1), "r"(a2), "r"(a3),
|
||||
"r"(b0), "r"(b1),
|
||||
"f"(c0), "f"(c1), "f"(c2), "f"(c3)
|
||||
);
|
||||
#else
|
||||
CUTE_INVALID_CONTROL_PATH("Attempting to use SM89_16x8x32_F32E4M3E4M3F32_TN without CUTE_ARCH_MMA_F32_SM89_ENABLED");
|
||||
#endif
|
||||
}
|
||||
};
|
||||
|
||||
// MMA 16x8x32 TN
|
||||
struct SM89_16x8x32_F16E4M3E4M3F16_TN
|
||||
{
|
||||
using DRegisters = uint32_t[2];
|
||||
using ARegisters = uint32_t[4];
|
||||
using BRegisters = uint32_t[2];
|
||||
using CRegisters = uint32_t[2];
|
||||
|
||||
CUTE_HOST_DEVICE static void
|
||||
fma(uint32_t & d0, uint32_t & d1,
|
||||
uint32_t const& a0, uint32_t const& a1, uint32_t const& a2, uint32_t const& a3,
|
||||
uint32_t const& b0, uint32_t const& b1,
|
||||
uint32_t const& c0, uint32_t const& c1)
|
||||
{
|
||||
#if defined(CUTE_ARCH_MMA_F16_SM89_ENABLED)
|
||||
asm(
|
||||
"mma.sync.aligned.m16n8k32.row.col.f16.e4m3.e4m3.f16 "
|
||||
"{%0,%1}, {%2,%3,%4,%5}, {%6,%7}, {%8,%9};\n"
|
||||
: "=r"(d0), "=r"(d1)
|
||||
:
|
||||
"r"(a0), "r"(a1), "r"(a2), "r"(a3),
|
||||
"r"(b0), "r"(b1),
|
||||
"r"(c0), "r"(c1)
|
||||
);
|
||||
#else
|
||||
CUTE_INVALID_CONTROL_PATH("Attempting to use SM89_16x8x32_F32E4M3E4M3F32_TN without CUTE_ARCH_MMA_F16_SM89_ENABLED");
|
||||
#endif
|
||||
}
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,50 @@
|
||||
#pragma once
|
||||
|
||||
#include <cute/atom/mma_traits.hpp>
|
||||
#include <cute/layout.hpp>
|
||||
#include <cute/numeric/numeric_types.hpp>
|
||||
#include "mma_sm89_fp16.hpp"
|
||||
|
||||
namespace cute
|
||||
{
|
||||
|
||||
namespace {
|
||||
|
||||
// (T32,V4) -> (M16,N8)
|
||||
using SM80_16x8_Row = Layout<Shape <Shape < _4,_8>,Shape < _2,_2>>,
|
||||
Stride<Stride<_32,_1>,Stride<_16,_8>>>;
|
||||
|
||||
}
|
||||
|
||||
template <>
|
||||
struct MMA_Traits<SM89_16x8x32_F32E4M3E4M3F32_TN> {
|
||||
using ValTypeD = float;
|
||||
using ValTypeA = float_e4m3_t;
|
||||
using ValTypeB = float_e4m3_t;
|
||||
using ValTypeC = float;
|
||||
|
||||
using Shape_MNK = Shape<_16,_8,_32>;
|
||||
using ThrID = Layout<_32>;
|
||||
using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _4,_2, _2>>,
|
||||
Stride<Stride<_64,_1>,Stride<_16,_8,_256>>>;
|
||||
using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_4, _2>>,
|
||||
Stride<Stride<_32,_1>,Stride<_8,_128>>>;
|
||||
using CLayout = SM80_16x8_Row;
|
||||
};
|
||||
|
||||
template <>
|
||||
struct MMA_Traits<SM89_16x8x32_F16E4M3E4M3F16_TN> {
|
||||
using ValTypeD = half_t;
|
||||
using ValTypeA = float_e4m3_t;
|
||||
using ValTypeB = float_e4m3_t;
|
||||
using ValTypeC = half_t;
|
||||
|
||||
using Shape_MNK = Shape<_16,_8,_32>;
|
||||
using ThrID = Layout<_32>;
|
||||
using ALayout = Layout<Shape <Shape < _4,_8>,Shape < _4,_2, _2>>,
|
||||
Stride<Stride<_64,_1>,Stride<_16,_8,_256>>>;
|
||||
using BLayout = Layout<Shape <Shape < _4,_8>,Shape <_4, _2>>,
|
||||
Stride<Stride<_32,_1>,Stride<_8,_128>>>;
|
||||
using CLayout = SM80_16x8_Row;
|
||||
};
|
||||
}
|
||||
@@ -0,0 +1,35 @@
|
||||
#pragma once
|
||||
|
||||
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
|
||||
if (COND) { \
|
||||
constexpr static bool CONST_NAME = true; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
__VA_ARGS__ \
|
||||
}
|
||||
//K/128
|
||||
#define BLOCK_K_SWITCH(COSNT_NAME, ...) \
|
||||
if (K == 2048) { \
|
||||
constexpr static int COSNT_NAME = 2048; \
|
||||
__VA_ARGS__ \
|
||||
} \
|
||||
else if (K == 4096) { \
|
||||
constexpr static int COSNT_NAME = 4096; \
|
||||
__VA_ARGS__ \
|
||||
} else if (K == 8192) { \
|
||||
constexpr static int COSNT_NAME = 8192; \
|
||||
__VA_ARGS__ \
|
||||
} else if (K == 16384) { \
|
||||
constexpr static int COSNT_NAME = 16384; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
TORCH_CHECK(false, "Unsupported K value: ", K); \
|
||||
}
|
||||
|
||||
#define M_SWITCH(...) \
|
||||
constexpr static int BM = 64; \
|
||||
constexpr static int BN = 128; \
|
||||
constexpr static int WARP_ROW = 2; \
|
||||
constexpr static int WARP_COL = 4; \
|
||||
__VA_ARGS__
|
||||
@@ -0,0 +1,206 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include <torch/python.h>
|
||||
#include "exceptions.hpp"
|
||||
|
||||
namespace blockwise {
|
||||
template <typename T>
|
||||
static T ceil_div(const T& a, const T& b) {
|
||||
return (a + b - 1) / b;
|
||||
}
|
||||
template <typename T>
|
||||
static constexpr T align(const T& a, const T& b) {
|
||||
return ceil_div(a, b) * b;
|
||||
}
|
||||
|
||||
static int get_tma_aligned_size(const int& x, const int& element_size) {
|
||||
constexpr int kNumTMAAlignmentBytes = 16;
|
||||
DG_HOST_ASSERT(kNumTMAAlignmentBytes % element_size == 0);
|
||||
return align(x, kNumTMAAlignmentBytes / element_size);
|
||||
}
|
||||
static std::pair<int, int> get_inner_outer_dims(const cute::UMMA::Major& major, const int& k, const int& mn) {
|
||||
return major == cute::UMMA::Major::K ? std::make_pair(k, mn) : std::make_pair(mn, k);
|
||||
}
|
||||
|
||||
static int get_non_contiguous_dim(const cute::UMMA::Major& major) {
|
||||
return major == cute::UMMA::Major::K ? -2 : -1;
|
||||
}
|
||||
|
||||
static int get_compiled_dim(const int& dim, const char& name, const std::string& compiled_dims) {
|
||||
for (const char& c: compiled_dims) {
|
||||
if (name == c)
|
||||
return dim;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType& dtype,
|
||||
const bool& allow_tf32) {
|
||||
if (allow_tf32 and dtype == torch::kFloat)
|
||||
return CU_TENSOR_MAP_DATA_TYPE_TFLOAT32;
|
||||
|
||||
switch (dtype) {
|
||||
case torch::kInt: return CU_TENSOR_MAP_DATA_TYPE_INT32;
|
||||
case torch::kFloat: return CU_TENSOR_MAP_DATA_TYPE_FLOAT32;
|
||||
case torch::kBFloat16: return CU_TENSOR_MAP_DATA_TYPE_BFLOAT16;
|
||||
case torch::kFloat8_e4m3fn: return CU_TENSOR_MAP_DATA_TYPE_UINT8;
|
||||
default: DG_HOST_UNREACHABLE("Unsupported dtype");
|
||||
}
|
||||
}
|
||||
|
||||
static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const int& base) {
|
||||
#if CUDA_VERSION >= 12080
|
||||
if (base != 0) {
|
||||
DG_HOST_ASSERT(base == 32 and mode == 128);
|
||||
return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B;
|
||||
}
|
||||
#endif
|
||||
|
||||
DG_HOST_ASSERT(base == 0);
|
||||
switch (mode) {
|
||||
case 0:
|
||||
case 16: return CU_TENSOR_MAP_SWIZZLE_NONE;
|
||||
case 32: return CU_TENSOR_MAP_SWIZZLE_32B;
|
||||
case 64: return CU_TENSOR_MAP_SWIZZLE_64B;
|
||||
case 128: return CU_TENSOR_MAP_SWIZZLE_128B;
|
||||
default: DG_HOST_UNREACHABLE("Unsupported swizzling mode");
|
||||
}
|
||||
}
|
||||
|
||||
static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
|
||||
int gmem_inner_dim, int gmem_outer_dim,
|
||||
int smem_inner_dim, int smem_outer_dim,
|
||||
const int& gmem_outer_stride,
|
||||
const int& swizzle_mode, const int& swizzle_base = 0,
|
||||
const bool& allow_tf32 = false) {
|
||||
const auto& elem_size = static_cast<int>(t.element_size());
|
||||
if (swizzle_mode != 0)
|
||||
smem_inner_dim = swizzle_mode / elem_size;
|
||||
|
||||
CUtensorMap tensor_map;
|
||||
const cuuint64_t gmem_dims[2] = {static_cast<cuuint64_t>(gmem_inner_dim), static_cast<cuuint64_t>(gmem_outer_dim)};
|
||||
const cuuint32_t smem_dims[2] = {static_cast<cuuint32_t>(smem_inner_dim), static_cast<cuuint32_t>(smem_outer_dim)};
|
||||
const cuuint64_t gmem_strides[1] = {static_cast<cuuint64_t>(gmem_outer_stride * elem_size), };
|
||||
const cuuint32_t elem_strides[2] = {1, 1};
|
||||
// if (get_env<int>("DG_JIT_DEBUG")) {
|
||||
// printf("Making TMA desc: global memory: %d %d, shared memory: %d %d, outer stride: %d, swizzle: %d (base: %d), elem size: %d\n",
|
||||
// gmem_inner_dim, gmem_outer_dim, smem_inner_dim, smem_outer_dim,
|
||||
// gmem_outer_stride, swizzle_mode, swizzle_base, elem_size);
|
||||
// }
|
||||
cuTensorMapEncodeTiled(
|
||||
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
|
||||
2, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
|
||||
CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
|
||||
CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
|
||||
return tensor_map;
|
||||
}
|
||||
|
||||
static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
|
||||
const int& gmem_dim_0, const int& gmem_dim_1, const int& gmem_dim_2,
|
||||
const int& smem_dim_0, const int& smem_dim_1, const int& smem_dim_2,
|
||||
const int& gmem_stride_0, const int& gmem_stride_1,
|
||||
const int& swizzle_mode, const int& swizzle_base = 0,
|
||||
const bool& allow_tf32 = false) {
|
||||
const auto& elem_size = static_cast<int>(t.element_size());
|
||||
if (swizzle_mode != 0)
|
||||
DG_HOST_ASSERT(smem_dim_0 == swizzle_mode / elem_size);
|
||||
|
||||
CUtensorMap tensor_map;
|
||||
const cuuint64_t gmem_dims[3] = {static_cast<cuuint64_t>(gmem_dim_0), static_cast<cuuint64_t>(gmem_dim_1), static_cast<cuuint64_t>(gmem_dim_2),};
|
||||
const cuuint32_t smem_dims[3] = {static_cast<cuuint32_t>(smem_dim_0), static_cast<cuuint32_t>(smem_dim_1), static_cast<cuuint32_t>(smem_dim_2)};
|
||||
const cuuint64_t gmem_strides[2] = {static_cast<cuuint64_t>(gmem_stride_0 * elem_size), static_cast<cuuint64_t>(gmem_stride_1 * elem_size)};
|
||||
const cuuint32_t elem_strides[3] = {1, 1, 1};
|
||||
// if (get_env<int>("DG_JIT_DEBUG")) {
|
||||
// printf("Making 3D TMA desc: global memory: %d %d %d, shared memory: %d %d %d, outer stride: %d %d, swizzle: %d, elem size: %d\n",
|
||||
// gmem_dim_0, gmem_dim_1, gmem_dim_2, smem_dim_0, smem_dim_1, smem_dim_2,
|
||||
// gmem_stride_0, gmem_stride_1, swizzle_mode, elem_size);
|
||||
// }
|
||||
cuTensorMapEncodeTiled(
|
||||
&tensor_map, aten_dtype_to_tensor_map_dtype(t.scalar_type(), allow_tf32),
|
||||
3, t.data_ptr(), gmem_dims, gmem_strides, smem_dims, elem_strides,
|
||||
CU_TENSOR_MAP_INTERLEAVE_NONE, mode_into_tensor_map_swizzle(swizzle_mode, swizzle_base),
|
||||
CU_TENSOR_MAP_L2_PROMOTION_L2_256B, CU_TENSOR_MAP_FLOAT_OOB_FILL_NONE);
|
||||
return tensor_map;
|
||||
}
|
||||
|
||||
static CUtensorMap make_tma_a_desc(const cute::UMMA::Major& major,
|
||||
const torch::Tensor& t,
|
||||
const int& shape_m, const int& shape_k,
|
||||
const int& block_m, const int& block_k,
|
||||
const int& outer_stride,
|
||||
const int& num_groups,
|
||||
const int& swizzle_mode, const int& swizzle_base = 0,
|
||||
const bool& allow_tf32 = false) {
|
||||
if (num_groups > 1)
|
||||
DG_HOST_ASSERT(major == cute::UMMA::Major::K);
|
||||
const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_m * num_groups);
|
||||
const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_m);
|
||||
return make_tma_2d_desc(t,
|
||||
gmem_inner_dim, gmem_outer_dim,
|
||||
smem_inner_dim, smem_outer_dim,
|
||||
outer_stride,
|
||||
swizzle_mode, swizzle_base,
|
||||
allow_tf32);
|
||||
}
|
||||
|
||||
static CUtensorMap make_tma_b_desc(const cute::UMMA::Major& major,
|
||||
const torch::Tensor& t,
|
||||
const int& shape_n, const int& shape_k,
|
||||
const int& block_n, const int& block_k,
|
||||
const int& outer_stride,
|
||||
const int& num_groups,
|
||||
const int& swizzle_mode, const int& swizzle_base = 0,
|
||||
const bool& allow_tf32 = false) {
|
||||
const auto& [gmem_inner_dim, gmem_outer_dim] = get_inner_outer_dims(major, shape_k, shape_n);
|
||||
const auto& [smem_inner_dim, smem_outer_dim] = get_inner_outer_dims(major, block_k, block_n);
|
||||
|
||||
// `num_groups` is always applied into the outer dimensions
|
||||
return make_tma_2d_desc(t,
|
||||
gmem_inner_dim, gmem_outer_dim * num_groups,
|
||||
smem_inner_dim, smem_outer_dim,
|
||||
outer_stride,
|
||||
swizzle_mode, swizzle_base,
|
||||
allow_tf32);
|
||||
}
|
||||
|
||||
static CUtensorMap make_tma_cd_desc(const torch::Tensor& t,
|
||||
const int& shape_m, const int& shape_n,
|
||||
const int& block_m, const int& block_n,
|
||||
const int& outer_stride,
|
||||
const int& num_groups,
|
||||
const int& swizzle_mode, const int& swizzle_base = 0,
|
||||
const bool& allow_tf32 = false) {
|
||||
// Swizzling requires the inner box dim to be less or equal than `kSwizzleCDMode`
|
||||
// bytes, so `BLOCK_N * sizeof(T) / kSwizzleCDMode` TMA stores are required
|
||||
return make_tma_2d_desc(t,
|
||||
shape_n, shape_m * num_groups,
|
||||
block_n, block_m,
|
||||
outer_stride,
|
||||
swizzle_mode, swizzle_base,
|
||||
allow_tf32);
|
||||
}
|
||||
|
||||
static CUtensorMap make_tma_sf_desc(const cute::UMMA::Major& major,
|
||||
const torch::Tensor& t,
|
||||
int shape_mn, int shape_k,
|
||||
const int& block_mn, const int& block_k,
|
||||
const int& num_groups,
|
||||
const int& swizzle_mode, const int& swizzle_base = 0,
|
||||
const bool& allow_tf32 = false) {
|
||||
DG_HOST_ASSERT(major == cute::UMMA::Major::MN);
|
||||
|
||||
// TODO: maybe swizzle SF as well
|
||||
DG_HOST_ASSERT(swizzle_mode == 0);
|
||||
|
||||
shape_mn = get_tma_aligned_size(shape_mn, static_cast<int>(t.element_size()));
|
||||
return make_tma_2d_desc(t,
|
||||
shape_mn, ceil_div(shape_k, block_k * (t.scalar_type() == torch::kFloat ? 1 : 4)) * num_groups,
|
||||
block_mn, 1,
|
||||
shape_mn,
|
||||
swizzle_mode, swizzle_base,
|
||||
allow_tf32);
|
||||
}
|
||||
|
||||
} // namespace deep_gemm
|
||||
@@ -0,0 +1,39 @@
|
||||
#pragma once
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <nvrtc.h>
|
||||
|
||||
#include <torch/python.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include "kernels/geforce/static_switch.h"
|
||||
|
||||
namespace sm89 {
|
||||
template<bool use_fast_accum>
|
||||
void fp8_bias_gemm_cuda(void* Aptr, void* SFA, void* Bptr, void* SFB, void* bias_ptr, void* out, int M, int N, int K, cudaStream_t stream);
|
||||
}
|
||||
|
||||
namespace blockwise {
|
||||
static void sm89_fp8_gemm_1d2d_bias(const torch::Tensor& a, const torch::Tensor& sfa,
|
||||
const torch::Tensor& b, const torch::Tensor& sfb,
|
||||
const torch::Tensor& bias,
|
||||
const torch::Tensor& d,
|
||||
const int& m, const int& n, const int& k,
|
||||
const bool use_fast_accum) {
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
|
||||
if (use_fast_accum) {
|
||||
sm89::fp8_bias_gemm_cuda<true>(
|
||||
a.data_ptr(), sfa.data_ptr(),
|
||||
b.data_ptr(), sfb.data_ptr(),
|
||||
bias.data_ptr(), d.data_ptr(),
|
||||
m, n, k, stream);
|
||||
} else {
|
||||
sm89::fp8_bias_gemm_cuda<false>(
|
||||
a.data_ptr(), sfa.data_ptr(),
|
||||
b.data_ptr(), sfb.data_ptr(),
|
||||
bias.data_ptr(), d.data_ptr(),
|
||||
m, n, k, stream);
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,66 @@
|
||||
|
||||
|
||||
#pragma once
|
||||
#include <cuda.h>
|
||||
#include <cuda_runtime.h>
|
||||
#include <nvrtc.h>
|
||||
|
||||
#include <torch/python.h>
|
||||
#include <ATen/cuda/CUDAContext.h>
|
||||
|
||||
#include <cute/arch/mma_sm100_desc.hpp>
|
||||
#include "runtime_utils.hpp"
|
||||
|
||||
#include "config.hpp"
|
||||
#include "static_switch.hpp"
|
||||
|
||||
namespace deep_gemm{
|
||||
template<int N, int K>
|
||||
void sm90_fp8_gemm_1d2d_bias_launch(int num_sms, int num_threads, int cluster_dim, int smem_size, cudaStream_t stream, float* sfb, float* bias, int* grouped_layout,
|
||||
uint32_t shape_m, uint32_t shape_n, uint32_t shape_k,
|
||||
const CUtensorMap tensor_map_a,
|
||||
const CUtensorMap tensor_map_b,
|
||||
const CUtensorMap tensor_map_d,
|
||||
const CUtensorMap tensor_map_sfa);
|
||||
};
|
||||
|
||||
namespace blockwise{
|
||||
|
||||
static void sm90_fp8_gemm_1d2d_bias(const torch::Tensor& a, const torch::Tensor& sfa,
|
||||
const torch::Tensor& b, const torch::Tensor& sfb,
|
||||
const torch::Tensor& bias,
|
||||
const std::optional<torch::Tensor>& c,
|
||||
const torch::Tensor& d,
|
||||
const int& m, const int& n, const int& k, const int num_sms) {
|
||||
// DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
|
||||
const auto& config = GemmConfig<90>();
|
||||
|
||||
// Requires no TMA splits
|
||||
// DG_HOST_ASSERT(config.smem_config.swizzle_a_mode == config.block_k);
|
||||
// DG_HOST_ASSERT(config.smem_config.swizzle_b_mode == config.block_k);
|
||||
int smem_size = k == 16384 || k == 8192 ? 216624 : config.smem_config.smem_size;
|
||||
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::K, a, m, k,
|
||||
config.block_m,
|
||||
config.block_k,
|
||||
static_cast<int>(a.stride(-2)), 1,
|
||||
config.smem_config.swizzle_a_mode);
|
||||
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::K, b, n, k,
|
||||
config.block_n,
|
||||
config.block_k,
|
||||
static_cast<int>(b.stride(-2)), 1,
|
||||
config.smem_config.swizzle_b_mode);
|
||||
const auto& tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
|
||||
config.block_m,
|
||||
config.block_n,
|
||||
static_cast<int>(d.stride(-2)), 1,
|
||||
config.smem_config.swizzle_cd_mode);
|
||||
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, k,
|
||||
config.block_m, config.block_k, 1, 0);
|
||||
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
||||
// Launch
|
||||
DIM_SWITCH(k, K,
|
||||
DIM_SWITCH(n, N,
|
||||
deep_gemm::sm90_fp8_gemm_1d2d_bias_launch<N, K>(num_sms, config.thread_config.num_threads, config.multicast_config.num_multicast, smem_size, stream, (float*)sfb.data_ptr(), (float*)bias.data_ptr(), nullptr, m, n, k, tensor_map_a, tensor_map_b, tensor_map_d, tensor_map_sfa);)
|
||||
)
|
||||
}
|
||||
};
|
||||
@@ -0,0 +1,60 @@
|
||||
#pragma once
|
||||
|
||||
#define DIM_SWITCH(VAR_NAME, CONST_NAME, ...) \
|
||||
if (VAR_NAME == 4096) { \
|
||||
constexpr static int CONST_NAME = 4096; \
|
||||
__VA_ARGS__ \
|
||||
} else if (VAR_NAME == 2048){ \
|
||||
constexpr static int CONST_NAME = 2048; \
|
||||
__VA_ARGS__ \
|
||||
} else if (VAR_NAME == 8192){ \
|
||||
constexpr static int CONST_NAME = 8192; \
|
||||
__VA_ARGS__ \
|
||||
} else if(VAR_NAME == 16384) { \
|
||||
constexpr static int CONST_NAME = 16384; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
TORCH_CHECK(false, "Unsupported DIM_SWITCH value: ", VAR_NAME); \
|
||||
}
|
||||
|
||||
#define BOOL_SWITCH(COND, CONST_NAME, ...) \
|
||||
if (COND) { \
|
||||
constexpr static bool CONST_NAME = true; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
constexpr static bool CONST_NAME = false; \
|
||||
__VA_ARGS__ \
|
||||
} \
|
||||
//K/128
|
||||
#define BLOCK_K_SWITCH(COSNT_NAME, ...) \
|
||||
if (K == 2048) { \
|
||||
constexpr static int COSNT_NAME = 16; \
|
||||
__VA_ARGS__ \
|
||||
} \
|
||||
else if (K == 4096) { \
|
||||
constexpr static int COSNT_NAME = 32; \
|
||||
__VA_ARGS__ \
|
||||
} else if (K == 8192) { \
|
||||
constexpr static int COSNT_NAME = 64; \
|
||||
__VA_ARGS__ \
|
||||
} else if (K == 16384) { \
|
||||
constexpr static int COSNT_NAME = 128; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
TORCH_CHECK(false, "Unsupported K value: ", K); \
|
||||
}
|
||||
|
||||
#define M_SWITCH(...) \
|
||||
if (M <= 1024) { \
|
||||
constexpr static int BM = 128; \
|
||||
constexpr static int BN = 128; \
|
||||
constexpr static int WARP_ROW = 2; \
|
||||
constexpr static int WARP_COL = 2; \
|
||||
__VA_ARGS__ \
|
||||
} else { \
|
||||
constexpr static int BM = 128; \
|
||||
constexpr static int BN = 256; \
|
||||
constexpr static int WARP_ROW = 2; \
|
||||
constexpr static int WARP_COL = 4; \
|
||||
__VA_ARGS__ \
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
#pragma once
|
||||
|
||||
#include <cuda_bf16.h>
|
||||
#include <cuda_fp8.h>
|
||||
#include <cuda/std/cstdint>
|
||||
#include <cuda/std/utility>
|
||||
#include <cute/container/tuple.hpp>
|
||||
|
||||
#ifdef __CLION_IDE__
|
||||
|
||||
__host__ __device__ __forceinline__ void host_device_printf(const char* format, ...) {
|
||||
asm volatile("trap;");
|
||||
}
|
||||
|
||||
#define printf host_device_printf
|
||||
#endif
|
||||
|
||||
#ifndef DG_DEVICE_ASSERT
|
||||
#define DG_DEVICE_ASSERT(cond) \
|
||||
do { \
|
||||
if (not (cond)) { \
|
||||
printf("Assertion failed: %s:%d, condition: %s\n", __FILE__, __LINE__, #cond); \
|
||||
asm("trap;"); \
|
||||
} \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#ifndef DG_TRAP_ONLY_DEVICE_ASSERT
|
||||
#define DG_TRAP_ONLY_DEVICE_ASSERT(cond) \
|
||||
do { \
|
||||
if (not (cond)) \
|
||||
asm("trap;"); \
|
||||
} while (0)
|
||||
#endif
|
||||
|
||||
#ifndef DG_STATIC_ASSERT
|
||||
#define DG_STATIC_ASSERT(cond, ...) static_assert(cond, __VA_ARGS__)
|
||||
#endif
|
||||
Reference in New Issue
Block a user