chenchenghao/DeepGEMM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Multiple updates and refactorings (#231)

Browse Source
This commit is contained in:
Ray Wang
2025-11-21 17:49:47 +08:00
committed by GitHub
parent bb4424aad4
commit 38f8ef73a4
80 changed files with 3767 additions and 2103 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/apis/attention.hpp
+58 -21
View File
@@ -1,21 +1,25 @@
#pragma once
#include "../utils/compatibility.hpp"
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
#include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
#include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
#include "../jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp"
#include "../jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp"
#include "../jit_kernels/impls/smxx_fp8_mqa_logits.hpp"
#include "../jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp"
#include "../jit_kernels/impls/smxx_clean_logits.hpp"
#endif
#include "layout.hpp"
namespace deep_gemm::attention {
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tensor>& a,
const std::pair<torch::Tensor, torch::Tensor>& b,
const torch::Tensor& d,
const std::tuple<int, int, int> &head_splits,
const std::tuple<int, int, int>& head_splits,
std::optional<std::tuple<int, int, int>> recipe,
const std::string& compiled_dims,
const bool& disable_ue8m0_cast) {
@@ -59,11 +63,10 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
const auto& arch_major = device_runtime->get_arch_major();
const auto& epilogue_type = fmt::format("EpilogueHeadSplits<{}, {}, {}>", left, mid, right);
if (arch_major == 9 and sfa.scalar_type() == torch::kFloat and std::get<1>(recipe.value()) != 1) {
sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k, major_a, major_b, compiled_dims, epilogue_type);
const auto& major_sfb = get_major_type_ab(sfb);
sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k, major_a, major_b, major_sfb, compiled_dims, epilogue_type);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
sm100_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k, major_a, major_b, compiled_dims, epilogue_type);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kFloat) {
sm100_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, std::nullopt, d, m, n, k, major_a, major_b, compiled_dims, epilogue_type);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
@@ -74,7 +77,8 @@ static torch::Tensor fp8_mqa_logits(const torch::Tensor& q,
const torch::Tensor& weights,
const torch::Tensor& cu_seq_len_k_start,
const torch::Tensor& cu_seq_len_k_end,
const bool& clean_logits) {
const bool& clean_logits,
const int& max_seqlen_k) {
const auto& [seq_len, num_heads, head_dim] = get_shape<3>(q);
const auto& [seq_len_kv, head_dim_] = get_shape<2>(kv.first);
const auto& [seq_len_, num_heads_] = get_shape<2>(weights);
@@ -102,27 +106,45 @@ static torch::Tensor fp8_mqa_logits(const torch::Tensor& q,
constexpr int seq_len_alignment = 4;
constexpr int block_kv = 256;
const auto aligned_seq_len = align(seq_len, seq_len_alignment);
const auto aligned_seq_len_kv = align(seq_len_kv + block_kv, 4);
auto logits = torch::empty({aligned_seq_len, aligned_seq_len_kv}, q.options().dtype(torch::kFloat));
logits = logits.index({torch::indexing::Slice(0, seq_len), torch::indexing::Slice(0, seq_len_kv)});
torch::Tensor logits;
int stride_logits;
if (max_seqlen_k == 0) {
stride_logits = align(seq_len_kv + block_kv, 4);
logits = torch::empty({aligned_seq_len, stride_logits}, q.options().dtype(torch::kFloat));
logits = logits.index({torch::indexing::Slice(0, seq_len), torch::indexing::Slice(0, seq_len_kv)});
} else {
stride_logits = align(max_seqlen_k, block_kv);
logits = torch::empty({aligned_seq_len, stride_logits}, q.options().dtype(torch::kFloat));
logits = logits.index({torch::indexing::Slice(0, seq_len), torch::indexing::Slice(0, max_seqlen_k)});
DG_HOST_ASSERT(not clean_logits);
}
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 9 or arch_major == 10) {
smxx_fp8_mqa_logits(q, kv.first, kv.second, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits,
seq_len, seq_len_kv, aligned_seq_len_kv, num_heads, head_dim, seq_len_alignment);
seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, seq_len_alignment);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
// Clean unfilled logits
if (clean_logits)
smxx_clean_logits(logits, cu_seq_len_k_start, cu_seq_len_k_end, 1, seq_len, seq_len_kv, aligned_seq_len_kv);
smxx_clean_logits(logits, cu_seq_len_k_start, cu_seq_len_k_end, 1, seq_len, seq_len_kv, stride_logits);
return logits;
}
static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_lens, int block_kv, int num_sms) {
const auto& [batch_size] = get_shape<1>(context_lens);
const bool is_context_lens_2d = context_lens.dim() == 2;
int batch_size = 0, next_n = 0;
if (is_context_lens_2d) {
batch_size = context_lens.size(0);
next_n = context_lens.size(1);
} else {
DG_HOST_ASSERT(context_lens.dim() == 1);
batch_size = context_lens.size(0);
}
DG_HOST_ASSERT(context_lens.scalar_type() == torch::kInt);
DG_HOST_ASSERT(context_lens.is_contiguous());
@@ -131,7 +153,7 @@ static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 9 or arch_major == 10) {
smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, block_kv, num_sms);
smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
@@ -149,15 +171,24 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
const bool& clean_logits) {
const auto& [batch_size, next_n, num_heads, head_dim] = get_shape<4>(q);
const auto& [num_kv_blocks, block_kv, num_heads_kv, head_dim_with_sf] = get_shape<4>(fused_kv_cache);
const auto& [batch_size_] = get_shape<1>(context_lens);
const auto& [batch_size_next_n, num_heads_] = get_shape<2>(weights);
const auto& [batch_size__, max_block_len] = get_shape<2>(block_table);
const auto& [batch_size_, max_block_len] = get_shape<2>(block_table);
const auto& [schedule_meta_size, meta_info_size] = get_shape<2>(schedule_meta);
const auto& num_sms = device_runtime->get_num_sms();
const auto& kv_cache_stride_bytes = fused_kv_cache.stride(0);
const auto& block_table_stride = block_table.stride(0);
DG_HOST_ASSERT(batch_size == batch_size_ and batch_size == batch_size__);
const bool is_context_lens_2d = context_lens.dim() == 2;
if (is_context_lens_2d) {
const auto& [batch_size__, next_n_] = get_shape<2>(context_lens);
DG_HOST_ASSERT(batch_size == batch_size__ and next_n == next_n_);
} else {
DG_HOST_ASSERT(context_lens.dim() == 1);
const auto& [batch_size__] = get_shape<1>(context_lens);
DG_HOST_ASSERT(batch_size == batch_size__);
}
DG_HOST_ASSERT(batch_size == batch_size_);
DG_HOST_ASSERT(batch_size_next_n == batch_size * next_n);
DG_HOST_ASSERT(num_heads == num_heads_ and num_heads_kv == 1);
DG_HOST_ASSERT(head_dim_with_sf == head_dim + static_cast<int>(sizeof(float)));
@@ -207,34 +238,40 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 9 or arch_major == 10) {
smxx_fp8_paged_mqa_logits(q, kv_cache, kv_cache_scales, weights, context_lens, logits, block_table, schedule_meta,
batch_size, next_n, num_heads, head_dim, num_kv_blocks, block_kv,
batch_size, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,
kv_cache_stride_bytes, aligned_max_context_len, block_table_stride, num_sms, num_math_warp_groups);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
// Clean unfilled logits
if (clean_logits)
if (clean_logits) {
DG_HOST_ASSERT(not is_context_lens_2d);
smxx_clean_logits(logits, std::nullopt, context_lens, next_n, batch_size * next_n, max_context_len, aligned_max_context_len);
}
return logits;
}
#endif
static void register_apis(pybind11::module_& m) {
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
m.def("fp8_gemm_nt_skip_head_mid", &fp8_gemm_nt_skip_head_mid,
py::arg("a"), py::arg("b"), py::arg("d"), py::arg("head_splits"),
py::arg("recipe") = std::nullopt,
py::arg("compiled_dims") = "nk",
py::arg("disable_ue8m0_cast") = false);
m.def("fp8_mqa_logits", &fp8_mqa_logits,
py::arg("q"), py::arg("kv"), py::arg("weights"),
py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"),
py::arg("clean_logits") = true);
py::arg("q"), py::arg("kv"), py::arg("weights"),
py::arg("cu_seq_len_k_start"), py::arg("cu_seq_len_k_end"),
py::arg("clean_logits") = true,
py::arg("max_seqlen_k") = 0);
m.def("get_paged_mqa_logits_metadata", &get_paged_mqa_logits_metadata,
py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"));
m.def("fp8_paged_mqa_logits", &fp8_paged_mqa_logits,
py::arg("q"), py::arg("kv_cache"), py::arg("weights"),
py::arg("context_lens"), py::arg("block_table"), py::arg("schedule_meta"),
py::arg("max_context_len"), py::arg("clean_logits") = false);
#endif
}
} // namespace deep_gemm::attention
csrc/apis/einsum.hpp
+121 -8
View File
@@ -6,13 +6,20 @@
#include "../utils/exception.hpp"
#include "../utils/format.hpp"
#include "../utils/layout.hpp"
#include "../utils/compatibility.hpp"
#include "gemm.hpp"
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
#include "../jit_kernels/impls/sm90_bmk_bnk_mn.hpp"
#include "../jit_kernels/impls/sm100_bmk_bnk_mn.hpp"
#include "../jit_kernels/impls/sm90_bf16_gemm.hpp"
#include "../jit_kernels/impls/sm100_bf16_gemm.hpp"
#include "../jit_kernels/impls/smxx_cublaslt.hpp"
#endif
namespace deep_gemm::einsum {
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const torch::Tensor& d,
const std::optional<torch::Tensor>& c) {
// Currently FP32 only support the accumulated expression
@@ -51,7 +58,7 @@ static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const tor
}
}
static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D) {
static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D, const bool& use_cublaslt) {
const auto& [b , h , r ] = get_shape<3>(A);
const auto& [h_, d , r_] = get_shape<3>(B);
const auto& [b_, h__, d_] = get_shape<3>(D);
@@ -61,10 +68,20 @@ static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const to
DG_HOST_ASSERT(B.scalar_type() == torch::kBFloat16 and B.stride(2) == 1);
DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1);
cublaslt_bhr_hdr_bhd(A, B, D, b, h, r, d);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (use_cublaslt) {
cublaslt_bhr_hdr_bhd(A, B, D, b, h, r, d);
} else if (arch_major == 9) {
sm90_bf16_bhr_hdr_bhd(A, B, D, b, h, r, d);
} else if (arch_major == 10) {
sm100_bf16_bhr_hdr_bhd(A, B, D, b, h, r, d);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D) {
static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D, const bool& use_cublaslt) {
const auto& [b , h , d ] = get_shape<3>(A);
const auto& [h_, d_ , r ] = get_shape<3>(B);
const auto& [b_, h__, r_] = get_shape<3>(D);
@@ -74,14 +91,25 @@ static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const to
DG_HOST_ASSERT(B.scalar_type() == torch::kBFloat16 and B.stride(2) == 1);
DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1);
cublaslt_bhd_hdr_bhr(A, B, D, b, h, r, d);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (use_cublaslt) {
cublaslt_bhd_hdr_bhr(A, B, D, b, h, r, d);
} else if (arch_major == 9) {
sm90_bf16_bhd_hdr_bhr(A, B, D, b, h, r, d);
} else if (arch_major == 10) {
sm100_bf16_bhd_hdr_bhr(A, B, D, b, h, r, d);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
static void einsum(const std::string& expr,
const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& d,
const std::optional<torch::Tensor>& c) {
const std::optional<torch::Tensor>& c,
const bool& use_cublaslt) {
DG_HOST_ASSERT(a.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(b.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat);
@@ -94,22 +122,107 @@ static void einsum(const std::string& expr,
// TODO: support any expression
// TODO: canonicalize expression
if (expr == "bmk,bnk->mn") {
DG_HOST_ASSERT(not use_cublaslt);
bmk_bnk_mn(a, b, d, c);
} else if (expr == "bhr,hdr->bhd") {
DG_HOST_ASSERT(not c.has_value());
bhr_hdr_bhd(a, b, d);
bhr_hdr_bhd(a, b, d, use_cublaslt);
} else if (expr == "bhd,hdr->bhr") {
DG_HOST_ASSERT(not c.has_value());
bhd_hdr_bhr(a, b, d);
bhd_hdr_bhr(a, b, d, use_cublaslt);
} else {
DG_HOST_UNREACHABLE(fmt::format("Unsupported einsum expression: {}", expr));
}
}
static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const std::optional<torch::Tensor>& c,
const std::tuple<int, int, int>& recipe,
const std::string& compiled_dims) {
// Shape must be `[B, M, K] @ [B, N, K].T`
const auto& major_a = a.stride(-1) == 1 ? cute::UMMA::Major::K : cute::UMMA::Major::MN;
const auto& major_b = b.stride(-1) == 1 ? cute::UMMA::Major::K : cute::UMMA::Major::MN;
DG_HOST_ASSERT(a.stride(-1) == 1 or a.stride(-2) == 1);
DG_HOST_ASSERT(b.stride(-1) == 1 or b.stride(-2) == 1);
DG_HOST_ASSERT(d.stride(-1) == 1);
// Type and shape checks
const auto& [batch_size , m , k ] = get_shape<3>(a);
const auto& [batch_size_ , n , k_] = get_shape<3>(b);
const auto& [batch_size__, m_, n_] = get_shape<3>(d);
DG_HOST_ASSERT(batch_size == batch_size_ and batch_size == batch_size_);
DG_HOST_ASSERT(m == m_ and n == n_ and k == k_);
DG_HOST_ASSERT(a.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(b.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat);
// Early return for trivial cases
if (batch_size == 0 or gemm::early_return(m, n, k, d, c))
return;
// Transform scaling factors
const auto& transformed_sfa = layout::transform_sf_into_required_layout(sfa, m, k, recipe, batch_size, true, false);
const auto& transformed_sfb = layout::transform_sf_into_required_layout(sfb, n, k, recipe, batch_size, false, false);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 10) {
sm100_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d, batch_size, m, n, k, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
static void fp8_einsum(const std::string& expr,
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>& c,
const std::tuple<int, int, int>& recipe) {
// Some hardcoded Einstein sum kernels
if (expr == "bhr,hdr->bhd") {
// Permute dims to satisfy the order of (batch_size, m, n, k)
// (batch_size, m, n, k): (h, b, d, r)
const auto& perm_a = a.first.permute({1, 0, 2});
const auto& perm_sfa = a.second.permute({1, 0, 2});
const auto& perm_d = d.permute({1, 0, 2});
const auto& perm_c = c.has_value() ? std::make_optional(c.value().permute({1, 0, 2})) : std::nullopt;
fp8_bmm(perm_a, perm_sfa, b.first, b.second, perm_d, perm_c, recipe, "nk");
} else if (expr == "bhd,hdr->bhr") {
// (batch_size, m, n, k): (h, b, r, d)
const auto& perm_a = a.first.permute({1, 0, 2});
const auto& perm_sfa = a.second.permute({1, 0, 2});
const auto& perm_b = b.first.permute({0, 2, 1});
const auto& perm_sfb = b.second.permute({0, 2, 1});
const auto& perm_d = d.permute({1, 0, 2});
const auto& perm_c = c.has_value() ? std::make_optional(c.value().permute({1, 0, 2})) : std::nullopt;
fp8_bmm(perm_a, perm_sfa, perm_b, perm_sfb, perm_d, perm_c, recipe, "nk");
} else if (expr == "bhd,bhr->hdr") {
// (batch_size, m, n, k): (h, d, r, b)
const auto& perm_a = a.first.permute({1, 2, 0});
const auto& perm_sfa = a.second.permute({1, 2, 0});
const auto& perm_b = b.first.permute({1, 2, 0});
const auto& perm_sfb = b.second.permute({1, 2, 0});
fp8_bmm(perm_a, perm_sfa, perm_b, perm_sfb, d, c, recipe, "mn");
} else {
DG_HOST_UNREACHABLE(fmt::format("Unsupported einsum expression: {}", expr));
}
}
#endif
static void register_apis(pybind11::module_& m) {
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
m.def("einsum", &einsum,
py::arg("expr"), py::arg("a"), py::arg("b"),
py::arg("d"), py::arg("c") = std::nullopt);
py::arg("d"), py::arg("c") = std::nullopt,
py::arg("use_cublaslt") = false);
m.def("fp8_einsum", &fp8_einsum,
py::arg("expr"), py::arg("a"), py::arg("b"),
py::arg("d"), py::arg("c") = std::nullopt,
py::arg("recipe") = std::make_tuple(1, 128, 128));
#endif
}
} // namespace deep_gemm::einsum
csrc/apis/gemm.hpp
+127 -55
View File
@@ -1,16 +1,51 @@
#pragma once
#include "../utils/compatibility.hpp"
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
#include "../jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp"
#include "../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp"
#include "../jit_kernels/impls/sm90_bf16_gemm.hpp"
#include "../jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp"
#include "../jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp"
#include "../jit_kernels/impls/sm100_bf16_gemm.hpp"
#endif
#include "../jit_kernels/impls/smxx_cublaslt.hpp"
#include "layout.hpp"
namespace deep_gemm::gemm {
static bool early_return(const int& m, const int &n, const int& k,
const torch::Tensor& d, const std::optional<torch::Tensor>& c) {
// Do nothing if the problem is empty
if (m == 0 or n == 0)
return true;
// Checks
const bool& is_cd_same = c.has_value() and c->data_ptr() == d.data_ptr();
if (is_cd_same)
DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides());
if (c.has_value()) {
check_major_type_cd(c.value());
DG_HOST_ASSERT(d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(c.value().scalar_type() == torch::kFloat);
}
// No accumulation
if (k == 0) {
if (not is_cd_same)
c.has_value() ? d.copy_(c.value()) : d.zero_();
return true;
}
// With accumulation, do copy before GEMM (assuming the GEMM kernel does not support different C/D)
if (c.has_value() and not is_cd_same)
d.copy_(c.value());
return false;
}
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
static void fp8_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
const std::pair<torch::Tensor, torch::Tensor>& b,
const torch::Tensor& d,
@@ -34,20 +69,12 @@ static void fp8_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
const auto& [n , k_] = get_shape<2>(b.first);
const auto& [m_, n_] = get_shape<2>(d);
DG_HOST_ASSERT(m == m_ and n == n_ and k == k_);
DG_HOST_ASSERT(n > 0 and k > 0);
DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(b.first.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat);
// Check C as well
if (c.has_value()) {
check_major_type_cd(c.value());
DG_HOST_ASSERT(d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(c.value().scalar_type() == torch::kFloat);
}
// Do nothing if the problem is empty
if (m == 0)
// Early return for trivial cases
if (early_return(m, n, k, d, c))
return;
// Transform SFA and SFB into compute-required layout
@@ -63,12 +90,11 @@ static void fp8_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
if (std::get<1>(recipe.value()) == 1) {
sm90_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
} else {
sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
const auto& major_sfb = get_major_type_ab(sfb);
sm90_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, major_sfb, compiled_dims);
}
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
sm100_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kFloat) {
sm100_fp8_gemm_1d2d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
@@ -151,14 +177,12 @@ static void m_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torc
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
const auto& major_sfb = get_major_type_ab(sfb);
sm90_m_grouped_fp8_gemm_contiguous_1d2d(a.first, sfa, b.first, sfb, d, m_indices,
num_groups, m, n, k, major_a, major_b, compiled_dims);
num_groups, m, n, k, major_a, major_b, major_sfb, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
sm100_m_grouped_fp8_gemm_contiguous_1d1d(a.first, sfa, b.first, sfb, d, m_indices,
num_groups, m, n, k, major_a, major_b, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kFloat) {
sm100_m_grouped_fp8_gemm_contiguous_1d2d(a.first, sfa, b.first, sfb, d, m_indices,
num_groups, m, n, k, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
@@ -214,14 +238,12 @@ static void m_grouped_fp8_gemm_nt_masked(const std::pair<torch::Tensor, torch::T
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
const auto& major_sfb = get_major_type_ab(sfb);
sm90_m_grouped_fp8_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims);
num_groups, m, n, k, expected_m, major_a, major_b, major_sfb, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
sm100_m_grouped_fp8_gemm_masked_1d1d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kFloat) {
sm100_m_grouped_fp8_gemm_masked_1d2d(a.first, sfa, b.first, sfb, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture or scaling factor types");
}
@@ -238,22 +260,23 @@ static void k_grouped_fp8_gemm_tn_contiguous(const std::pair<torch::Tensor, torc
// Must be 1D1D kernel
DG_HOST_ASSERT(recipe == std::make_tuple(1, 1, 128));
// Shape checks
const auto& [num_groups, m, n] = get_shape<3>(d);
const auto& [_, m_] = get_shape<2>(a.first);
const auto& [__, n_] = get_shape<2>(b.first);
DG_HOST_ASSERT(m == m_ and n == n_);
// Contiguity checks
DG_HOST_ASSERT(a.first.is_contiguous());
DG_HOST_ASSERT(b.first.is_contiguous());
DG_HOST_ASSERT(d.is_contiguous());
if (c.has_value()) {
DG_HOST_ASSERT(c.value().scalar_type() == torch::kFloat);
DG_HOST_ASSERT(c.value().is_contiguous());
}
DG_HOST_ASSERT(c.has_value() and c.value().is_contiguous());
// Do nothing if empty
if (std::accumulate(ks.begin(), ks.end(), 0) == 0)
// Early return for trivial cases
if (early_return(m, n, std::accumulate(ks.begin(), ks.end(), 0), d, c))
return;
// Transform SF with padding
const auto& [_, m] = get_shape<2>(a.first);
const auto& [__, n] = get_shape<2>(b.first);
const auto& sfa = layout::transform_k_grouped_sf_into_required_layout(a.second, ks, ks_tensor, recipe);
const auto& sfb = layout::transform_k_grouped_sf_into_required_layout(b.second, ks, ks_tensor, recipe);
@@ -285,20 +308,17 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torc
int sum_k = 0;
for (const auto& k: ks)
sum_k += k;
DG_HOST_ASSERT(sum_mk == static_cast<int64_t>(m) * sum_k);
DG_HOST_ASSERT(sum_nk == static_cast<int64_t>(n) * sum_k);
DG_HOST_ASSERT(sum_mk == static_cast<int64_t>(sum_k) * m);
DG_HOST_ASSERT(sum_nk == static_cast<int64_t>(sum_k) * n);
// Contiguity checks
DG_HOST_ASSERT(a.first.is_contiguous());
DG_HOST_ASSERT(b.first.is_contiguous());
DG_HOST_ASSERT(d.is_contiguous());
if (c.has_value()) {
DG_HOST_ASSERT(c.value().scalar_type() == torch::kFloat);
DG_HOST_ASSERT(c.value().is_contiguous());
}
DG_HOST_ASSERT(c.has_value() and c.value().is_contiguous());
// Do nothing if empty
if (std::accumulate(ks.begin(), ks.end(), 0) == 0)
// Early return for trivial cases
if (early_return(m, n, accumulate(ks.begin(), ks.end(), 0), d, c))
return;
// Transform SF with padding
@@ -320,7 +340,9 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torc
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
#endif
#if DG_TENSORMAP_COMPATIBLE
static void bf16_gemm_nt(const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& d,
@@ -338,20 +360,12 @@ static void bf16_gemm_nt(const torch::Tensor& a,
const auto& [n , k_] = get_shape<2>(b);
const auto& [m_, n_] = get_shape<2>(d);
DG_HOST_ASSERT(m == m_ and n == n_ and k == k_);
DG_HOST_ASSERT(n > 0 and k > 0);
DG_HOST_ASSERT(a.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(b.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat);
// Check C as well
if (c.has_value()) {
check_major_type_cd(c.value());
DG_HOST_ASSERT(d.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(c.value().scalar_type() == torch::kFloat);
}
// Do nothing if the problem is empty
if (m == 0)
// Early return for trivial cases
if (early_return(m, n, k, d, c))
return;
// Dispatch into different implements
@@ -396,7 +410,6 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc
const auto& major_a = get_major_type_ab(a);
const auto& major_b = get_major_type_ab(b);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
DG_HOST_ASSERT(major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(m_indices.is_contiguous());
// Type and shape checks
@@ -423,11 +436,21 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc
if (arch_major == 9) {
sm90_m_grouped_bf16_gemm_contiguous(a, b, d, m_indices,
num_groups, m, n, k, major_a, major_b, compiled_dims);
} else if (arch_major == 10) {
sm100_m_grouped_bf16_gemm_contiguous(a, b, d, m_indices,
num_groups, m, n, k, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
static void m_grouped_bf16_gemm_nn_contiguous(const torch::Tensor& a, const torch::Tensor& b,
const torch::Tensor& d, const torch::Tensor& m_indices,
const std::string& compiled_dims) {
m_grouped_bf16_gemm_nt_contiguous(a, b.transpose(1, 2),
d, m_indices, compiled_dims);
}
static void m_grouped_bf16_gemm_nt_masked(const torch::Tensor& a, const torch::Tensor& b,
const torch::Tensor& d, const torch::Tensor& masked_m,
const int& expected_m, const std::string& compiled_dims) {
@@ -458,11 +481,51 @@ static void m_grouped_bf16_gemm_nt_masked(const torch::Tensor& a, const torch::T
if (arch_major == 9) {
sm90_bf16_m_grouped_gemm_masked(a, b, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims);
} else if (arch_major == 10) {
sm100_m_grouped_bf16_gemm_masked(a, b, d, masked_m,
num_groups, m, n, k, expected_m, major_a, major_b, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& d,
const std::vector<int>& ks,
const torch::Tensor& ks_tensor,
const std::optional<torch::Tensor>& c,
const std::string& compiled_dims) {
// Shape checks
const auto& [num_groups, m, n] = get_shape<3>(d);
const auto& [_, m_] = get_shape<2>(a);
const auto& [__, n_] = get_shape<2>(b);
DG_HOST_ASSERT(m == m_ and n == n_);
// Contiguity checks
DG_HOST_ASSERT(a.is_contiguous());
DG_HOST_ASSERT(b.is_contiguous());
DG_HOST_ASSERT(d.is_contiguous());
DG_HOST_ASSERT(c.has_value() and c.value().is_contiguous());
// Early return for trivial cases
if (early_return(m, n, std::accumulate(ks.begin(), ks.end(), 0), d, c))
return;
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
sm90_bf16_k_grouped_gemm(a, b, c, d, m, n, ks, ks_tensor,
cute::UMMA::Major::MN, cute::UMMA::Major::MN, compiled_dims);
} else if (arch_major == 10) {
sm100_bf16_k_grouped_gemm(a, b, c, d, m, n, ks, ks_tensor,
cute::UMMA::Major::MN, cute::UMMA::Major::MN, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}
#endif
static void cublaslt_gemm_nt(const torch::Tensor& a, const torch::Tensor& b,
const torch::Tensor& d, const std::optional<torch::Tensor>& c) {
// Shape must be `[M, K] @ [N, K].T`
@@ -475,11 +538,8 @@ static void cublaslt_gemm_nt(const torch::Tensor& a, const torch::Tensor& b,
const auto& [m_, n_] = get_shape<2>(d);
DG_HOST_ASSERT(m == m_ and n == n_ and k == k_);
if (c.has_value())
DG_HOST_ASSERT(c.value().scalar_type() == d.scalar_type());
// Do nothing if the problem is empty
if (m == 0 or n == 0)
// Early return for trivial cases
if (early_return(m, n, k, d, c))
return;
cublaslt_gemm(a, b, c, d, m, n, k, major_a, major_b);
@@ -501,8 +561,10 @@ static void cublaslt_gemm_tt(const torch::Tensor& a, const torch::Tensor& b,
}
static void register_apis(pybind11::module_& m) {
#if DG_FP8_COMPATIBLE and DG_TENSORMAP_COMPATIBLE
// FP8 GEMMs
m.def("fp8_gemm_nt", &fp8_gemm_nt,
m.def("fp8_gemm_nt", &fp8_gemm_nt,
py::arg("a"), py::arg("b"), py::arg("d"),
py::arg("c") = std::nullopt, py::arg("recipe") = std::nullopt,
py::arg("compiled_dims") = "nk",
@@ -544,7 +606,9 @@ static void register_apis(pybind11::module_& m) {
py::arg("ks_tensor"), py::arg("c") = std::nullopt,
py::arg("recipe") = std::make_tuple(1, 1, 128),
py::arg("compiled_dims") = "mn");
#endif
#if DG_TENSORMAP_COMPATIBLE
// BF16 GEMMs
m.def("bf16_gemm_nt", &bf16_gemm_nt,
py::arg("a"), py::arg("b"), py::arg("d"),
@@ -565,9 +629,17 @@ static void register_apis(pybind11::module_& m) {
m.def("m_grouped_bf16_gemm_nt_contiguous", &m_grouped_bf16_gemm_nt_contiguous,
py::arg("a"), py::arg("b"), py::arg("d"), py::arg("m_indices"),
py::arg("compiled_dims") = "nk");
m.def("m_grouped_bf16_gemm_nn_contiguous", &m_grouped_bf16_gemm_nn_contiguous,
py::arg("a"), py::arg("b"), py::arg("d"), py::arg("m_indices"),
py::arg("compiled_dims") = "nk");
m.def("m_grouped_bf16_gemm_nt_masked", &m_grouped_bf16_gemm_nt_masked,
py::arg("a"), py::arg("b"), py::arg("d"), py::arg("masked_m"),
py::arg("expected_m"), py::arg("compiled_dims") = "nk");
m.def("k_grouped_bf16_gemm_tn_contiguous", &k_grouped_bf16_gemm_tn_contiguous,
py::arg("a"), py::arg("b"), py::arg("d"), py::arg("ks"),
py::arg("ks_tensor"), py::arg("c") = std::nullopt,
py::arg("compiled_dims") = "mn");
#endif
// cuBLASLt GEMMs
m.def("cublaslt_gemm_nt", &cublaslt_gemm_nt,
csrc/apis/layout.hpp
+1 -1
View File
@@ -28,7 +28,7 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
return get_mn_major_tma_aligned_packed_ue8m0_tensor(sf);
}
// (FP32, 128, 128) on SM90: no need to transform, check shape and contiguous
// (FP32, 128, 128) on SM90: no need to transform, check SFB requirements
if (sf.scalar_type() == torch::kFloat and gran_mn == 128 and gran_k == 128 and (arch_major == 9 or disable_ue8m0_cast))
return check_sf_layout(sf, mn, k, gran_mn, gran_k, num_groups, false, true, torch::kFloat);
csrc/indexing/main.cu
+1 -2
View File
@@ -4,14 +4,12 @@
#include <deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh>
#include <deep_gemm/impls/sm100_bf16_gemm.cuh>
#include <deep_gemm/impls/sm100_fp8_gemm_1d1d.cuh>
#include <deep_gemm/impls/sm100_fp8_gemm_1d2d.cuh>
// Attention kernels
#include <deep_gemm/impls/sm90_fp8_mqa_logits.cuh>
#include <deep_gemm/impls/sm90_fp8_paged_mqa_logits.cuh>
#include <deep_gemm/impls/sm100_fp8_mqa_logits.cuh>
#include <deep_gemm/impls/sm100_fp8_paged_mqa_logits.cuh>
#include <deep_gemm/impls/smxx_clean_logits.cuh>
// Einsum kernels
#include <deep_gemm/impls/sm90_bmk_bnk_mn.cuh>
@@ -19,6 +17,7 @@
// Layout kernels
#include <deep_gemm/impls/smxx_layout.cuh>
#include <deep_gemm/impls/smxx_clean_logits.cuh>
using namespace deep_gemm;
csrc/jit/compiler.hpp
+14 -5
View File
@@ -26,12 +26,17 @@ public:
static std::string library_version;
static std::string get_library_version() {
std::stringstream ss;
std::vector<char> buffer;
for (const auto& f: collect_files(library_include_path / "deep_gemm")) {
std::ifstream in(f, std::ios::binary);
ss << in.rdbuf();
DG_HOST_ASSERT(in.is_open());
// Append into the buffer
buffer.insert(buffer.end(),
std::istreambuf_iterator<char>(in),
std::istreambuf_iterator<char>());
}
return get_hex_digest(ss.str());
return get_hex_digest(buffer);
}
static void prepare_init(const std::string& library_root_path,
@@ -62,8 +67,8 @@ public:
flags = fmt::format("-std=c++{} --diag-suppress=39,161,174,177,186,940 "
"--ptxas-options=--register-usage-level=10",
get_env<int>("DG_JIT_CPP_STANDARD", 20));
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0))
flags += " --ptxas-options=--verbose";
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0) or get_env("DG_JIT_PTXAS_CHECK", 0))
flags += " --ptxas-options=--verbose,--warn-on-local-memory-usage";
if (get_env("DG_JIT_WITH_LINEINFO", 0))
flags += " -Xcompiler -rdynamic -lineinfo";
}
@@ -178,6 +183,10 @@ public:
DG_HOST_ASSERT(false and "NVCC compilation failed");
}
// Check local memory usage
if (get_env("DG_JIT_PTXAS_CHECK", 0))
DG_HOST_ASSERT(not std::regex_search(output, std::regex(R"(Local memory used)")));
// Print PTXAS log
if (get_env("DG_JIT_DEBUG", 0) or get_env("DG_JIT_PTXAS_VERBOSE", 0))
printf("%s", output.c_str());
csrc/jit/device_runtime.hpp
+21 -1
View File
@@ -1,6 +1,7 @@
#pragma once
#include <cublasLt.h>
#include <torch/version.h>
#include <ATen/cuda/CUDAContext.h>
#include "../utils/exception.hpp"
@@ -14,10 +15,24 @@ class DeviceRuntime {
// cuBLASLt utils
static constexpr size_t kCublasLtWorkspaceSize = 32 * 1024 * 1024;
public:
#if TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 3)
// For PyTorch 2.3+, share the PyTorch cuBLASLt handle
DeviceRuntime() = default;
static cublasLtHandle_t get_cublaslt_handle() {
return at::cuda::getCurrentCUDABlasLtHandle();
}
static torch::Tensor get_cublaslt_workspace() {
return torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA));
}
#else
// Otherwise, create the cuBLASLt handle ourselves
cublasLtHandle_t cublaslt_handle{};
std::shared_ptr<torch::Tensor> cublaslt_workspace;
public:
explicit DeviceRuntime() {
cublaslt_workspace = std::make_shared<torch::Tensor>(torch::empty({kCublasLtWorkspaceSize}, dtype(torch::kByte).device(at::kCUDA)));
DG_CUBLASLT_CHECK(cublasLtCreate(&cublaslt_handle));
@@ -34,6 +49,7 @@ public:
torch::Tensor get_cublaslt_workspace() const {
return *cublaslt_workspace;
}
#endif
std::shared_ptr<cudaDeviceProp> get_prop() {
if (cached_prop == nullptr) {
@@ -77,6 +93,10 @@ public:
return num_sms;
}
int get_l2_cache_size() {
return get_prop()->l2CacheSize;
}
void set_tc_util(const int& new_tc_util) {
DG_HOST_ASSERT(0 <= new_tc_util and new_tc_util <= 100);
tc_util = new_tc_util;
csrc/jit/handle.hpp
+42 -5
View File
@@ -2,12 +2,49 @@
#include <cuda.h>
#include <cuda_runtime.h>
#include <dlfcn.h>
#include <filesystem>
#include "../utils/exception.hpp"
#include "../utils/compatibility.hpp"
namespace deep_gemm {
// Lazy loading all driver symbols
static void* get_driver_handle() {
static void* handle = nullptr;
if (handle == nullptr) {
handle = dlopen("libcuda.so.1", RTLD_LAZY | RTLD_LOCAL);
DG_HOST_ASSERT(handle != nullptr and "Failed to load CUDA driver `libcuda.so.1`");
}
return handle;
}
// Macro to define wrapper functions named `lazy_cu{API name}`
#define DECL_LAZY_CUDA_DRIVER_FUNCTION(name) \
template <typename... Args> \
static auto lazy_##name(Args&&... args) -> decltype(name(args...)) { \
using FuncType = decltype(&name); \
static FuncType func = nullptr; \
if (func == nullptr) { \
func = reinterpret_cast<FuncType>(dlsym(get_driver_handle(), #name)); \
DG_HOST_ASSERT(func != nullptr and "Failed to load CUDA driver API"); \
} \
return func(std::forward<decltype(args)>(args)...); \
}
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorName);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuGetErrorString);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuFuncSetAttribute);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleLoad);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleUnload);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuModuleGetFunction);
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuLaunchKernelEx);
#if DG_TENSORMAP_COMPATIBLE
DECL_LAZY_CUDA_DRIVER_FUNCTION(cuTensorMapEncodeTiled);
#endif
#if CUDART_VERSION >= 12080 and defined(DG_JIT_USE_RUNTIME_API)
// Use CUDA runtime API
@@ -80,8 +117,8 @@ static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const s
LibraryHandle *library_opt = nullptr) {
LibraryHandle library;
KernelHandle kernel;
DG_CUDA_DRIVER_CHECK(cuModuleLoad(&library, cubin_path.c_str()));
DG_CUDA_DRIVER_CHECK(cuModuleGetFunction(&kernel, library, func_name.c_str()));
DG_CUDA_DRIVER_CHECK(lazy_cuModuleLoad(&library, cubin_path.c_str()));
DG_CUDA_DRIVER_CHECK(lazy_cuModuleGetFunction(&kernel, library, func_name.c_str()));
if (library_opt != nullptr)
*library_opt = library;
@@ -89,7 +126,7 @@ static KernelHandle load_kernel(const std::filesystem::path& cubin_path, const s
}
static void unload_library(const LibraryHandle& library) {
const auto& error = cuModuleUnload(library);
const auto& error = lazy_cuModuleUnload(library);
DG_HOST_ASSERT(error == CUDA_SUCCESS or error == CUDA_ERROR_DEINITIALIZED);
}
@@ -97,7 +134,7 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
const cudaStream_t& stream, const int& smem_size,
const dim3& grid_dim, const dim3& block_dim, const int& cluster_dim) {
if (smem_size > 0)
DG_CUDA_DRIVER_CHECK(cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
DG_CUDA_DRIVER_CHECK(lazy_cuFuncSetAttribute(kernel, CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, smem_size));
LaunchConfigHandle config;
config.gridDimX = grid_dim.x;
@@ -127,7 +164,7 @@ static LaunchConfigHandle construct_launch_config(const KernelHandle& kernel,
template<typename... ActTypes>
static auto launch_kernel(const KernelHandle& kernel, const LaunchConfigHandle& config, ActTypes&&... args) {
void *ptr_args[] = { &args... };
return cuLaunchKernelEx(&config, kernel, ptr_args, nullptr);
return lazy_cuLaunchKernelEx(&config, kernel, ptr_args, nullptr);
}
#endif
csrc/jit_kernels/heuristics/common.hpp
+5 -4
View File
@@ -4,6 +4,7 @@
#include "../../utils/math.hpp"
#include "../../utils/layout.hpp"
#include "../../utils/system.hpp"
namespace deep_gemm {
@@ -156,12 +157,12 @@ static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& k
DG_HOST_ASSERT(cd_dtype == torch::kBFloat16 or cd_dtype == torch::kFloat);
// Select M/N block sizes
auto block_ms = std::vector{64, 128, 256};
auto block_ms = ArchSpec::get_block_m_candidates(kernel_type, major_a, m);
if (gemm_type == GemmType::MGroupedContiguous)
block_ms = std::vector{get_mk_alignment_for_contiguous_layout()};
if (gemm_type == GemmType::MGroupedMasked) // Exclude 256 for performance
block_ms = std::vector{64, 128};
const auto block_ns = ArchSpec::get_block_n_candidates(cd_dtype);
const auto block_ns = ArchSpec::get_block_n_candidates(kernel_type, cd_dtype);
// K block size is selected in a fixed manner
const auto& block_k = 128 / static_cast<int>(c10::elementSize(ab_dtype));
@@ -185,7 +186,7 @@ static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& k
for (const auto& block_n: block_ns) {
const int& num_waves = get_num_waves(block_m, block_n);
const auto& last_util = get_last_wave_util(block_m, block_n);
if (not ArchSpec::is_block_size_legal(kernel_type, major_a, major_b, ab_dtype, cd_dtype, block_m, block_n, block_k))
if (not ArchSpec::is_block_size_legal(kernel_type, major_a, major_b, ab_dtype, cd_dtype, m, n, k, block_m, block_n, block_k))
continue;
bool success = false;
@@ -234,7 +235,7 @@ static GemmConfig get_best_config(const GemmType& gemm_type, const KernelType& k
constexpr int smem_capacity = ArchSpec::smem_capacity;
int best_num_stages = 0;
SharedMemoryConfig best_smem_config;
for (int num_stages = 12; num_stages > 0; -- num_stages) {
for (int num_stages = 32; num_stages > 0; -- num_stages) {
if (not ArchSpec::is_num_stages_legal(ab_dtype, cd_dtype, num_stages, best_block_m, best_block_n, block_k))
continue;
csrc/jit_kernels/heuristics/sm100.hpp
+24 -20
View File
@@ -12,7 +12,17 @@ namespace deep_gemm {
struct SM100ArchSpec {
static constexpr int smem_capacity = 232448;
static std::vector<int> get_block_n_candidates(const at::ScalarType& cd_dtype) {
static std::vector<int> get_block_m_candidates(const KernelType& kernel_type, const cute::UMMA::Major& major_a, const int& m) {
std::vector<int> candidates{128, 256};
if ((kernel_type == KernelType::Kernel1D1D or kernel_type == KernelType::KernelNoSF) and major_a == cute::UMMA::Major::K) {
// NOTES: `block_m = 32/64` is smaller than `LAYOUT_AD_M`, should be careful in handling this
if (m <= 32) candidates.push_back(32);
if (m <= 64) candidates.push_back(64);
}
return candidates;
}
static std::vector<int> get_block_n_candidates(const KernelType& kernel_type, const at::ScalarType& cd_dtype) {
// 16 is for better SM usage
// Stride 32 is due to low-performance swizzle-16/32B
std::vector<int> candidates = {16};
@@ -45,7 +55,6 @@ struct SM100ArchSpec {
static std::pair<int, int> get_sf_uttcp_aligned_block_sizes(
const int& block_m, const int& block_n, const at::ScalarType& ab_dtype) {
constexpr int num_utccp_aligned_elems = 128;
DG_HOST_ASSERT(block_m % num_utccp_aligned_elems == 0);
switch (ab_dtype) {
case torch::kBFloat16: return {0, 0};
case torch::kFloat8_e4m3fn: return {align(block_m, num_utccp_aligned_elems), align(block_n, num_utccp_aligned_elems)};
@@ -56,23 +65,18 @@ struct SM100ArchSpec {
static bool is_block_size_legal(const KernelType& kernel_type,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
const int& m, const int& n, const int& k,
const int& block_m, const int& block_n, const int& block_k) {
// TODO: consider more carefully for BF16 GEMMs
// 2SM BF16 UMMA does not support `N % 32 != 0`
if (ab_dtype == torch::kBFloat16 and block_n % 32 != 0)
return false;
// Layout A/D does not support `block_m == 64` and `block_n % 16 != 0`
if (block_m == 64 or block_n % 16 != 0)
// Layout A/D does not support `block_n % 16 != 0`
if (block_n % 16 != 0)
return false;
// Performance is lower with 1D1D and `block_m == 256`
if (kernel_type == KernelType::Kernel1D1D and major_b == cute::UMMA::Major::K and block_m != 128)
if (kernel_type == KernelType::Kernel1D1D and major_b == cute::UMMA::Major::K and block_m > 128)
return false;
// 1D2D kernels' maximum block N is 128
// 1D2D kernels require more friendly block Ns
if (kernel_type == KernelType::Kernel1D2D and (block_n > 128 or 128 % block_n != 0))
// For small K, fewer store blocks improve store/compute overlap and reduce epilogue bottleneck
if (k <= 256 and (block_n > 128 or block_m > 128))
return false;
// Check tensor memory validity
@@ -96,22 +100,23 @@ struct SM100ArchSpec {
}
static bool should_minimize_num_sms() {
return false;
return true;
}
static std::pair<bool, bool> get_multicast_legality(const GemmType& gemm_type, const int& num_groups,
const int& m, const int& n, const int& block_m, const int& block_n,
const int& num_sms) {
const int& m, const int& n, const int& block_m, const int& block_n,
const int& num_sms) {
// TODO: support other layouts
return {
false,
is_multicast_legal(m, block_m, 2, num_sms, true) and (gemm_type == GemmType::Normal or gemm_type == GemmType::KGroupedContiguous),
is_multicast_legal(m, block_m, 2, num_sms, true) and (gemm_type == GemmType::Normal or gemm_type == GemmType::KGroupedContiguous
or (gemm_type == GemmType::Batched and num_groups <= 32)),
};
}
static ThreadConfig get_thread_config(const KernelType& kernel_type,
const int& block_m, const int& block_n) {
return ThreadConfig::sm100(128, kernel_type == KernelType::Kernel1D2D ? block_m : 128);
return ThreadConfig::sm100(128, 128);
}
static int get_smem_cd_size(const KernelType& kernel_type,
@@ -119,7 +124,7 @@ struct SM100ArchSpec {
const int& swizzle_cd_mode,
const at::ScalarType& cd_dtype) {
constexpr static int layout_ad_m = 128;
return (kernel_type != KernelType::Kernel1D2D ? std::min(block_m, layout_ad_m) : block_m) * swizzle_cd_mode * 2;
return std::min(block_m, layout_ad_m) * swizzle_cd_mode * 2;
}
static std::pair<int, int> get_sf_smem_size_per_stage(const KernelType& kernel_type,
@@ -149,7 +154,6 @@ struct SM100ArchSpec {
static int get_barrier_smem_size(const int& num_stages) {
// TODO: remove SF barriers for BF16 GEMMs
// TMA full/empty barriers, with-SF full barriers, tensor memory full/empty barriers
// NOTES: 1D2D kernel will not use the with-SF full barriers
// NOTES: some shapes may only have 1 epilogue stage, but we still allocate space for 2 stages
// NOTES: the last barrier is for tensor core utilization control
return num_stages * 8 * 3 + 2 * 8 * 2 + 8;
csrc/jit_kernels/heuristics/sm90.hpp
+34 -13
View File
@@ -10,11 +10,28 @@ namespace deep_gemm {
struct SM90ArchSpec {
static constexpr int smem_capacity = 232448;
static std::vector<int> get_block_m_candidates(const KernelType& kernel_type, const cute::UMMA::Major& major_a, const int& m) {
std::vector<int> candidates{64, 128, 256};
if ((kernel_type == KernelType::Kernel1D2D or kernel_type == KernelType::KernelNoSF) and major_a == cute::UMMA::Major::K) {
// NOTES: `block_m = 16/32` is smaller than MMA M size, should be careful in handling this
if (m <= 16) candidates.push_back(16);
if (m <= 32) candidates.push_back(32);
}
return candidates;
}
static std::vector<int> get_block_n_candidates(const at::ScalarType& cd_dtype) {
// Avoid bank conflicts for FP32 output
const auto& start = cd_dtype == torch::kFloat ? 8 : 16;
static std::vector<int> get_block_n_candidates(const KernelType& kernel_type, const at::ScalarType& cd_dtype) {
int start = 16;
// Avoid bank conflicts for 1D1D kernel FP32 output
std::vector<int> candidates;
if (kernel_type == KernelType::Kernel1D1D and cd_dtype == torch::kFloat) {
candidates.push_back(16);
start = 24;
}
// Push the strided options
for (int i = start; i <= 256; i += 16)
candidates.push_back(i);
return candidates;
@@ -44,6 +61,7 @@ struct SM90ArchSpec {
static bool is_block_size_legal(const KernelType& kernel_type,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const at::ScalarType& ab_dtype, const at::ScalarType& cd_dtype,
const int& m, const int& n, const int& k,
const int& block_m, const int& block_n, const int& block_k) {
// SM90 FP32 output does not support `block_m == 256`
if (cd_dtype == at::kFloat and block_m == 256)
@@ -58,15 +76,15 @@ struct SM90ArchSpec {
return false;
}
// When B is N Major, use swizzle 128B for better performance; only affects SM90 BF16 GEMM
if (major_b == cute::UMMA::Major::MN and block_n >= 128 and block_n % 64 != 0)
return false;
// Too many scaling factors in a single block: `block_n > block_k and std::gcd(block_n, block_k) != block_n - block_k`
// Or too many register spills
if (block_n > 128 and kernel_type == KernelType::Kernel1D2D and (block_n != 144 and block_n != 160 and block_n != 192))
return false;
// Avoid bank conflicts for FP32 output
if (cd_dtype == torch::kFloat and block_n % 16 == 0)
return false;
// The block sizes cannot be too large (for enough registers), so at least one dim less than 128
return block_m <= 128 or block_n <= 128;
}
@@ -91,6 +109,9 @@ struct SM90ArchSpec {
if (gemm_type == GemmType::KGroupedContiguous and num_groups > 4)
return {false, false};
if (gemm_type == GemmType::Batched)
return {false, false};
return {
is_multicast_legal(n, block_n, 2, num_sms, gemm_type == GemmType::MGroupedMasked),
// For masked GEMM layout, divisibility on N is also required as we must ensure the total number of blocks is even
@@ -101,13 +122,14 @@ struct SM90ArchSpec {
static ThreadConfig get_thread_config(const KernelType& kernel_type,
const int& block_m, const int& block_n) {
return ThreadConfig::sm90(128, (block_m == 64 ? 1 : 2) * 128);
return ThreadConfig::sm90(128, (block_m <= 64 ? 1 : 2) * 128);
}
static int get_smem_cd_size(const KernelType& kernel_type,
const int& block_m, const int& block_n,
const int& swizzle_cd_mode, const at::ScalarType& cd_dtype) {
return block_m * block_n * static_cast<int>(c10::elementSize(cd_dtype));
// NOTES: 1024 is for TMA swizzling alignment requirement
return align(block_m * block_n * static_cast<int>(c10::elementSize(cd_dtype)), 1024);
}
static std::pair<int, int> get_sf_smem_size_per_stage(const KernelType& kernel_type,
@@ -116,12 +138,11 @@ struct SM90ArchSpec {
if (ab_dtype == torch::kBFloat16)
return {0, 0};
int smem_sfa_per_stage = block_m * static_cast<int>(sizeof(float));
// NOTES: 128 is for 2D TMA alignment requirement
int smem_sfa_per_stage = align(block_m * static_cast<int>(sizeof(float)), 128);
int smem_sfb_per_stage = 0;
if (kernel_type == KernelType::Kernel1D1D) {
// NOTES: `128` is for 2D TMA alignment requirement
if (kernel_type == KernelType::Kernel1D1D)
smem_sfb_per_stage = align(block_n * 4, 128);
}
return {smem_sfa_per_stage, smem_sfb_per_stage};
}
csrc/jit_kernels/impls/runtime_utils.hpp
+9 -7
View File
@@ -3,8 +3,9 @@
#include <cuda.h>
#include <torch/python.h>
#include "../../utils/math.hpp"
#include "../heuristics/sm90.hpp"
#include "../../jit/handle.hpp"
#include "../../utils/math.hpp"
#include "../../utils/system.hpp"
#include "../../utils/exception.hpp"
@@ -40,6 +41,7 @@ static std::string to_string(const GemmType& type) {
case GemmType::MGroupedContiguous: return "GemmType::MGroupedContiguous";
case GemmType::MGroupedMasked: return "GemmType::MGroupedMasked";
case GemmType::KGroupedContiguous: return "GemmType::KGroupedContiguous";
case GemmType::Batched: return "GemmType::Batched";
}
DG_HOST_UNREACHABLE("Unknown GEMM type");
}
@@ -68,7 +70,7 @@ static CUtensorMapDataType aten_dtype_to_tensor_map_dtype(const at::ScalarType&
}
static CUtensorMapSwizzle mode_into_tensor_map_swizzle(const int& mode, const int& base) {
#if CUDA_VERSION >= 12080
#if CUDART_VERSION >= 12080
if (base != 0) {
DG_HOST_ASSERT(base == 32 and mode == 128);
return CU_TENSOR_MAP_SWIZZLE_128B_ATOM_32B;
@@ -106,7 +108,7 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
gmem_inner_dim, gmem_outer_dim, smem_inner_dim, smem_outer_dim,
gmem_outer_stride, swizzle_mode, swizzle_base, elem_size);
}
DG_CUDA_DRIVER_CHECK(cuTensorMapEncodeTiled(
DG_CUDA_DRIVER_CHECK(lazy_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),
@@ -115,14 +117,14 @@ static CUtensorMap make_tma_2d_desc(const torch::Tensor& t,
}
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,
int gmem_dim_0, int gmem_dim_1, int gmem_dim_2,
int smem_dim_0, int smem_dim_1, 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);
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),};
@@ -134,7 +136,7 @@ static CUtensorMap make_tma_3d_desc(const torch::Tensor& t,
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);
}
DG_CUDA_DRIVER_CHECK(cuTensorMapEncodeTiled(
DG_CUDA_DRIVER_CHECK(lazy_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),
csrc/jit_kernels/impls/sm100_bf16_gemm.hpp
+263 -26
View File
@@ -25,8 +25,7 @@ public:
void* grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
CUtensorMap tensor_map_c;
CUtensorMap tensor_map_d;
CUtensorMap tensor_map_cd;
};
static std::string generate_impl(const Args& args) {
@@ -69,7 +68,7 @@ static void __instantiate_kernel() {{
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.grouped_layout, args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_c, args.tensor_map_d));
args.tensor_map_cd));
}
};
@@ -87,7 +86,6 @@ static void sm100_bf16_gemm(const torch::Tensor& a,
torch::kBFloat16, d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto& cd = c.value_or(d);
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
@@ -98,26 +96,11 @@ static void sm100_bf16_gemm(const torch::Tensor& a,
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_c = make_tma_cd_desc(cd, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(cd.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
// Duplicate the accumulator if necessary
if (c.has_value()) {
if (c->data_ptr() == d.data_ptr()) {
DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides());
} else {
// ReSharper disable once CppExpressionWithoutSideEffects
d.copy_(c.value());
}
}
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
@@ -131,12 +114,266 @@ static void sm100_bf16_gemm(const torch::Tensor& a,
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_c = tensor_map_c,
.tensor_map_d = tensor_map_d
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_gemm", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
static void sm100_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& d,
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& aligned_k = align(k, 64);
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedContiguous, KernelType::KernelNoSF,
// NOTES: `num_groups` is 1, since the contiguous layout is seen as a whole
m, n, k, 1, major_a, major_b,
torch::kBFloat16, d.scalar_type(), false,
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = aligned_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_m_grouped_gemm_contiguous", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
static void sm100_m_grouped_bf16_gemm_masked(const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& d,
const torch::Tensor& masked_m,
const int& num_groups, const int& m, const int& n, const int& k,
const int& expected_m,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& aligned_k = align(k, 64);
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedMasked, KernelType::KernelNoSF,
expected_m, n, k, num_groups, major_a, major_b,
torch::kBFloat16, d.scalar_type(), false,
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = aligned_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_m_grouped_gemm_masked", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
static void sm100_bf16_k_grouped_gemm(const torch::Tensor& a,
const torch::Tensor& b,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n,
const std::vector<int>& ks, const torch::Tensor& ks_tensor,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
int sum_k = 0;
for (const auto& k: ks) {
sum_k += k;
DG_HOST_ASSERT(k % 128 == 0);
}
const auto& num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::KGroupedContiguous, KernelType::KernelNoSF,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
torch::kBFloat16, d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
// Create tensor descriptors
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(0)), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
// Launch kernel
const SM100BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_k_grouped_gemm", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
static void sm100_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_b,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, d, r, h, cute::UMMA::Major::K, cute::UMMA::Major::K,
torch::kBFloat16, tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = b, .n = d, .k = r,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_bhr_hdr_bhd", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
static void sm100_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_b,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, r, d, h, cute::UMMA::Major::K, cute::UMMA::Major::MN,
torch::kBFloat16, tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
// Launch
const SM100BF16GemmRuntime::Args& args = {
.m = b, .n = r, .k = d,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm100_bf16_bhd_hdr_bhr", code);
SM100BF16GemmRuntime::launch(runtime, args);
}
} // namespace deep_gemm
csrc/jit_kernels/impls/sm100_bmk_bnk_mn.hpp
+1 -1
View File
@@ -134,4 +134,4 @@ static void sm100_bmn_bnk_mn_gemm(const torch::Tensor &a,
SM100BmkBnkMnRuntime::launch(runtime, args);
}
} // namespace deep_gemm
} // namespace deep_gemm
csrc/jit_kernels/impls/sm100_fp8_gemm_1d1d.hpp
+92 -58
View File
@@ -30,8 +30,7 @@ public:
CUtensorMap tensor_map_b;
CUtensorMap tensor_map_sfa;
CUtensorMap tensor_map_sfb;
CUtensorMap tensor_map_c;
CUtensorMap tensor_map_d;
CUtensorMap tensor_map_cd;
};
static std::string generate_impl(const Args& args) {
@@ -75,7 +74,7 @@ static void __instantiate_kernel() {{
args.grouped_layout, args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_sfa, args.tensor_map_sfb,
args.tensor_map_c, args.tensor_map_d));
args.tensor_map_cd));
}
};
@@ -105,31 +104,16 @@ static void sm100_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_c = make_tma_cd_desc(cd, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(cd.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(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);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.block_n, config.block_k, 1, 0);
// Duplicate the accumulator if necessary
if (c.has_value()) {
if (c->data_ptr() == d.data_ptr()) {
DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides());
} else {
// ReSharper disable once CppExpressionWithoutSideEffects
d.copy_(c.value());
}
}
// Launch
const SM100FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = aligned_k,
@@ -145,8 +129,7 @@ static void sm100_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_c = tensor_map_c,
.tensor_map_d = tensor_map_d
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_fp8_gemm_1d1d", code);
@@ -163,6 +146,7 @@ static void sm100_m_grouped_fp8_gemm_contiguous_1d1d(const torch::Tensor& a, con
const auto& aligned_k = align(k, 128);
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedContiguous, KernelType::Kernel1D1D,
// NOTES: `num_groups` is 1, since the contiguous layout is seen as a whole
m, n, k, 1, major_a, major_b,
torch::kFloat8_e4m3fn, d.scalar_type(), false,
device_runtime->get_num_sms());
@@ -178,11 +162,11 @@ static void sm100_m_grouped_fp8_gemm_contiguous_1d1d(const torch::Tensor& a, con
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(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);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
@@ -203,8 +187,7 @@ static void sm100_m_grouped_fp8_gemm_contiguous_1d1d(const torch::Tensor& a, con
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_c = tensor_map_d,
.tensor_map_d = tensor_map_d
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_m_grouped_fp8_gemm_contiguous_1d1d", code);
@@ -237,11 +220,11 @@ static void sm100_m_grouped_fp8_gemm_masked_1d1d(const torch::Tensor& a, const t
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
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, num_groups, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
@@ -262,8 +245,7 @@ static void sm100_m_grouped_fp8_gemm_masked_1d1d(const torch::Tensor& a, const t
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_c = tensor_map_d,
.tensor_map_d = tensor_map_d
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_fp8_m_grouped_gemm_masked_1d1d", code);
@@ -296,7 +278,6 @@ static void fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor&
device_runtime->get_num_sms());
// Create tensor descriptors
const auto& cd = c.value_or(d);
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
@@ -307,27 +288,16 @@ static void fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor&
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_c = make_tma_cd_desc(cd, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(cd.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_sfa = make_tma_sf_desc(cute::UMMA::Major::MN, sfa, m, sum_sf_k * 512,
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * 512,
config.block_n, config.block_k, 1, 0);
// Duplicate the accumulator if necessary
if (c.has_value()) {
DG_HOST_ASSERT(c->data_ptr() == d.data_ptr());
DG_HOST_ASSERT(c->sizes() == d.sizes() and c->strides() == d.strides());
}
// Launch kernel
const SM100FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
@@ -343,12 +313,76 @@ static void fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Tensor&
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_c = tensor_map_c,
.tensor_map_d = tensor_map_d
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_fp8_k_grouped_gemm_1d1d", code);
SM100FP8Gemm1D1DRuntime::launch(runtime, args);
}
static void sm100_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& batch_size, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Batched, KernelType::Kernel1D1D,
m, n, k, batch_size, major_a, major_b,
torch::kFloat8_e4m3fn, d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const int& load_block_m = SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& [inner_dim_a, outer_dim_a] = get_inner_outer_dims(major_a, k, m);
const auto& [inner_block_a, outer_block_a] = get_inner_outer_dims(major_a, config.block_k, load_block_m);
const auto& tensor_map_a = make_tma_3d_desc(a, inner_dim_a, outer_dim_a, batch_size,
inner_block_a, outer_block_a, 1,
a.stride(major_a == cute::UMMA::Major::K ? 1 : 2),
a.stride(0),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& [inner_dim_b, outer_dim_b] = get_inner_outer_dims(major_b, k, n);
const auto& [inner_block_b, outer_block_b] = get_inner_outer_dims(major_b, config.block_k, load_block_n);
const auto& tensor_map_b = make_tma_3d_desc(b, inner_dim_b, outer_dim_b, batch_size,
inner_block_b, outer_block_b, 1,
b.stride(major_b == cute::UMMA::Major::K ? 1 : 2),
b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM100ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM100ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(d, n, m, batch_size,
store_block_n, store_block_m, 1,
d.stride(1), d.stride(0),
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, batch_size, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.block_n, config.block_k, batch_size, 0);
// Launch
const SM100FP8Gemm1D1DRuntime::Args& args = {
.m = m, .n = n, .k = k,
.num_groups = batch_size,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_cd = tensor_map_cd
};
const auto& code = SM100FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_fp8_gemm_1d1d", code);
SM100FP8Gemm1D1DRuntime::launch(runtime, args);
}
} // namespace deep_gemm
csrc/jit_kernels/impls/sm100_fp8_gemm_1d2d.hpp
-244
View File
@@ -1,244 +0,0 @@
#pragma once
#include <torch/python.h>
#include "../../jit/compiler.hpp"
#include "../../jit/device_runtime.hpp"
#include "../../jit/kernel_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "../../utils/math.hpp"
#include "../heuristics/sm100.hpp"
#include "runtime_utils.hpp"
namespace deep_gemm {
class SM100FP8Gemm1D2DRuntime final: public LaunchRuntime<SM100FP8Gemm1D2DRuntime> {
public:
struct Args {
int m, n, k, num_groups;
const std::string& compiled_dims;
const std::optional<std::string>& epilogue_type;
GemmConfig gemm_config;
LaunchArgs launch_args;
void *sfb, *grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
CUtensorMap tensor_map_d;
CUtensorMap tensor_map_sfa;
};
static std::string generate_impl(const Args& args) {
return fmt::format(R"(
#include <deep_gemm/impls/sm100_fp8_gemm_1d2d.cuh>
using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp8_gemm_1d2d_impl<
{}, {},
{}, {}, {},
{}, {}, {},
{},
{}, {}, {},
{}, {},
{}, {},
{}, {},
{},
{}, {},
{}
>);
}};
)",
to_string(args.gemm_config.major_a), to_string(args.gemm_config.major_b),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.num_groups,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages, args.gemm_config.num_last_stages,
args.gemm_config.thread_config.num_non_epilogue_threads, args.gemm_config.thread_config.num_epilogue_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), to_string(args.gemm_config.cd_dtype),
get_default_epilogue_type(args.epilogue_type));
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.sfb, args.grouped_layout,
args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_d, args.tensor_map_sfa));
}
};
static void sm100_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims,
const std::optional<std::string>& epilogue_type = std::nullopt) {
DG_HOST_ASSERT(not c.has_value());
const auto& aligned_k = align(k, 128);
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::Normal, KernelType::Kernel1D2D,
m, n, k, 1, major_a, major_b,
torch::kFloat8_e4m3fn, d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, static_cast<int>(d.size(-1)),
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(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);
// Launch
const SM100FP8Gemm1D2DRuntime::Args& args = {
.m = m, .n = n, .k = aligned_k,
.num_groups = 1,
.compiled_dims = compiled_dims,
.epilogue_type = epilogue_type,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.sfb = sfb.data_ptr(),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM100FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_fp8_gemm_1d2d", code);
SM100FP8Gemm1D2DRuntime::launch(runtime, args);
}
static void sm100_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& aligned_k = align(k, 128);
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedContiguous, KernelType::Kernel1D2D,
m, n, k, 1, major_a, major_b,
torch::kFloat8_e4m3fn, d.scalar_type(), false,
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(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);
// Launch
const SM100FP8Gemm1D2DRuntime::Args& args = {
.m = m, .n = n, .k = aligned_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.sfb = sfb.data_ptr(),
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM100FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_m_grouped_fp8_gemm_contiguous_1d2d", code);
SM100FP8Gemm1D2DRuntime::launch(runtime, args);
}
static void sm100_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d,
const torch::Tensor& masked_m,
const int& num_groups, const int& m, const int& n, const int& k,
const int& expected_m,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
const auto& aligned_k = align(k, 128);
const auto& config = get_best_config<SM100ArchSpec>(
GemmType::MGroupedMasked, KernelType::Kernel1D2D,
expected_m, n, k, num_groups, major_a, major_b,
torch::kFloat8_e4m3fn, d.scalar_type(), false,
device_runtime->get_num_sms());
const auto& tensor_map_a = make_tma_a_desc(major_a, a, m, k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(get_non_contiguous_dim(major_a))), num_groups,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(major_b, b, n, k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
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, num_groups, 0);
// Launch
const SM100FP8Gemm1D2DRuntime::Args& args = {
.m = m, .n = n, .k = aligned_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.epilogue_type = std::nullopt,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.sfb = sfb.data_ptr(),
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.tensor_map_sfa = tensor_map_sfa,
};
const auto& code = SM100FP8Gemm1D2DRuntime::generate(args);
const auto& runtime = compiler->build("sm100_fp8_m_grouped_gemm_masked_1d2d", code);
SM100FP8Gemm1D2DRuntime::launch(runtime, args);
}
} // namespace deep_gemm
csrc/jit_kernels/impls/sm90_bf16_gemm.hpp
+171 -14
View File
@@ -23,7 +23,7 @@ public:
void *grouped_layout;
CUtensorMap tensor_map_a;
CUtensorMap tensor_map_b;
CUtensorMap tensor_map_d;
CUtensorMap tensor_map_cd;
};
static std::string generate_impl(const Args& args) {
@@ -34,26 +34,31 @@ using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm90_bf16_gemm_impl<
{}, {},
{}, {}, {},
{},
{}, {}, {},
{}, {}, {},
{},
{}, {},
{}, {},
{},
{}, {},
{}, {}, {}
{}
>);
}};
)",
// TODO: add CD dtype
to_string(args.gemm_config.major_a), to_string(args.gemm_config.major_b),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages, args.gemm_config.num_last_stages,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
args.gemm_config.num_sms, to_string(args.gemm_config.gemm_type),
args.gemm_config.num_sms,
to_string(args.gemm_config.gemm_type), args.gemm_config.with_accumulation,
to_string(args.gemm_config.cd_dtype));
}
@@ -63,7 +68,7 @@ static void __instantiate_kernel() {{
args.grouped_layout,
args.m, args.n, args.k,
args.tensor_map_a, args.tensor_map_b,
args.tensor_map_d));
args.tensor_map_cd));
}
};
@@ -75,7 +80,6 @@ static void sm90_bf16_gemm(const torch::Tensor& a,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
DG_HOST_ASSERT(not c.has_value());
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
const auto& aligned_k = align(k, 64);
const auto& config = get_best_config<SM90ArchSpec>(
@@ -95,7 +99,7 @@ static void sm90_bf16_gemm(const torch::Tensor& a,
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
@@ -113,7 +117,7 @@ static void sm90_bf16_gemm(const torch::Tensor& a,
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_gemm", code);
@@ -128,7 +132,7 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
DG_HOST_ASSERT(k % 64 == 0);
const auto& config = get_best_config<SM90ArchSpec>(
@@ -148,7 +152,7 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
@@ -166,7 +170,7 @@ static void sm90_m_grouped_bf16_gemm_contiguous(const torch::Tensor& a,
.grouped_layout = m_indices.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_m_grouped_bf16_gemm_contiguous", code);
@@ -202,7 +206,7 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
config.block_k,
static_cast<int>(b.stride(get_non_contiguous_dim(major_b))), num_groups,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
@@ -220,11 +224,164 @@ static void sm90_bf16_m_grouped_gemm_masked(const torch::Tensor& a,
.grouped_layout = masked_m.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_d = tensor_map_d,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_m_grouped_gemm_masked", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
static void sm90_bf16_k_grouped_gemm(const torch::Tensor& a,
const torch::Tensor& b,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n,
const std::vector<int>& ks, const torch::Tensor& ks_tensor,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const std::string& compiled_dims) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::MN and major_b == cute::UMMA::Major::MN);
int sum_k = 0;
for (const auto& k: ks) {
sum_k += k;
DG_HOST_ASSERT(k % 128 == 0);
}
const auto& num_groups = static_cast<int>(ks.size());
// Get config using max K for better performance
const auto& max_k = *std::max_element(ks.begin(), ks.end());
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::KGroupedContiguous, KernelType::KernelNoSF,
m, n, max_k, num_groups, cute::UMMA::Major::MN, cute::UMMA::Major::MN,
torch::kBFloat16, d.scalar_type(), c.has_value(),
device_runtime->get_num_sms());
// Create tensor descriptors
const auto& tensor_map_a = make_tma_a_desc(cute::UMMA::Major::MN, a, m, sum_k,
SM100ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m),
config.block_k,
static_cast<int>(a.stride(0)), 1,
config.smem_config.swizzle_a_mode);
const auto& tensor_map_b = make_tma_b_desc(cute::UMMA::Major::MN, b, n, sum_k,
SM100ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n),
config.block_k,
static_cast<int>(b.stride(0)), 1,
config.smem_config.swizzle_b_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM100ArchSpec::get_cd_store_block_m(config.block_m),
SM100ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(1)), num_groups,
config.smem_config.swizzle_cd_mode);
// Launch kernel
const SM90BF16GemmRuntime::Args& args = {
.m = m, .n = n, .k = sum_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = ks_tensor.data_ptr(),
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_k_grouped_gemm", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
static void sm90_bf16_bhr_hdr_bhd(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_b,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, d, r, h, cute::UMMA::Major::K, cute::UMMA::Major::K,
torch::kBFloat16, tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, r, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
config.block_k, load_block_n, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, d, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = b, .n = d, .k = r,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_bhr_hdr_bhd", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
static void sm90_bf16_bhd_hdr_bhr(const torch::Tensor& tensor_a,
const torch::Tensor& tensor_b,
const torch::Tensor& tensor_d,
const int& b, const int& h, const int& r, const int& d,
const std::string& compiled_dims = "nk") {
const auto& config = get_best_config<SM90ArchSpec>(
GemmType::Batched, KernelType::KernelNoSF,
b, r, d, h, cute::UMMA::Major::K, cute::UMMA::Major::MN,
torch::kBFloat16, tensor_d.scalar_type(), false,
device_runtime->get_num_sms());
const int& load_block_m = SM90ArchSpec::get_ab_load_block_m(config.multicast_config, config.block_m);
const auto& tensor_map_a = make_tma_3d_desc(tensor_a, d, b, h,
config.block_k, load_block_m, 1,
tensor_a.stride(0), tensor_a.stride(1),
config.smem_config.swizzle_a_mode);
const int& load_block_n = SM90ArchSpec::get_ab_load_block_n(config.multicast_config, config.block_n);
const auto& tensor_map_b = make_tma_3d_desc(tensor_b, r, d, h,
load_block_n, config.block_k, 1,
tensor_b.stride(1), tensor_b.stride(0),
config.smem_config.swizzle_b_mode);
const int& store_block_m = SM90ArchSpec::get_cd_store_block_m(config.block_m);
const int& store_block_n = SM90ArchSpec::get_cd_store_block_n(config.block_n);
const auto& tensor_map_cd = make_tma_3d_desc(tensor_d, r, b, h,
store_block_n, store_block_m, 1,
tensor_d.stride(0), tensor_d.stride(1),
config.smem_config.swizzle_cd_mode);
// Launch
const SM90BF16GemmRuntime::Args& args = {
.m = b, .n = r, .k = d,
.num_groups = h,
.compiled_dims = compiled_dims,
.gemm_config = config,
.launch_args = LaunchArgs(config.num_sms, config.thread_config.num_threads,
config.smem_config.smem_size,
config.multicast_config.num_multicast),
.grouped_layout = nullptr,
.tensor_map_a = tensor_map_a,
.tensor_map_b = tensor_map_b,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90BF16GemmRuntime::generate(args);
const auto& runtime = compiler->build("sm90_bf16_bhd_hdr_bhr", code);
SM90BF16GemmRuntime::launch(runtime, args);
}
} // namespace deep_gemm
csrc/jit_kernels/impls/sm90_fp8_gemm_1d1d.hpp
+16 -14
View File
@@ -29,7 +29,7 @@ public:
CUtensorMap tensor_map_b_base;
CUtensorMap tensor_map_sfa;
CUtensorMap tensor_map_sfb;
CUtensorMap tensor_map_d;
CUtensorMap tensor_map_cd;
};
static std::string generate_impl(const Args& args) {
@@ -43,6 +43,7 @@ static void __instantiate_kernel() {{
{}, {}, {},
{},
{}, {}, {},
{}, {},
{},
{}, {},
{}, {},
@@ -54,6 +55,7 @@ static void __instantiate_kernel() {{
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode,
args.gemm_config.num_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
@@ -69,7 +71,7 @@ static void __instantiate_kernel() {{
args.m, args.n, args.k,
args.tensor_map_a_base, args.tensor_map_b_base,
args.tensor_map_sfa, args.tensor_map_sfb,
args.tensor_map_d));
args.tensor_map_cd));
}
};
@@ -105,11 +107,11 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, k,
config.block_n, config.block_k, 1, 0);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
0);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), 1,
0);
// Launch
const SM90FP8Gemm1D1DRuntime::Args& args = {
@@ -128,7 +130,7 @@ static void sm90_fp8_gemm_1d1d(const torch::Tensor& a, const torch::Tensor& sfa,
.tensor_map_b_base = tensor_map_b,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_d = tensor_map_d,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
@@ -180,11 +182,11 @@ static void sm90_fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Te
config.block_m, config.block_k, 1, 0);
const auto& tensor_map_sfb = make_tma_sf_desc(cute::UMMA::Major::MN, sfb, n, sum_sf_k * 128,
config.block_n, config.block_k, 1, 0);
const auto& tensor_map_d = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
const auto& tensor_map_cd = make_tma_cd_desc(d, m, n,
SM90ArchSpec::get_cd_store_block_m(config.block_m, true),
SM90ArchSpec::get_cd_store_block_n(config.block_n),
static_cast<int>(d.stride(-2)), num_groups,
config.smem_config.swizzle_cd_mode);
// Launch
const SM90FP8Gemm1D1DRuntime::Args& args = {
@@ -203,7 +205,7 @@ static void sm90_fp8_k_grouped_gemm_1d1d(const torch::Tensor& a, const torch::Te
.tensor_map_b_base = tensor_map_b_base,
.tensor_map_sfa = tensor_map_sfa,
.tensor_map_sfb = tensor_map_sfb,
.tensor_map_d = tensor_map_d,
.tensor_map_cd = tensor_map_cd,
};
const auto& code = SM90FP8Gemm1D1DRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_gemm_1d1d", code);
csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
+13 -6
View File
@@ -17,6 +17,7 @@ namespace deep_gemm {
class SM90FP8Gemm1D2DRuntime final: public LaunchRuntime<SM90FP8Gemm1D2DRuntime> {
public:
struct Args {
cute::UMMA::Major major_sfb;
int m, n, k, num_groups;
const std::string& compiled_dims;
const std::optional<std::string>& epilogue_type;
@@ -39,22 +40,25 @@ using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm90_fp8_gemm_1d2d_impl<
{}, {}, {},
{},
{}, {}, {},
{},
{}, {}, {},
{}, {}, {},
{}, {},
{}, {},
{}, {},
{}, {}, {}
{}, {},
{}
>);
}};
)",
// TODO: add CD dtype
to_string(args.major_sfb),
get_compiled_dim(args.m, 'm', args.compiled_dims), get_compiled_dim(args.n, 'n', args.compiled_dims), get_compiled_dim(args.k, 'k', args.compiled_dims),
args.num_groups,
args.gemm_config.block_m, args.gemm_config.block_n, args.gemm_config.block_k,
args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.smem_config.swizzle_a_mode, args.gemm_config.smem_config.swizzle_b_mode, args.gemm_config.smem_config.swizzle_cd_mode,
args.gemm_config.num_stages, args.gemm_config.num_last_stages,
args.gemm_config.thread_config.num_tma_threads, args.gemm_config.thread_config.num_math_threads,
args.gemm_config.multicast_config.num_multicast, args.gemm_config.multicast_config.is_multicast_on_a,
@@ -77,7 +81,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
const std::optional<torch::Tensor>& c,
const torch::Tensor& d,
const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
const std::string& compiled_dims,
const std::optional<std::string>& epilogue_type = std::nullopt) {
DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
@@ -113,6 +117,7 @@ static void sm90_fp8_gemm_1d2d(const torch::Tensor& a, const torch::Tensor& sfa,
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = aligned_k,
.num_groups = 1,
.compiled_dims = compiled_dims,
@@ -138,7 +143,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
const torch::Tensor& d,
const torch::Tensor& m_indices,
const int& num_groups, const int& m, const int& n, const int& k,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
const std::string& compiled_dims) {
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
@@ -173,6 +178,7 @@ static void sm90_m_grouped_fp8_gemm_contiguous_1d2d(const torch::Tensor& a, cons
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = aligned_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
@@ -199,7 +205,7 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
const torch::Tensor& masked_m,
const int& num_groups, const int& m, const int& n, const int& k,
const int& expected_m,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b,
const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
const std::string& compiled_dims) {
const auto& aligned_k = align(k, 128);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
@@ -234,6 +240,7 @@ static void sm90_m_grouped_fp8_gemm_masked_1d2d(const torch::Tensor& a, const to
// Launch
const SM90FP8Gemm1D2DRuntime::Args& args = {
.major_sfb = major_sfb,
.m = m, .n = n, .k = aligned_k,
.num_groups = num_groups,
.compiled_dims = compiled_dims,
csrc/jit_kernels/impls/smxx_clean_logits.hpp
+4 -4
View File
@@ -13,7 +13,7 @@ public:
int next_n;
int seq_len;
int seq_len_kv;
uint64_t stride_kv;
uint64_t stride_logits;
int* cu_seq_len_k_start;
int* cu_seq_len_k_end;
@@ -41,7 +41,7 @@ static void __instantiate_kernel() {{
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.seq_len, args.seq_len_kv, static_cast<int64_t>(args.stride_kv),
args.seq_len, args.seq_len_kv, static_cast<int64_t>(args.stride_logits),
args.cu_seq_len_k_start, args.cu_seq_len_k_end, args.logits
));
}
@@ -52,7 +52,7 @@ static void smxx_clean_logits(const torch::Tensor& logits,
const torch::Tensor& cu_seq_len_k_end,
const int& next_n,
const int& seq_len, const int& seq_len_kv,
const uint64_t &stride_kv) {
const uint64_t &stride_logits) {
const int block_kv = 8192;
const int num_warps = 8;
const int smem_size = block_kv * sizeof(float);
@@ -62,7 +62,7 @@ static void smxx_clean_logits(const torch::Tensor& logits,
.next_n = next_n,
.seq_len = seq_len,
.seq_len_kv = seq_len_kv,
.stride_kv = stride_kv,
.stride_logits = stride_logits,
.cu_seq_len_k_start = cu_seq_len_k_start.has_value() ? cu_seq_len_k_start.value().data_ptr<int>() : nullptr,
.cu_seq_len_k_end = cu_seq_len_k_end.data_ptr<int>(),
.logits = logits.data_ptr<float>(),
csrc/jit_kernels/impls/smxx_cublaslt.hpp
+9 -1
View File
@@ -3,6 +3,11 @@
#include <cublasLt.h>
#include <ATen/cuda/CUDAContext.h>
#include <ATen/cuda/CUDADataType.h>
#include <cute/arch/mma_sm100_umma.hpp>
#include "../../jit/device_runtime.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/compatibility.hpp"
namespace deep_gemm {
@@ -33,7 +38,6 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F_FAST_TF32;
cudaDataType_t scale_type = CUDA_R_32F;
const int& math_sms = device_runtime->get_num_sms();
bool fp8_fast_accumulate = false;
// Operation description
cublasLtMatmulDesc_t desc;
@@ -42,8 +46,12 @@ static void call_cublaslt_api(const cublasOperation_t& trans_a,
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_TRANSB, &trans_b, sizeof(trans_b)));
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SCALE_TYPE, &scale_type, sizeof(scale_type)));
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_SM_COUNT_TARGET, &math_sms, sizeof(math_sms)));
#if DG_FP8_COMPATIBLE
bool fp8_fast_accumulate = false;
if (a.scalar_type() == torch::kFloat8_e4m3fn)
DG_CUBLASLT_CHECK(cublasLtMatmulDescSetAttribute(desc, CUBLASLT_MATMUL_DESC_FAST_ACCUM, &fp8_fast_accumulate, sizeof(fp8_fast_accumulate)));
#endif
// Get cuBLASLt handle, workspace, and stream
const auto& handle = device_runtime->get_cublaslt_handle();
csrc/jit_kernels/impls/smxx_fp8_mqa_logits.hpp
+23 -11
View File
@@ -9,16 +9,18 @@
namespace deep_gemm {
class SM90FP8MQALogitsRuntime final: public LaunchRuntime<SM90FP8MQALogitsRuntime> {
class SMXXFP8MQALogitsRuntime final: public LaunchRuntime<SMXXFP8MQALogitsRuntime> {
public:
struct Args {
int seq_len;
int seq_len_kv;
int stride_kv;
int max_seqlen_k;
int stride_logits;
int num_heads, head_dim;
bool is_compressed_logits;
int num_q_stages;
int num_kv_stages;
int block_q;
int block_kv;
@@ -52,6 +54,7 @@ using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_mqa_logits<
{}, {},
{},
{}, {},
{}, {},
{}, {}
@@ -59,6 +62,7 @@ static void __instantiate_kernel() {{
}};
)", arch, arch,
args.num_heads, args.head_dim,
args.is_compressed_logits,
args.block_q, args.block_kv,
args.num_q_stages, args.num_kv_stages,
args.num_specialized_threads, args.num_math_threads);
@@ -66,7 +70,8 @@ static void __instantiate_kernel() {{
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.seq_len, args.seq_len_kv, static_cast<int64_t>(args.stride_kv),
args.seq_len, args.seq_len_kv,
args.max_seqlen_k, static_cast<int64_t>(args.stride_logits),
args.cu_seq_len_k_start, args.cu_seq_len_k_end,
args.logits,
args.tensor_map_q, args.tensor_map_kv,
@@ -81,18 +86,22 @@ static void smxx_fp8_mqa_logits(const torch::Tensor& q,
const torch::Tensor& cu_seq_len_k_start,
const torch::Tensor& cu_seq_len_k_end,
const torch::Tensor& logits,
const int& seq_len, const int& seq_len_kv, const int& stride_kv,
const int& seq_len, const int& seq_len_kv,
const int& max_seqlen_k, const int& stride_logits,
const int& num_heads, const int& head_dim,
const int& seq_len_alignment) {
constexpr int block_qh = 128;
constexpr int block_kv = 256;
constexpr int num_specialized_threads = 128;
const int num_math_threads = (device_runtime->get_arch_major() == 10 ? 256 : 512);
constexpr int num_q_stages = 3, num_kv_stages = 3;
const int num_math_threads = (device_runtime->get_arch_major() == 10 ? 256 : 512);
const int block_q = block_qh / num_heads;
DG_HOST_ASSERT(block_qh % num_heads == 0);
DG_HOST_ASSERT(seq_len_alignment % block_q == 0);
// Use compressed logits format when max_seqlen_k is specified
const bool is_compressed_logits = (max_seqlen_k > 0);
// Construct TMAs
DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
const auto& tensor_map_q = make_tma_2d_desc(q, head_dim, seq_len * num_heads,
@@ -120,13 +129,16 @@ static void smxx_fp8_mqa_logits(const torch::Tensor& q,
smem_size += (num_q_stages * 2 + num_kv_stages * 2 + (num_math_threads / 128) * 2) * 8;
smem_size += 4;
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
// Launch
const SM90FP8MQALogitsRuntime::Args& args = {
const SMXXFP8MQALogitsRuntime::Args& args = {
.seq_len = seq_len,
.seq_len_kv = seq_len_kv,
.stride_kv = stride_kv,
.max_seqlen_k = max_seqlen_k,
.stride_logits = stride_logits,
.num_heads = num_heads, .head_dim = head_dim,
.is_compressed_logits = is_compressed_logits,
.num_q_stages = num_q_stages,
.num_kv_stages = num_kv_stages,
.block_q = block_q,
@@ -144,9 +156,9 @@ static void smxx_fp8_mqa_logits(const torch::Tensor& q,
num_specialized_threads + num_math_threads,
smem_size)
};
const auto& code = SM90FP8MQALogitsRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_mqa_logits", code);
SM90FP8MQALogitsRuntime::launch(runtime, args);
const auto& code = SMXXFP8MQALogitsRuntime::generate(args);
const auto& runtime = compiler->build("smxx_fp8_mqa_logits", code);
SMXXFP8MQALogitsRuntime::launch(runtime, args);
}
} // namespace deep_gemm
csrc/jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp
+17 -4
View File
@@ -14,8 +14,10 @@ public:
int aligned_batch_size;
int split_kv;
int num_sms;
int batch_size;
int next_n;
bool is_context_lens_2d;
int* context_lens;
int* schedule_metadata;
@@ -41,6 +43,8 @@ static void __instantiate_kernel() {{
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.batch_size,
args.next_n,
args.is_context_lens_2d,
args.context_lens,
args.schedule_metadata
));
@@ -49,12 +53,14 @@ static void __instantiate_kernel() {{
static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens,
const torch::Tensor& schedule_metadata,
const int& batch_size, const int& block_kv, const int& num_sms) {
const int& batch_size, const int& next_n,
const int& block_kv, const int& num_sms,
const bool& is_context_lens_2d) {
constexpr int num_math_warpgroups = 4;
constexpr int num_threads = 32;
const int aligned_batch_size = align(batch_size, 32);
const int split_kv = block_kv * num_math_warpgroups;
// Calculate shared memory size
const int smem_size = aligned_batch_size * static_cast<int>(sizeof(int));
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
@@ -66,6 +72,8 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens,
.split_kv = split_kv,
.num_sms = num_sms,
.batch_size = batch_size,
.next_n = next_n,
.is_context_lens_2d = is_context_lens_2d,
.context_lens = context_lens.data_ptr<int>(),
.schedule_metadata = schedule_metadata.data_ptr<int>(),
.launch_args = LaunchArgs(1, num_threads, smem_size)
@@ -83,6 +91,7 @@ public:
int num_heads;
int head_dim;
int block_kv;
bool is_context_lens_2d;
int block_table_stride;
int logits_stride;
@@ -121,6 +130,7 @@ static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{},
{}, {}
@@ -129,6 +139,7 @@ static void __instantiate_kernel() {{
)", arch, arch,
args.next_n, args.num_heads,
args.head_dim, args.block_kv,
args.is_context_lens_2d,
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads);
@@ -158,6 +169,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
const int& batch_size, const int& next_n,
const int& num_heads, const int& head_dim,
const int& num_kv_blocks, const int& block_kv,
const bool& is_context_lens_2d,
const int& kv_cache_stride_bytes,
const int& logits_stride,
const int& block_table_stride,
@@ -209,6 +221,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
.num_heads = num_heads,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
@@ -229,7 +242,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
smem_size)
};
const auto& code = SMXXFP8PagedMQALogitsRuntime::generate(args);
const auto& runtime = compiler->build("sm90_fp8_paged_mqa_logits", code);
const auto& runtime = compiler->build("smxx_fp8_paged_mqa_logits", code);
SMXXFP8PagedMQALogitsRuntime::launch(runtime, args);
}
csrc/jit_kernels/impls/smxx_layout.hpp
+1
View File
@@ -3,6 +3,7 @@
#include <torch/python.h>
#include "../../jit/kernel_runtime.hpp"
#include "../../jit/compiler.hpp"
#include "../../utils/exception.hpp"
#include "../../utils/format.hpp"
#include "../../utils/math.hpp"
csrc/python_api.cpp
+1 -1
View File
@@ -8,7 +8,7 @@
#include "apis/runtime.hpp"
#ifndef TORCH_EXTENSION_NAME
#define TORCH_EXTENSION_NAME deep_gemm_cpp
#define TORCH_EXTENSION_NAME _C
#endif
// ReSharper disable once CppParameterMayBeConstPtrOrRef
csrc/utils/compatibility.hpp
+10
View File
@@ -0,0 +1,10 @@
#pragma once
#include <torch/version.h>
#include <cuda.h>
// `torch::kFloat8_e4m3fn` is supported since PyTorch 2.1
#define DG_FP8_COMPATIBLE (TORCH_VERSION_MAJOR > 2 or (TORCH_VERSION_MAJOR == 2 and TORCH_VERSION_MINOR >= 1))
// `cuTensorMapEncodeTiled` is supported since CUDA Driver API 12.1
#define DG_TENSORMAP_COMPATIBLE (CUDA_VERSION >= 12010)
csrc/utils/exception.hpp
+1 -1
View File
@@ -54,7 +54,7 @@ do { \
if (e != CUDA_SUCCESS) { \
std::stringstream ss; \
const char *name, *info; \
cuGetErrorName(e, &name), cuGetErrorString(e, &info); \
lazy_cuGetErrorName(e, &name), lazy_cuGetErrorString(e, &info); \
ss << static_cast<int>(e) << " (" << name << ", " << info << ")"; \
throw DGException("CUDA driver", __FILE__, __LINE__, ss.str()); \
} \
csrc/utils/hash.hpp
+6 -2
View File
@@ -4,7 +4,7 @@
namespace deep_gemm {
static uint64_t fnv1a(const std::string& data, const uint64_t& seed) {
static uint64_t fnv1a(const std::vector<char>& data, const uint64_t& seed) {
uint64_t h = seed;
const uint64_t& prime = 0x100000001b3ull;
for (const char& c: data) {
@@ -14,7 +14,7 @@ static uint64_t fnv1a(const std::string& data, const uint64_t& seed) {
return h;
}
static std::string get_hex_digest(const std::string& data) {
static std::string get_hex_digest(const std::vector<char>& data) {
const auto& state_0 = fnv1a(data, 0xc6a4a7935bd1e995ull);
const auto& state_1 = fnv1a(data, 0x9e3779b97f4a7c15ull);
@@ -32,4 +32,8 @@ static std::string get_hex_digest(const std::string& data) {
return oss.str();
}
static std::string get_hex_digest(const std::string& data) {
return get_hex_digest(std::vector<char>{data.begin(), data.end()});
}
} // namespace deep_gemm
csrc/utils/layout.hpp
+11 -6
View File
@@ -51,7 +51,7 @@ get_default_recipe(const torch::ScalarType& sfa_dtype, const torch::ScalarType&
} else if (arch_major == 10) {
DG_HOST_ASSERT(sfb_dtype == torch::kFloat or sfb_dtype == torch::kInt);
return sfb_dtype == torch::kFloat ?
std::make_tuple(1, 128, 128): // Legacy format or 1D2D kernels
std::make_tuple(1, 128, 128): // Legacy format
std::make_tuple(1, 1, 128); // 1D1D kernels
}
DG_HOST_UNREACHABLE("Unknown recipe");
@@ -63,7 +63,7 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf,
const int& gran_mn, const int& gran_k,
const std::optional<int>& num_groups,
const bool& tma_stride_check = false,
const bool& contiguous_check = false,
const bool& sm90_sfb_check = false,
const std::optional<torch::ScalarType>& type_check = std::nullopt) {
// Type check
if (type_check.has_value())
@@ -82,13 +82,18 @@ static torch::Tensor check_sf_layout(const torch::Tensor& sf,
if (tma_stride_check) {
if (num_groups.has_value())
DG_HOST_ASSERT(sf.stride(-3) == sf.stride(-1) * sf.size(-1));
DG_HOST_ASSERT(sf.stride(-2) == 1);
// Check contiguity in the MN direction
DG_HOST_ASSERT(sf.stride(-2) == 1 or mn == 1);
DG_HOST_ASSERT(sf.stride(-1) == get_tma_aligned_size(mn, sf.element_size()));
}
// Hopper SFB must be contiguous
if (contiguous_check)
DG_HOST_ASSERT(sf.is_contiguous());
// SM90 SFB must be contiguous, or contiguous after transposing the last two dimensions
if (sm90_sfb_check) {
if (num_groups.has_value())
DG_HOST_ASSERT(sf.stride(-3) == sf.size(-2) * sf.size(-1));
DG_HOST_ASSERT((sf.stride(-1) == 1 and sf.stride(-2) == sf.size(-1)) or
(sf.stride(-1) == sf.size(-2) and sf.stride(-2) == 1));
}
return sf;
}