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

[Public release 26/04] Introducing Mega MoE, FP4 Indexer and other features/fixes (#304)

Browse Source 
* Merge with private repo

* Update README

* Update README

* Update README

* Add PyTorch requirements

* Fix sync scopes for MQA logits (#256)

* Update README
This commit is contained in:
Chenggang Zhao
2026-04-17 09:45:14 +08:00
committed by GitHub
parent d30fc36c8f
commit 7f2a703ed5
109 changed files with 12101 additions and 3219 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/apis/attention.hpp
+276 -125
View File
@@ -5,9 +5,9 @@
#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/smxx_fp8_mqa_logits.hpp"
#include "../jit_kernels/impls/smxx_fp8_paged_mqa_logits.hpp"
#include "../jit_kernels/impls/sm100_fp8_fp4_gemm_1d1d.hpp"
#include "../jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp"
#include "../jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp"
#include "../jit_kernels/impls/smxx_clean_logits.hpp"
#endif
@@ -24,8 +24,8 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
const std::string& compiled_dims,
const bool& disable_ue8m0_cast) {
// Shape must be `[M, K] @ [N, K].T`
const auto& major_a = get_major_type_ab(a.first);
const auto& major_b = get_major_type_ab(b.first);
const auto major_a = get_major_type_ab(a.first);
const auto major_b = get_major_type_ab(b.first);
if (fp8_requires_k_major()) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
DG_HOST_ASSERT(major_b == cute::UMMA::Major::K);
@@ -35,9 +35,9 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
check_major_type_cd(d);
// Type and shape checks
const auto& [m , k ] = get_shape<2>(a.first);
const auto& [n , k_] = get_shape<2>(b.first);
const auto& [m_, n_] = get_shape<2>(d);
const auto [m , k ] = get_shape<2>(a.first);
const auto [n , k_] = get_shape<2>(b.first);
const auto [m_, n_] = get_shape<2>(d);
DG_HOST_ASSERT(m == m_ and k == k_);
DG_HOST_ASSERT(n > 0 and k > 0);
DG_HOST_ASSERT(a.first.scalar_type() == torch::kFloat8_e4m3fn);
@@ -45,7 +45,7 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat);
// Check head splits and N
const auto& [left, mid, right] = head_splits;
const auto [left, mid, right] = head_splits;
DG_HOST_ASSERT(n % (left + right) == 0 and n_ == n + n / (left + right) * mid);
// Do nothing if the problem is empty
@@ -53,16 +53,16 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
return;
// Transform SFA and SFB into compute-required layout
const auto& [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout(
const auto [sfa, sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout(
a.second, b.second, m, n, k, recipe, std::nullopt, std::nullopt,
std::nullopt, std::nullopt, disable_ue8m0_cast);
DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128);
// Dispatch into different implements
const auto& arch_major = device_runtime->get_arch_major();
const auto& epilogue_type = fmt::format("EpilogueHeadSplits<{}, {}, {}>", left, mid, right);
const auto arch_major = device_runtime->get_arch_major();
const auto epilogue_type = fmt::format("epilogue::transform::EpilogueHeadSplits<{}, {}, {}>", left, mid, right);
if (arch_major == 9 and sfa.scalar_type() == torch::kFloat and std::get<1>(recipe.value()) != 1) {
const auto& major_sfb = get_major_type_ab(sfb);
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) {
// NOTES: Only granularity 128 and FP8 are exposed in the API
@@ -73,59 +73,113 @@ static void fp8_gemm_nt_skip_head_mid(const std::pair<torch::Tensor, torch::Tens
}
}
static torch::Tensor fp8_mqa_logits(const torch::Tensor& q,
const std::pair<torch::Tensor, torch::Tensor>& kv,
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 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);
const auto& [seq_len_kv_] = get_shape<1>(kv.second);
static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::optional<torch::Tensor>>& q,
const std::tuple<torch::Tensor, torch::Tensor>& kv,
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 int& max_seqlen_k,
const at::ScalarType& logits_dtype) {
const auto [q_fp, q_sf] = q;
const auto [kv_fp, kv_sf] = kv;
const bool is_fp4 = q_sf.has_value();
int seq_len, seq_len_kv, num_heads, head_dim;
DG_HOST_ASSERT(seq_len == seq_len_);
DG_HOST_ASSERT(num_heads == num_heads_ and head_dim == head_dim_);
DG_HOST_ASSERT(seq_len_kv == seq_len_kv_);
DG_HOST_ASSERT(cu_seq_len_k_start.size(0) == seq_len);
DG_HOST_ASSERT(cu_seq_len_k_end.size(0) == seq_len);
if (is_fp4) {
// Check FP4 Q
std::tie(seq_len, num_heads, head_dim) = get_shape<3>(q_fp);
head_dim *= 2;
DG_HOST_ASSERT(num_heads == 32 or num_heads == 64);
DG_HOST_ASSERT(head_dim == 128);
DG_HOST_ASSERT(q_fp.is_contiguous());
DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4);
DG_HOST_ASSERT(q.is_contiguous() and kv.first.is_contiguous());
DG_HOST_ASSERT(kv.second.is_contiguous());
DG_HOST_ASSERT(weights.is_contiguous());
DG_HOST_ASSERT(cu_seq_len_k_start.is_contiguous());
DG_HOST_ASSERT(cu_seq_len_k_end.is_contiguous());
// Check SF Q
auto [_seq_len, _num_heads] = get_shape<2>(q_sf.value());
DG_HOST_ASSERT(seq_len == _seq_len and num_heads == _num_heads);
DG_HOST_ASSERT(q_sf.value().is_contiguous());
DG_HOST_ASSERT(q_sf.value().scalar_type() == torch::kInt32);
DG_HOST_ASSERT(q.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(kv.first.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(kv.second.scalar_type() == torch::kFloat);
// Check FP4 KV
int _head_dim;
std::tie(seq_len_kv, _head_dim) = get_shape<2>(kv_fp);
_head_dim *= 2;
DG_HOST_ASSERT(head_dim == _head_dim);
DG_HOST_ASSERT(kv_fp.is_contiguous());
DG_HOST_ASSERT(kv_fp.scalar_type() == kPackedFP4);
// Check SF KV
auto [_seq_len_kv] = get_shape<1>(kv_sf);
DG_HOST_ASSERT(seq_len_kv == _seq_len_kv);
DG_HOST_ASSERT(kv_sf.is_contiguous());
DG_HOST_ASSERT(kv_sf.scalar_type() == torch::kInt32);
} else {
// Check FP8 Q
std::tie(seq_len, num_heads, head_dim) = get_shape<3>(q_fp);
DG_HOST_ASSERT(num_heads == 32 or num_heads == 64);
DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
DG_HOST_ASSERT(q_fp.is_contiguous());
DG_HOST_ASSERT(q_fp.scalar_type() == torch::kFloat8_e4m3fn);
// Check FP4 KV
int _head_dim;
std::tie(seq_len_kv, _head_dim) = get_shape<2>(kv_fp);
DG_HOST_ASSERT(head_dim == _head_dim);
DG_HOST_ASSERT(kv_fp.is_contiguous());
DG_HOST_ASSERT(kv_fp.scalar_type() == torch::kFloat8_e4m3fn);
// Check SF KV
auto [_seq_len_kv] = get_shape<1>(kv_sf);
DG_HOST_ASSERT(seq_len_kv == _seq_len_kv);
DG_HOST_ASSERT(kv_sf.is_contiguous());
DG_HOST_ASSERT(kv_sf.scalar_type() == torch::kFloat);
}
// Check weights
auto [_seq_len, _num_heads] = get_shape<2>(weights);
DG_HOST_ASSERT(seq_len == _seq_len and num_heads == _num_heads);
DG_HOST_ASSERT(weights.stride(1) == 1);
DG_HOST_ASSERT(weights.scalar_type() == torch::kFloat);
// Check cu_seq_len_k_start
DG_HOST_ASSERT(cu_seq_len_k_start.size(0) == seq_len);
DG_HOST_ASSERT(cu_seq_len_k_start.is_contiguous());
DG_HOST_ASSERT(cu_seq_len_k_start.scalar_type() == torch::kInt);
// Check cu_seq_len_k_end
DG_HOST_ASSERT(cu_seq_len_k_end.size(0) == seq_len);
DG_HOST_ASSERT(cu_seq_len_k_end.is_contiguous());
DG_HOST_ASSERT(cu_seq_len_k_end.scalar_type() == torch::kInt);
constexpr int seq_len_alignment = 4;
// Allocate output
constexpr int block_qh = 128;
constexpr int block_kv = 256;
const auto aligned_seq_len = align(seq_len, seq_len_alignment);
const int block_q = block_qh / num_heads;
DG_HOST_ASSERT(block_qh % num_heads == 0);
torch::Tensor logits;
int stride_logits;
int aligned_seq_len = align(seq_len, block_q), 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 stride must be 16-byte aligned
stride_logits = align(seq_len_kv + block_kv, 8);
logits = torch::empty({aligned_seq_len, stride_logits}, q_fp.options().dtype(logits_dtype));
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 = torch::empty({aligned_seq_len, stride_logits}, q_fp.options().dtype(logits_dtype));
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, max_seqlen_k, stride_logits, num_heads, head_dim, seq_len_alignment);
const auto arch_major = device_runtime->get_arch_major();
if (is_fp4 and arch_major == 10) {
sm100_fp4_mqa_logits(q_fp, q_sf.value(), kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype,
seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv);
} else if (not is_fp4 and (arch_major == 9 or arch_major == 10)) {
smxx_fp8_mqa_logits(q_fp, kv_fp, kv_sf, weights, cu_seq_len_k_start, cu_seq_len_k_end, logits, logits_dtype,
seq_len, seq_len_kv, max_seqlen_k, stride_logits, num_heads, head_dim, block_q, block_kv);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
@@ -137,23 +191,21 @@ static torch::Tensor fp8_mqa_logits(const torch::Tensor& q,
}
static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_lens, int block_kv, int num_sms) {
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);
}
// NOTES: Only 2D context lens is supported for now
DG_HOST_ASSERT(context_lens.dim() == 2);
const bool is_context_lens_2d = true;
const int batch_size = context_lens.size(0);
const int next_n = context_lens.size(1);
DG_HOST_ASSERT(context_lens.scalar_type() == torch::kInt);
DG_HOST_ASSERT(context_lens.is_contiguous());
// Create metadata tensor
auto schedule_metadata = torch::empty({num_sms + 1, 2}, context_lens.options());
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9 or arch_major == 10) {
DG_HOST_ASSERT(block_kv == 64 or (arch_major == 10 and block_kv == 32));
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");
@@ -162,85 +214,145 @@ static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_
return schedule_metadata;
}
static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& fused_kv_cache,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const int& max_context_len,
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_next_n, num_heads_] = get_shape<2>(weights);
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);
static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, std::optional<torch::Tensor>>& q,
const torch::Tensor& fused_kv_cache,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const int& max_context_len,
const bool& clean_logits,
const at::ScalarType& logits_dtype) {
const auto [q_fp, q_sf] = q;
const bool is_fp4 = q_sf.has_value();
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_);
torch::Tensor kv_cache, kv_cache_sf;
int batch_size, next_n, num_heads, head_dim;
int num_kv_blocks, block_kv;
int kv_cache_stride_bytes;
int block_table_stride = block_table.stride(0);
int num_sms = device_runtime->get_num_sms();
if (is_fp4) {
// Check FP4 Q
std::tie(batch_size, next_n, num_heads, head_dim) = get_shape<4>(q_fp);
head_dim *= 2;
DG_HOST_ASSERT(next_n >= 1);
DG_HOST_ASSERT(num_heads == 32 or num_heads == 64);
DG_HOST_ASSERT(head_dim == 128);
DG_HOST_ASSERT(q_fp.is_contiguous());
DG_HOST_ASSERT(q_fp.scalar_type() == kPackedFP4);
// Check SF Q
auto [_batch_size, _next_n, _num_heads] = get_shape<3>(q_sf.value());
DG_HOST_ASSERT(batch_size == _batch_size and next_n == _next_n and num_heads == _num_heads);
DG_HOST_ASSERT(q_sf.value().is_contiguous());
DG_HOST_ASSERT(q_sf.value().scalar_type() == torch::kInt32);
// Check fused KV cache
int num_heads_kv, fp4_with_sf_bytes;
std::tie(num_kv_blocks, block_kv, num_heads_kv, fp4_with_sf_bytes) = get_shape<4>(fused_kv_cache);
DG_HOST_ASSERT(block_kv == 32 or block_kv == 64);
DG_HOST_ASSERT(num_heads_kv == 1 and fp4_with_sf_bytes == head_dim / 2 + static_cast<int>(sizeof(int)));
DG_HOST_ASSERT(fused_kv_cache.stride(1) == fp4_with_sf_bytes and fused_kv_cache.stride(3) == 1);
DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kByte);
// Derive FP4 values and SF tensor
kv_cache_stride_bytes = fused_kv_cache.stride(0);
DG_HOST_ASSERT(kv_cache_stride_bytes % sizeof(int) == 0);
kv_cache = torch::from_blob(
fused_kv_cache.data_ptr(),
{num_kv_blocks, block_kv, head_dim / 2},
{kv_cache_stride_bytes, head_dim / 2, 1},
torch::TensorOptions().dtype(kPackedFP4)
);
kv_cache_sf = torch::from_blob(
fused_kv_cache.data_ptr<uint8_t>() + block_kv * head_dim / 2,
{num_kv_blocks, block_kv},
{kv_cache_stride_bytes / static_cast<int>(sizeof(int)), 1},
torch::TensorOptions().dtype(torch::kInt32)
);
} else {
DG_HOST_ASSERT(context_lens.dim() == 1);
const auto& [batch_size__] = get_shape<1>(context_lens);
DG_HOST_ASSERT(batch_size == batch_size__);
// Check FP8 Q
std::tie(batch_size, next_n, num_heads, head_dim) = get_shape<4>(q_fp);
DG_HOST_ASSERT(next_n >= 1);
DG_HOST_ASSERT(num_heads == 32 or num_heads == 64);
DG_HOST_ASSERT(head_dim == 32 or head_dim == 64 or head_dim == 128);
DG_HOST_ASSERT(q_fp.is_contiguous());
DG_HOST_ASSERT(q_fp.scalar_type() == torch::kFloat8_e4m3fn);
// Check fused KV cache
int num_heads_kv, head_dim_with_sf;
std::tie(num_kv_blocks, block_kv, num_heads_kv, head_dim_with_sf) = get_shape<4>(fused_kv_cache);
DG_HOST_ASSERT(block_kv == 32 or block_kv == 64);
DG_HOST_ASSERT(num_heads_kv == 1 and head_dim_with_sf == head_dim + static_cast<int>(sizeof(float)));
DG_HOST_ASSERT(fused_kv_cache.stride(1) == head_dim_with_sf and fused_kv_cache.stride(3) == 1);
DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kByte);
// Derive FP8 values and SF tensor
kv_cache_stride_bytes = fused_kv_cache.stride(0);
DG_HOST_ASSERT(kv_cache_stride_bytes % sizeof(float) == 0);
kv_cache = torch::from_blob(
fused_kv_cache.data_ptr(),
{num_kv_blocks, block_kv, head_dim},
{kv_cache_stride_bytes, head_dim, 1},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn)
);
kv_cache_sf = torch::from_blob(
fused_kv_cache.data_ptr<uint8_t>() + block_kv * head_dim,
{num_kv_blocks, block_kv},
{kv_cache_stride_bytes / static_cast<int>(sizeof(float)), 1},
torch::TensorOptions().dtype(torch::kFloat32)
);
// Weights must be contiguous for FP8
DG_HOST_ASSERT(weights.is_contiguous());
}
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)));
DG_HOST_ASSERT(schedule_meta_size == num_sms + 1 and meta_info_size == 2);
DG_HOST_ASSERT(next_n == 1 or next_n == 2);
DG_HOST_ASSERT(block_kv == 64);
DG_HOST_ASSERT(q.is_contiguous());
DG_HOST_ASSERT(kv_cache_stride_bytes % sizeof(float) == 0);
DG_HOST_ASSERT(fused_kv_cache.stride(1) == head_dim_with_sf);
DG_HOST_ASSERT(fused_kv_cache.stride(2) == head_dim_with_sf);
DG_HOST_ASSERT(fused_kv_cache.stride(3) == 1);
DG_HOST_ASSERT(weights.is_contiguous());
DG_HOST_ASSERT(context_lens.is_contiguous());
DG_HOST_ASSERT(block_table.stride(1) == 1);
DG_HOST_ASSERT(schedule_meta.is_contiguous());
DG_HOST_ASSERT(q.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(fused_kv_cache.scalar_type() == torch::kByte);
// Check weights
auto [_batch_size_next_n, _num_heads] = get_shape<2>(weights);
DG_HOST_ASSERT(_batch_size_next_n == batch_size * next_n and _num_heads == num_heads);
DG_HOST_ASSERT(weights.stride(1) == 1);
DG_HOST_ASSERT(weights.scalar_type() == torch::kFloat);
DG_HOST_ASSERT(context_lens.scalar_type() == torch::kInt);
// Check block table
auto [_batch_size, _max_block_len] = get_shape<2>(block_table);
DG_HOST_ASSERT(_batch_size == batch_size);
DG_HOST_ASSERT(block_table.stride(1) == 1);
DG_HOST_ASSERT(block_table.scalar_type() == torch::kInt);
// Check schedule metadata
auto [_schedule_meta_size, _meta_info_size] = get_shape<2>(schedule_meta);
DG_HOST_ASSERT(_schedule_meta_size == num_sms + 1 and _meta_info_size == 2);
DG_HOST_ASSERT(schedule_meta.is_contiguous());
DG_HOST_ASSERT(schedule_meta.scalar_type() == torch::kInt);
// Derive FP8 values and SF tensor from KV cache
const auto& kv_cache = torch::from_blob(
fused_kv_cache.data_ptr(),
{num_kv_blocks, block_kv, head_dim},
{kv_cache_stride_bytes, head_dim, 1},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn)
);
const auto& kv_cache_scales = torch::from_blob(
fused_kv_cache.data_ptr<uint8_t>() + block_kv * head_dim,
{num_kv_blocks, block_kv},
{kv_cache_stride_bytes / static_cast<int>(sizeof(float)), 1},
torch::TensorOptions().dtype(torch::kFloat32)
);
// Check context lengths
// NOTES: Only 2D context lens is supported for now
DG_HOST_ASSERT(context_lens.dim() == 2);
const bool is_context_lens_2d = true;
const auto [__batch_size, _next_n] = get_shape<2>(context_lens);
DG_HOST_ASSERT(batch_size == __batch_size and next_n == _next_n);
DG_HOST_ASSERT(context_lens.is_contiguous());
DG_HOST_ASSERT(context_lens.scalar_type() == torch::kInt);
// Allocate output
constexpr int split_kv = 256;
const auto& aligned_max_context_len = align(max_context_len, split_kv);
auto logits = torch::empty({batch_size * next_n, aligned_max_context_len}, q.options().dtype(torch::kFloat));
const auto aligned_max_context_len = align(max_context_len, split_kv);
auto logits = torch::empty({batch_size * next_n, aligned_max_context_len}, q_fp.options().dtype(logits_dtype));
logits = logits.slice(-1, 0, max_context_len);
DG_HOST_ASSERT(logits_dtype == torch::kFloat32 or logits_dtype == torch::kBFloat16);
// Dispatch implementation
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, is_context_lens_2d,
kv_cache_stride_bytes, aligned_max_context_len, block_table_stride, num_sms, split_kv);
const auto arch_major = device_runtime->get_arch_major();
if (is_fp4 and arch_major == 10) {
sm100_fp4_paged_mqa_logits(q_fp, q_sf.value(), kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, schedule_meta,
logits_dtype, batch_size, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,
aligned_max_context_len, block_table_stride, num_sms, split_kv);
} else if (not is_fp4 and (arch_major == 9 or arch_major == 10)) {
smxx_fp8_paged_mqa_logits(q_fp, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, schedule_meta,
logits_dtype, batch_size, next_n, num_heads, head_dim, num_kv_blocks, block_kv, is_context_lens_2d,
aligned_max_context_len, block_table_stride, num_sms, split_kv);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
@@ -253,6 +365,32 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
return logits;
}
// Legacy API wrappers
static torch::Tensor fp8_mqa_logits(const torch::Tensor& q,
const std::tuple<torch::Tensor, torch::Tensor>& kv,
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 int& max_seqlen_k) {
return fp8_fp4_mqa_logits(std::make_tuple(q, std::nullopt), kv, weights,
cu_seq_len_k_start, cu_seq_len_k_end,
clean_logits, max_seqlen_k, torch::kFloat);
}
static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& fused_kv_cache,
const torch::Tensor& weights,
const torch::Tensor& context_lens,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const int& max_context_len,
const bool& clean_logits) {
return fp8_fp4_paged_mqa_logits(std::make_tuple(q, std::nullopt), fused_kv_cache, weights,
context_lens, block_table, schedule_meta,
max_context_len, clean_logits, torch::kFloat);
}
#endif
static void register_apis(pybind11::module_& m) {
@@ -262,13 +400,26 @@ static void register_apis(pybind11::module_& m) {
py::arg("recipe") = std::nullopt,
py::arg("compiled_dims") = "nk",
py::arg("disable_ue8m0_cast") = false);
m.def("fp8_fp4_mqa_logits", &fp8_fp4_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("max_seqlen_k") = 0,
py::arg("logits_dtype") = torch::kFloat32);
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_fp4_paged_mqa_logits", &fp8_fp4_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,
py::arg("logits_dtype") = torch::kFloat32);
// Legacy API
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("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"),
csrc/apis/einsum.hpp
+35 -34
View File
@@ -29,7 +29,7 @@ static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const tor
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(not c.has_value());
const auto& workspace = torch::empty_like(d, d.options().dtype(torch::kFloat32));
const auto workspace = torch::empty_like(d, d.options().dtype(torch::kFloat32));
DG_CUDA_RUNTIME_CHECK(cudaMemsetAsync(workspace.data_ptr(), 0, workspace.nbytes(),
c10::cuda::getCurrentCUDAStream()));
bmk_bnk_mn(a, b, workspace, workspace);
@@ -43,12 +43,12 @@ static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const tor
DG_HOST_ASSERT(b.is_contiguous());
DG_HOST_ASSERT(d.is_contiguous());
const auto& [s , m, k ] = get_shape<3>(a);
const auto& [s_, n, k_] = get_shape<3>(b);
const auto [s , m, k ] = get_shape<3>(a);
const auto [s_, n, k_] = get_shape<3>(b);
DG_HOST_ASSERT(s == s_ and k == k_);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
sm90_bmn_bnk_mn_gemm(a, b, d, s, m, n, k);
} else if (arch_major == 10) {
@@ -59,9 +59,9 @@ 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, 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);
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);
DG_HOST_ASSERT(b == b_ and h == h_ and r == r_ and d == d_ and h == h__);
DG_HOST_ASSERT(A.scalar_type() == torch::kBFloat16 and A.stride(2) == 1);
@@ -69,7 +69,7 @@ static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const to
DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
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) {
@@ -82,9 +82,9 @@ static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const to
}
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);
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);
DG_HOST_ASSERT(b == b_ and h == h_ and r == r_ and d == d_ and h == h__);
DG_HOST_ASSERT(A.scalar_type() == torch::kBFloat16 and A.stride(2) == 1);
@@ -92,7 +92,7 @@ static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const to
DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
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) {
@@ -142,16 +142,16 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
std::optional<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;
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);
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);
@@ -163,15 +163,16 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
return;
// Transform scaling factors
const auto& [transformed_sfa, transformed_sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout(
const auto [transformed_sfa, transformed_sfb, gran_k_a, gran_k_b] = layout::transform_sf_pair_into_required_layout(
sfa, sfb, m, n, k, recipe, std::nullopt, std::nullopt, batch_size, batch_size, 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);
sm100_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d, batch_size, m, n, k, gran_k_a, gran_k_b, major_a, major_b, compiled_dims);
} else {
const auto& major_sfb = get_major_type_ab(sfb);
const auto major_sfb = get_major_type_ab(sfb);
DG_HOST_ASSERT(gran_k_a == 128 and gran_k_b == 128);
sm90_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d, batch_size, m, n, k, major_a, major_b, major_sfb, compiled_dims);
}
}
@@ -187,26 +188,26 @@ static void fp8_einsum(const std::string& expr,
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;
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" and arch_major == 10) {
// (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;
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" and arch_major == 10) {
// (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});
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));
csrc/apis/gemm.hpp
+68 -65
View File
@@ -6,7 +6,7 @@
#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_fp4_gemm_1d1d.hpp"
#include "../jit_kernels/impls/sm100_bf16_gemm.hpp"
#endif
@@ -23,7 +23,7 @@ static bool early_return(const int& m, const int &n, const int& k,
return true;
// Checks
const bool& is_cd_same = c.has_value() and c->data_ptr() == d.data_ptr();
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()) {
@@ -57,8 +57,8 @@ static void fp8_fp4_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
const std::string& compiled_dims,
const bool& disable_ue8m0_cast) {
// Shape must be `[M, K] @ [N, K].T`
const auto& major_a = get_major_type_ab(a.first);
const auto& major_b = get_major_type_ab(b.first);
const auto major_a = get_major_type_ab(a.first);
const auto major_b = get_major_type_ab(b.first);
if (fp8_requires_k_major()) {
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
DG_HOST_ASSERT(major_b == cute::UMMA::Major::K);
@@ -89,7 +89,7 @@ static void fp8_fp4_gemm_nt(const std::pair<torch::Tensor, torch::Tensor>& a,
if (gran_n == 1) {
sm90_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, k, major_a, major_b, compiled_dims);
} else {
const auto& major_sfb = get_major_type_ab(sfb);
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) {
@@ -152,8 +152,8 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
// Shape must be `[M, K] @ [G, N, K].mT`
const auto& major_a = get_major_type_ab(a.first);
const auto& major_b = get_major_type_ab(b.first);
const auto major_a = get_major_type_ab(a.first);
const auto major_b = get_major_type_ab(b.first);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K);
if (fp8_requires_k_major())
DG_HOST_ASSERT(major_b == cute::UMMA::Major::K);
@@ -171,10 +171,10 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
// Layout checks
if (use_psum_layout) {
const auto& [num_groups_] = get_shape<1>(grouped_layout);
const auto [num_groups_] = get_shape<1>(grouped_layout);
DG_HOST_ASSERT(num_groups == num_groups_);
} else {
const auto& [m__] = get_shape<1>(grouped_layout);
const auto [m__] = get_shape<1>(grouped_layout);
DG_HOST_ASSERT(m == m__);
DG_HOST_ASSERT(not expected_m_for_psum_layout.has_value());
}
@@ -192,10 +192,10 @@ static void m_grouped_fp8_fp4_gemm_nt_contiguous(const std::pair<torch::Tensor,
// Dispatch implementation
if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
const auto& major_sfb = get_major_type_ab(sfb);
DG_HOST_ASSERT(not use_psum_layout);
const auto major_sfb = get_major_type_ab(sfb);
sm90_m_grouped_fp8_gemm_contiguous_1d2d(a.first, sfa, b.first, sfb, d, grouped_layout,
num_groups, m, n, k, major_a, major_b, major_sfb, compiled_dims);
num_groups, m, n, k, major_a, major_b, major_sfb,
compiled_dims, use_psum_layout, expected_m_for_psum_layout);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
sm100_m_grouped_fp8_fp4_gemm_contiguous_1d1d(a.first, sfa, b.first, sfb, d, grouped_layout,
num_groups, m, n, k, gran_k_a, gran_k_b, major_a, major_b,
@@ -230,8 +230,8 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair<torch::Tensor, torc
const std::string& compiled_dims,
const bool& disable_ue8m0_cast) {
// Shape must be `[G, M, K] @ [G, N, K].mT`
const auto& major_a = get_major_type_ab(a.first);
const auto& major_b = get_major_type_ab(b.first);
const auto major_a = get_major_type_ab(a.first);
const auto major_b = get_major_type_ab(b.first);
DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(masked_m.is_contiguous());
@@ -256,7 +256,7 @@ static void m_grouped_fp8_fp4_gemm_nt_masked(const std::pair<torch::Tensor, torc
// Dispatch implementation
if (arch_major == 9 and sfa.scalar_type() == torch::kFloat) {
const auto& major_sfb = get_major_type_ab(sfb);
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, major_sfb, compiled_dims);
} else if (arch_major == 10 and sfa.scalar_type() == torch::kInt) {
@@ -277,12 +277,15 @@ static void k_grouped_fp8_gemm_tn_contiguous(const std::pair<torch::Tensor, torc
const std::tuple<int, int, int>& recipe,
const std::string& compiled_dims) {
// Must be 1D1D kernel
DG_HOST_ASSERT(recipe == std::make_tuple(1, 1, 128));
DG_HOST_ASSERT(std::get<0>(recipe) == 1 and std::get<1>(recipe) == 1);
const int gran_k = std::get<2>(recipe);
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
// Shape checks
const auto& [num_groups, m, n] = get_shape<3>(d);
const auto& [sum_k_ , m_] = get_shape<2>(a.first);
const auto& [sum_k__, n_] = get_shape<2>(b.first);
const auto [num_groups, m, n] = get_shape<3>(d);
const auto [sum_k_ , m_] = get_shape<2>(a.first);
const auto [sum_k__, n_] = get_shape<2>(b.first);
const int sum_k = std::accumulate(ks.begin(), ks.end(), 0);
DG_HOST_ASSERT(m == m_ and n == n_ and sum_k == sum_k_ and sum_k == sum_k__);
@@ -297,13 +300,13 @@ static void k_grouped_fp8_gemm_tn_contiguous(const std::pair<torch::Tensor, torc
return;
// Transform SF with padding
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);
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);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 10) {
sm100_k_grouped_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, ks, ks_tensor,
sm100_k_grouped_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, ks, ks_tensor, gran_k,
cute::UMMA::Major::MN, cute::UMMA::Major::MN, compiled_dims);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
@@ -322,9 +325,9 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torc
DG_HOST_ASSERT(recipe == std::make_tuple(1, 1, 128));
// Shape checks
const auto& [num_groups, m, n] = get_shape<3>(d);
const auto& sum_mk = a.first.numel();
const auto& sum_nk = b.first.numel();
const auto [num_groups, m, n] = get_shape<3>(d);
const auto sum_mk = a.first.numel();
const auto sum_nk = b.first.numel();
const int sum_k = std::accumulate(ks.begin(), ks.end(), 0);
DG_HOST_ASSERT(sum_mk == static_cast<int64_t>(sum_k) * m);
DG_HOST_ASSERT(sum_nk == static_cast<int64_t>(sum_k) * n);
@@ -340,17 +343,17 @@ static void k_grouped_fp8_gemm_nt_contiguous(const std::pair<torch::Tensor, torc
return;
// Transform SF with padding
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);
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);
// Allocate tensormap buffer
// `4` means the double buffering for both A and B operands (2 * 2)
const auto& num_sms = device_runtime->get_num_sms();
const auto& tensor_map_buffer = torch::empty({num_sms * 4 * static_cast<int>(sizeof(CUtensorMap))},
a.first.options().dtype(torch::kByte));
const auto num_sms = device_runtime->get_num_sms();
const auto tensor_map_buffer = torch::empty({num_sms * 4 * static_cast<int>(sizeof(CUtensorMap))},
a.first.options().dtype(torch::kByte));
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
sm90_k_grouped_fp8_gemm_1d1d(a.first, sfa, b.first, sfb, c, d, m, n, ks, ks_tensor, tensor_map_buffer,
cute::UMMA::Major::K, cute::UMMA::Major::K, compiled_dims);
@@ -367,16 +370,16 @@ static void bf16_gemm_nt(const torch::Tensor& a,
const std::optional<torch::Tensor>& c,
const std::string& compiled_dims) {
// Shape must be `[M, K] @ [N, K].T`
const auto& major_a = get_major_type_ab(a);
const auto& major_b = get_major_type_ab(b);
const auto major_a = get_major_type_ab(a);
const auto major_b = get_major_type_ab(b);
// C/D must be N-major
check_major_type_cd(d);
// Type and shape checks
const auto& [m , k ] = get_shape<2>(a);
const auto& [n , k_] = get_shape<2>(b);
const auto& [m_, n_] = get_shape<2>(d);
const auto [m , k ] = get_shape<2>(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(a.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(b.scalar_type() == torch::kBFloat16);
@@ -387,7 +390,7 @@ static void bf16_gemm_nt(const torch::Tensor& a,
return;
// Dispatch into different implements
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
sm90_bf16_gemm(a, b, c, d, m, n, k, major_a, major_b, compiled_dims);
} else if (arch_major == 10) {
@@ -427,15 +430,15 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc
const bool& use_psum_layout,
const std::optional<int>& expected_m_for_psum_layout) {
// Shape must be `[M, K] @ [G, N, K].mT`
const auto& major_a = get_major_type_ab(a);
const auto& major_b = get_major_type_ab(b);
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(grouped_layout.is_contiguous());
// Type and shape checks
const auto& [m, k] = get_shape<2>(a);
const auto& [num_groups, n, k_] = get_shape<3>(b);
const auto& [m_, n_] = get_shape<2>(d);
const auto [m, k] = get_shape<2>(a);
const auto [num_groups, n, k_] = get_shape<3>(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 and num_groups > 0);
DG_HOST_ASSERT(a.scalar_type() == torch::kBFloat16);
@@ -445,10 +448,10 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc
// Layout checks
if (use_psum_layout) {
const auto& [num_groups_] = get_shape<1>(grouped_layout);
const auto [num_groups_] = get_shape<1>(grouped_layout);
DG_HOST_ASSERT(num_groups == num_groups_);
} else {
const auto& [m__] = get_shape<1>(grouped_layout);
const auto [m__] = get_shape<1>(grouped_layout);
DG_HOST_ASSERT(m == m__);
DG_HOST_ASSERT(not expected_m_for_psum_layout.has_value());
}
@@ -461,11 +464,11 @@ static void m_grouped_bf16_gemm_nt_contiguous(const torch::Tensor& a, const torc
return;
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
DG_HOST_ASSERT(not use_psum_layout);
sm90_m_grouped_bf16_gemm_contiguous(a, b, d, grouped_layout,
num_groups, m, n, k, major_a, major_b, compiled_dims);
num_groups, m, n, k, major_a, major_b, compiled_dims,
use_psum_layout, expected_m_for_psum_layout);
} else if (arch_major == 10) {
sm100_m_grouped_bf16_gemm_contiguous(a, b, d, grouped_layout,
num_groups, m, n, k, major_a, major_b, compiled_dims,
@@ -487,16 +490,16 @@ static void m_grouped_bf16_gemm_nt_masked(const torch::Tensor& a, const torch::T
const torch::Tensor& d, const torch::Tensor& masked_m,
const int& expected_m, const std::string& compiled_dims) {
// Shape must be `[G, M, K] @ [G, N, K].mT`
const auto& major_a = get_major_type_ab(a);
const auto& major_b = get_major_type_ab(b);
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 and major_b == cute::UMMA::Major::K);
DG_HOST_ASSERT(masked_m.is_contiguous());
// Type and shape checks
const auto& [num_groups, m, k] = get_shape<3>(a);
const auto& [num_groups_, n, k_] = get_shape<3>(b);
const auto& [num_groups__, m_, n_] = get_shape<3>(d);
const auto& num_groups___ = static_cast<int>(masked_m.numel());
const auto [num_groups, m, k] = get_shape<3>(a);
const auto [num_groups_, n, k_] = get_shape<3>(b);
const auto [num_groups__, m_, n_] = get_shape<3>(d);
const auto num_groups___ = static_cast<int>(masked_m.numel());
DG_HOST_ASSERT(num_groups == num_groups_ and num_groups == num_groups__ and num_groups == num_groups___);
DG_HOST_ASSERT(m == m_ and n == n_ and k == k_);
DG_HOST_ASSERT(expected_m > 0 and m > 0 and n > 0 and k > 0 and num_groups > 0);
@@ -509,7 +512,7 @@ static void m_grouped_bf16_gemm_nt_masked(const torch::Tensor& a, const torch::T
check_major_type_cd(d);
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
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);
@@ -529,9 +532,9 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a,
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& [sum_k_ , m_] = get_shape<2>(a);
const auto& [sum_k__, n_] = get_shape<2>(b);
const auto [num_groups, m, n] = get_shape<3>(d);
const auto [sum_k_ , m_] = get_shape<2>(a);
const auto [sum_k__, n_] = get_shape<2>(b);
const int sum_k = std::accumulate(ks.begin(), ks.end(), 0);
DG_HOST_ASSERT(m == m_ and n == n_ and sum_k == sum_k_ and sum_k == sum_k__);
@@ -546,7 +549,7 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a,
return;
// Dispatch implementation
const auto& arch_major = device_runtime->get_arch_major();
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);
@@ -562,20 +565,20 @@ static void k_grouped_bf16_gemm_tn_contiguous(const torch::Tensor& a,
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`
const auto& major_a = get_major_type_ab(a);
const auto& major_b = get_major_type_ab(b);
const auto major_a = get_major_type_ab(a);
const auto major_b = get_major_type_ab(b);
// Type and shape checks
const auto& [m , k ] = get_shape<2>(a);
const auto& [n , k_] = get_shape<2>(b);
const auto& [m_, n_] = get_shape<2>(d);
const auto [m , k ] = get_shape<2>(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_);
// 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);
cublaslt_gemm(a, b, d, m, n, k, major_a, major_b, c.has_value());
}
static void cublaslt_gemm_nn(const torch::Tensor& a, const torch::Tensor& b,
csrc/apis/hyperconnection.hpp
+7 -7
View File
@@ -24,16 +24,16 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a,
DG_HOST_ASSERT(sqr_sum.is_contiguous());
// Type and shape checks
const auto& [m, k ] = get_shape<2>(a);
const auto& [n, k_] = get_shape<2>(b);
const auto [m, k ] = get_shape<2>(a);
const auto [n, k_] = get_shape<2>(b);
if (num_splits.has_value()) {
const auto& [num_splits_, m_, n_] = get_shape<3>(d);
const auto& [num_splits__, m__] = get_shape<2>(sqr_sum);
const auto [num_splits_, m_, n_] = get_shape<3>(d);
const auto [num_splits__, m__] = get_shape<2>(sqr_sum);
DG_HOST_ASSERT(num_splits.value() == num_splits_ and num_splits.value() == num_splits__ and num_splits.value() >= 1);
DG_HOST_ASSERT(m == m_ and m == m__ and n == n_ and k == k_);
} else {
const auto& [m_, n_] = get_shape<2>(d);
const auto& [m__] = get_shape<1>(sqr_sum);
const auto [m_, n_] = get_shape<2>(d);
const auto [m__] = get_shape<1>(sqr_sum);
DG_HOST_ASSERT(m == m_ and m == m__ and n == n_ and k == k_);
}
DG_HOST_ASSERT(n > 0 and k > 0);
@@ -47,7 +47,7 @@ static void tf32_hc_prenorm_gemm(const torch::Tensor& a,
return;
// Dispatch into different implements
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
if (arch_major == 9) {
sm90_tf32_hc_prenorm_gemm(a, b, d, sqr_sum, m, n, k, num_splits.has_value() ? num_splits.value() : 1);
} else if (arch_major == 10) {
csrc/apis/layout.hpp
+46 -23
View File
@@ -1,5 +1,6 @@
#pragma once
#include "../jit_kernels/heuristics/runtime.hpp"
#include "../utils/layout.hpp"
#include "../utils/compatibility.hpp"
@@ -12,21 +13,24 @@ namespace deep_gemm::layout {
#if DG_TENSORMAP_COMPATIBLE
static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
const int& mn, const int& k,
const std::optional<std::tuple<int, int, int>>& recipe,
const std::optional<std::tuple<int, int>>& recipe_ab,
const std::variant<std::tuple<int, int, int>,
std::tuple<int, int>>& recipe,
const std::optional<int>& num_groups,
const bool& is_sfa,
const std::optional<bool>& is_sfa,
const bool& disable_ue8m0_cast) {
const auto& arch_major = device_runtime->get_arch_major();
const auto arch_major = device_runtime->get_arch_major();
// Get granularity MN/K from recipe
int gran_mn, gran_k;
if (recipe.has_value()) {
DG_HOST_ASSERT(not recipe_ab.has_value());
gran_mn = is_sfa ? std::get<0>(recipe.value()) : std::get<1>(recipe.value());
gran_k = std::get<2>(recipe.value());
if (auto p = std::get_if<std::tuple<int, int, int>>(&recipe)) {
DG_HOST_ASSERT(is_sfa.has_value());
gran_mn = is_sfa.value() ? std::get<0>(*p) : std::get<1>(*p);
gran_k = std::get<2>(*p);
} else if (auto p = std::get_if<std::tuple<int, int>>(&recipe)) {
DG_HOST_ASSERT(not is_sfa.has_value());
std::tie(gran_mn, gran_k) = *p;
} else {
DG_HOST_ASSERT(recipe_ab.has_value());
std::tie(gran_mn, gran_k) = recipe_ab.value();
DG_HOST_UNREACHABLE("Invalid recipe");
}
// Pre-transform checks
@@ -43,8 +47,8 @@ static torch::Tensor transform_sf_into_required_layout(const torch::Tensor& sf,
// (FP32, x, gran_k) on SM100: transform to (INT, 1, gran_k), TMA-aligned and MN-major
if (sf.scalar_type() == torch::kFloat and (gran_k == 32 or gran_k == 128) and arch_major == 10) {
DG_HOST_ASSERT(not disable_ue8m0_cast);
const auto& broadcasted = gran_mn == 1 ? sf :
sf.index_select(-2, torch::arange(mn, at::TensorOptions().device(sf.device())).floor_divide_(gran_mn));
const auto broadcasted = gran_mn == 1 ? sf :
sf.index_select(-2, torch::arange(mn, at::TensorOptions().device(sf.device())).floor_divide_(gran_mn));
return get_mn_major_tma_aligned_packed_ue8m0_tensor(broadcasted);
}
@@ -64,11 +68,19 @@ static std::tuple<torch::Tensor, torch::Tensor, int, int> transform_sf_pair_into
const std::optional<int>& num_groups_a,
const std::optional<int>& num_groups_b,
const bool& disable_ue8m0_cast = false) {
DG_HOST_ASSERT(recipe_a.has_value() == recipe_b.has_value());
// Use default recipe, if none is specified
if (not recipe_a.has_value() and not recipe.has_value())
recipe = get_default_recipe(sfa.scalar_type(), sfb.scalar_type());
const auto transformed_sfa = transform_sf_into_required_layout(sfa, m, k, recipe, recipe_a, num_groups_a, true, disable_ue8m0_cast);
const auto transformed_sfb = transform_sf_into_required_layout(sfb, n, k, recipe, recipe_b, num_groups_b, false, disable_ue8m0_cast);
// Must be either 'recipe' or the 'recipe_a' + 'recipe_b' pair.
DG_HOST_ASSERT(recipe_a.has_value() == recipe_b.has_value());
DG_HOST_ASSERT(recipe_a.has_value() != recipe.has_value());
// Transform SFA and SFB layout
const auto transformed_sfa = recipe.has_value() ? transform_sf_into_required_layout(sfa, m, k, recipe.value(), num_groups_a, true, disable_ue8m0_cast)
: transform_sf_into_required_layout(sfa, m, k, recipe_a.value(), num_groups_a, std::nullopt, disable_ue8m0_cast);
const auto transformed_sfb = recipe.has_value() ? transform_sf_into_required_layout(sfb, n, k, recipe.value(), num_groups_b, false, disable_ue8m0_cast)
: transform_sf_into_required_layout(sfb, n, k, recipe_b.value(), num_groups_b, std::nullopt, disable_ue8m0_cast);
const int gran_k_a = recipe_a.has_value() ? std::get<1>(recipe_a.value()) : std::get<2>(recipe.value());
const int gran_k_b = recipe_b.has_value() ? std::get<1>(recipe_b.value()) : std::get<2>(recipe.value());
return std::make_tuple(transformed_sfa, transformed_sfb, gran_k_a, gran_k_b);
@@ -79,8 +91,12 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
const torch::Tensor& ks_tensor,
const std::tuple<int, int, int>& recipe) {
DG_HOST_ASSERT(sf.dim() == 2);
DG_HOST_ASSERT(recipe == std::make_tuple(1, 1, 128));
const auto& arch_major = device_runtime->get_arch_major();
DG_HOST_ASSERT(std::get<0>(recipe) == 1 and std::get<1>(recipe) == 1);
const int gran_k = std::get<2>(recipe);
DG_HOST_ASSERT(gran_k == 32 or gran_k == 128);
const auto arch_major = device_runtime->get_arch_major();
// FP32 on SM90
if (sf.scalar_type() == torch::kFloat and arch_major == 9)
@@ -88,7 +104,7 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
// FP32 on SM100
if (sf.scalar_type() == torch::kFloat and arch_major == 10)
return get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf, ks_tensor, ks);
return get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor(sf, ks_tensor, ks, gran_k);
// INT on SM100
if (sf.scalar_type() == torch::kInt and arch_major == 10)
@@ -100,12 +116,11 @@ static torch::Tensor transform_k_grouped_sf_into_required_layout(const torch::Te
#endif
static void register_apis(pybind11::module_& m) {
#if DG_TENSORMAP_COMPATIBLE
m.def("transform_sf_into_required_layout", &transform_sf_into_required_layout,
py::arg("sf"), py::arg("mn"), py::arg("k"),
py::arg("recipe") = std::nullopt, py::arg("recipe_ab") = std::nullopt,
py::arg("num_groups") = std::nullopt, py::arg("is_sfa") = false,
py::arg("sf"), py::arg("mn"), py::arg("k"), py::arg("recipe"),
py::arg("num_groups") = std::nullopt,
py::arg("is_sfa") = std::nullopt,
py::arg("disable_ue8m0_cast") = false);
m.def("get_tma_aligned_size", &get_tma_aligned_size);
@@ -114,7 +129,15 @@ static void register_apis(pybind11::module_& m) {
m.def("get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor", &get_k_grouped_mn_major_tma_aligned_packed_ue8m0_tensor);
#endif
m.def("get_mk_alignment_for_contiguous_layout", &get_mk_alignment_for_contiguous_layout);
m.def("set_mk_alignment_for_contiguous_layout", [&](const int& new_value) {
heuristics_runtime->set_mk_alignment_for_contiguous_layout(new_value);
});
m.def("get_mk_alignment_for_contiguous_layout", [&]() {
return heuristics_runtime->get_mk_alignment_for_contiguous_layout();
});
m.def("get_theoretical_mk_alignment_for_contiguous_layout", [&](const std::optional<int>& expected_m) {
return heuristics_runtime->get_theoretical_mk_alignment_for_contiguous_layout(expected_m);
}, py::arg("expected_m") = std::nullopt);
}
} // namespace deep_gemm::layout
csrc/apis/mega.hpp
+216
View File
@@ -0,0 +1,216 @@
#pragma once
#include <functional>
#include <pybind11/functional.h>
#if DG_TENSORMAP_COMPATIBLE
#include "../jit/compiler.hpp"
#endif
#include "../jit/device_runtime.hpp"
#include "../jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp"
namespace deep_gemm::mega {
static std::tuple<int64_t, std::function<std::tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>(const torch::Tensor&)>>
get_symm_buffer_size_for_mega_moe(
const int& num_ranks, const int& num_experts,
const int& num_max_tokens_per_rank, const int& num_topk,
const int& hidden, const int& intermediate_hidden,
const bool& use_fp8_dispatch, const std::string& activation) {
DG_HOST_ASSERT(num_experts % num_ranks == 0);
// Workspace bytes
const auto block_m = get_block_m_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
const auto workspace = layout::Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk, block_m);
// Layouts
const auto fp8_token_layout = layout::Data(hidden);
const auto bf16_token_layout = layout::Data(hidden * 2);
const auto fp8_intermediate_token_layout = layout::Data(intermediate_hidden);
const auto fp8_sf_layout = layout::Data(hidden / 32);
const auto fp8_intermediate_sf_layout = layout::Data(intermediate_hidden / 32);
const auto input_topk_idx_layout = layout::Data(num_topk * sizeof(int64_t), false);
const auto input_topk_weights_layout = layout::Data(num_topk * sizeof(float), false);
const auto l1_topk_weights_layout = layout::Data(sizeof(float), false);
// Input buffers
const auto input_token_buffer = layout::Buffer(
fp8_token_layout, 1, num_max_tokens_per_rank,
workspace.get_end_ptr());
const auto input_sf_buffer = layout::Buffer(
fp8_sf_layout, 1, num_max_tokens_per_rank,
input_token_buffer.get_end_ptr());
const auto input_topk_idx_buffer = layout::Buffer(
input_topk_idx_layout, 1, num_max_tokens_per_rank,
input_sf_buffer.get_end_ptr());
const auto input_topk_weights_buffer = layout::Buffer(
input_topk_weights_layout, 1, num_max_tokens_per_rank,
input_topk_idx_buffer.get_end_ptr());
// Buffer configs
const auto num_max_pool_tokens = static_cast<int>(workspace.num_max_pool_tokens);
const auto num_padded_sf_pool_tokens = layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m);
// L1 input buffer
const auto l1_token_buffer = layout::Buffer(
fp8_token_layout, 1, num_max_pool_tokens,
input_topk_weights_buffer.get_end_ptr());
const auto l1_sf_buffer = layout::Buffer(
fp8_sf_layout, 1, num_padded_sf_pool_tokens,
l1_token_buffer.get_end_ptr());
const auto l1_topk_weights_buffer = layout::Buffer(
l1_topk_weights_layout, 1, num_max_pool_tokens,
l1_sf_buffer.get_end_ptr());
// L2 input buffer
const auto l2_token_buffer = layout::Buffer(
fp8_intermediate_token_layout, 1, num_max_pool_tokens,
l1_topk_weights_buffer.get_end_ptr());
const auto l2_sf_buffer = layout::Buffer(
fp8_intermediate_sf_layout, 1, num_padded_sf_pool_tokens,
l2_token_buffer.get_end_ptr());
// Combine input buffer: BF16 tokens for cross-rank combine
const auto combine_token_buffer = layout::Buffer(
bf16_token_layout, num_topk, num_max_tokens_per_rank,
l2_sf_buffer.get_end_ptr());
// Check SF buffer requirements
DG_HOST_ASSERT(hidden % 128 == 0 and intermediate_hidden % 128 == 0);
DG_HOST_ASSERT(num_padded_sf_pool_tokens % 4 == 0);
// Slice function: creates `(x, x_sf, topk_weights, topk_idx, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf)` tensor views from the raw buffer
// NOTES: `x_sf` is K-major, while `l1_acts_sf` and `l2_acts_sf` are M-major
auto slice_input_buffers = [=](const torch::Tensor& buffer) {
auto x = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_token_buffer.base)),
{num_max_tokens_per_rank, hidden},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto x_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_sf_buffer.base)),
{num_max_tokens_per_rank, hidden / 128},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto topk_idx = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_topk_idx_buffer.base)),
{num_max_tokens_per_rank, num_topk},
torch::TensorOptions().dtype(torch::kInt64).device(buffer.device()));
auto topk_weights = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(input_topk_weights_buffer.base)),
{num_max_tokens_per_rank, num_topk},
torch::TensorOptions().dtype(torch::kFloat32).device(buffer.device()));
auto l1_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_token_buffer.base)),
{num_max_pool_tokens, hidden},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto l1_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l1_sf_buffer.base)),
{num_padded_sf_pool_tokens, hidden / 128},
{1, num_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
auto l2_acts = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_token_buffer.base)),
{num_max_pool_tokens, intermediate_hidden},
torch::TensorOptions().dtype(torch::kFloat8_e4m3fn).device(buffer.device()));
auto l2_acts_sf = torch::from_blob(
math::advance_ptr(buffer.data_ptr(), reinterpret_cast<int64_t>(l2_sf_buffer.base)),
{num_padded_sf_pool_tokens, intermediate_hidden / 128},
{1, num_padded_sf_pool_tokens},
torch::TensorOptions().dtype(torch::kInt).device(buffer.device()));
return std::make_tuple(x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf);
};
return {reinterpret_cast<int64_t>(combine_token_buffer.get_end_ptr()), slice_input_buffers};
}
static void fp8_fp4_mega_moe(
const torch::Tensor& y,
const std::tuple<torch::Tensor, torch::Tensor>& l1_weights_,
const std::tuple<torch::Tensor, torch::Tensor>& l2_weights_,
const torch::Tensor& sym_buffer,
const std::vector<int64_t>& sym_buffer_ptrs, const int& rank_idx,
const int& num_max_tokens_per_rank,
const int& num_experts, const int& num_topk,
const std::tuple<int, int, int>& recipe,
const std::string& activation,
const std::optional<float>& activation_clamp_opt,
const bool& fast_math) {
const auto [l1_weights, l1_weights_sf] = l1_weights_;
const auto [l2_weights, l2_weights_sf] = l2_weights_;
// Config checks
const auto num_tokens = static_cast<int>(y.size(0));
const auto [rm, rn, rk] = recipe;
DG_HOST_ASSERT(rm == 1 and rn == 1 and rk == 32);
DG_HOST_ASSERT(activation == "swiglu");
// Activation checks
const auto activation_clamp =
activation_clamp_opt.value_or(std::numeric_limits<float>::infinity());
DG_HOST_ASSERT(activation_clamp >= 0);
// Tensor checks
DG_HOST_ASSERT(get_major_type_ab(l1_weights) == cute::UMMA::Major::K);
DG_HOST_ASSERT(get_major_type_ab(l2_weights) == cute::UMMA::Major::K);
const auto arch_major = device_runtime->get_arch_major();
const auto [num_experts_per_rank, intermediate_hidden_2, hidden] =
check_grouped_ab_fp8_fp4(l1_weights, cute::UMMA::Major::K, arch_major);
const auto [num_experts_per_rank_, hidden_, intermediate_hidden] =
check_grouped_ab_fp8_fp4(l2_weights, cute::UMMA::Major::K, arch_major);
DG_HOST_ASSERT(num_tokens <= num_max_tokens_per_rank);
DG_HOST_ASSERT(num_experts_per_rank == num_experts_per_rank_);
DG_HOST_ASSERT(hidden == hidden_);
DG_HOST_ASSERT(intermediate_hidden_2 == 2 * intermediate_hidden);
DG_HOST_ASSERT(l1_weights.is_contiguous() and l2_weights.is_contiguous());
// Check weight SF layout for UE8M0 packing, MN-major, and TMA alignment
constexpr int kGranMN = 1, kGranK = 32;
check_sf_layout(l1_weights_sf, intermediate_hidden * 2, hidden, kGranMN, kGranK,
num_experts_per_rank, true, false, torch::kInt);
check_sf_layout(l2_weights_sf, hidden, intermediate_hidden, kGranMN, kGranK,
num_experts_per_rank, true, false, torch::kInt);
// Check buffer bytes
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts_ = num_experts_per_rank * num_ranks;
const auto [num_required_bytes, slice] = get_symm_buffer_size_for_mega_moe(
num_ranks, num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
true, "swiglu");
DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast<size_t>(num_required_bytes));
DG_HOST_ASSERT(num_experts == num_experts_);
// Already registered tensors
const auto [x, x_sf, topk_idx, topk_weights, l1_acts, l1_acts_sf, l2_acts, l2_acts_sf] = slice(sym_buffer);
// Dispatch into different architectures
if (arch_major == 10) {
sm100_fp8_fp4_mega_moe(y,
l1_acts, l1_acts_sf,
l2_acts, l2_acts_sf,
l1_weights, l2_weights,
l1_weights_sf, l2_weights_sf,
sym_buffer_ptrs,
rank_idx, num_max_tokens_per_rank,
num_experts_per_rank,
num_tokens, num_topk,
hidden, intermediate_hidden,
activation_clamp, fast_math);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
// Zero the entire symmetric buffer for debug mode
// NOTES: caller must re-copy inputs into the buffer before each kernel call
if (get_env<int>("DG_COMM_KERNEL_DEBUG"))
sym_buffer.zero_();
}
static void register_apis(pybind11::module_& m) {
#if DG_TENSORMAP_COMPATIBLE
m.def("get_block_m_for_mega_moe", &get_block_m_for_mega_moe);
m.def("get_symm_buffer_size_for_mega_moe", &get_symm_buffer_size_for_mega_moe);
m.def("fp8_fp4_mega_moe", &fp8_fp4_mega_moe);
#endif
}
} // namespace deep_gemm::mega
csrc/apis/runtime.hpp
+20
View File
@@ -4,6 +4,7 @@
#include "../jit/compiler.hpp"
#endif
#include "../jit/device_runtime.hpp"
#include "../jit_kernels/heuristics/runtime.hpp"
namespace deep_gemm::runtime {
@@ -20,10 +21,29 @@ static void register_apis(pybind11::module_& m) {
m.def("get_tc_util", [&]() {
return device_runtime->get_tc_util();
});
m.def("set_pdl", [&](const bool& new_enable_pdl) {
device_runtime->set_pdl(new_enable_pdl);
});
m.def("get_pdl", [&]() {
return device_runtime->get_pdl();
});
m.def("set_ignore_compile_dims", [&](const bool& new_value) {
heuristics_runtime->set_ignore_compile_dims(new_value);
});
m.def("set_block_size_multiple_of", [&](const std::variant<int, std::tuple<int, int>>& new_value) {
if (std::holds_alternative<int>(new_value)) {
auto x = std::get<int>(new_value);
heuristics_runtime->set_block_size_multiple_of(x, x);
} else {
auto [x, y] = std::get<std::tuple<int, int>>(new_value);
heuristics_runtime->set_block_size_multiple_of(x, y);
}
});
m.def("init", [&](const std::string& library_root_path, const std::string& cuda_home_path_by_python) {
#if DG_TENSORMAP_COMPATIBLE
Compiler::prepare_init(library_root_path, cuda_home_path_by_python);
KernelRuntime::prepare_init(cuda_home_path_by_python);
IncludeParser::prepare_init(library_root_path);
#endif
});
}