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

Add various optimizations and Mega MoE benchmarks (#316)

Browse Source 
* Merge with private repo

* Add Mega MoE Benchmark

* Minor fix

* Update

---------

Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com>
This commit is contained in:
Zhean Xu
2026-04-24 18:41:37 +08:00
committed by GitHub
parent 7f2a703ed5
commit 891d57b4db
21 changed files with 1276 additions and 372 deletions
Download Patch File Download Diff File Expand all files Collapse all files

51
csrc/apis/attention.hpp
View File

@@ -190,12 +190,13 @@ static torch::Tensor fp8_fp4_mqa_logits(const std::tuple<torch::Tensor, std::opt
return logits;
}
static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_lens, int block_kv, int num_sms) {
static torch::Tensor get_paged_mqa_logits_metadata(const torch::Tensor& context_lens, int block_kv, int num_sms, const std::optional<torch::Tensor>& indices) {
// 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);
const bool is_varlen = indices.has_value();
DG_HOST_ASSERT(context_lens.scalar_type() == torch::kInt);
DG_HOST_ASSERT(context_lens.is_contiguous());
@@ -204,9 +205,16 @@ 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) {
if (is_varlen) {
const auto& indices_tensor = indices.value();
DG_HOST_ASSERT(arch_major == 10 and next_n == 1 and (block_kv == 64 or block_kv == 32));
DG_HOST_ASSERT(indices_tensor.dim() == 1 and indices_tensor.size(0) == batch_size);
DG_HOST_ASSERT(indices_tensor.is_contiguous());
DG_HOST_ASSERT(indices_tensor.scalar_type() == torch::kInt);
smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, true, indices_tensor.data_ptr<int>());
} else 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);
smxx_paged_mqa_logits_metadata(context_lens, schedule_metadata, batch_size, next_n, block_kv, num_sms, is_context_lens_2d, false, nullptr);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
@@ -222,7 +230,8 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
const torch::Tensor& schedule_meta,
const int& max_context_len,
const bool& clean_logits,
const at::ScalarType& logits_dtype) {
const at::ScalarType& logits_dtype,
const std::optional<torch::Tensor>& indices) {
const auto [q_fp, q_sf] = q;
const bool is_fp4 = q_sf.has_value();
@@ -321,6 +330,17 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
DG_HOST_ASSERT(block_table.stride(1) == 1);
DG_HOST_ASSERT(block_table.scalar_type() == torch::kInt);
// Check indices
const bool is_varlen = indices.has_value();
const auto arch_major = device_runtime->get_arch_major();
const auto indices_tensor = indices.value_or(torch::Tensor());
if (is_varlen) {
DG_HOST_ASSERT(arch_major == 10 and next_n == 1);
DG_HOST_ASSERT(indices_tensor.dim() == 1 and indices_tensor.size(0) == batch_size);
DG_HOST_ASSERT(indices_tensor.is_contiguous());
DG_HOST_ASSERT(indices_tensor.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);
@@ -344,15 +364,14 @@ static torch::Tensor fp8_fp4_paged_mqa_logits(const std::tuple<torch::Tensor, st
DG_HOST_ASSERT(logits_dtype == torch::kFloat32 or logits_dtype == torch::kBFloat16);
// Dispatch implementation
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,
sm100_fp4_paged_mqa_logits(q_fp, q_sf.value(), kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, 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);
is_varlen, 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,
smxx_fp8_paged_mqa_logits(q_fp, kv_cache, kv_cache_sf, weights, context_lens, logits, block_table, indices_tensor, 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);
is_varlen, aligned_max_context_len, block_table_stride, num_sms, split_kv);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
@@ -386,10 +405,11 @@ static torch::Tensor fp8_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& block_table,
const torch::Tensor& schedule_meta,
const int& max_context_len,
const bool& clean_logits) {
const bool& clean_logits,
const std::optional<torch::Tensor>& indices) {
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);
max_context_len, clean_logits, torch::kFloat, indices);
}
#endif
@@ -407,13 +427,15 @@ static void register_apis(pybind11::module_& m) {
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"));
py::arg("context_lens"), py::arg("block_kv"), py::arg("num_sms"),
py::arg("indices") = std::nullopt);
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);
py::arg("logits_dtype") = torch::kFloat32,
py::arg("indices") = std::nullopt);
// Legacy API
m.def("fp8_mqa_logits", &fp8_mqa_logits,
py::arg("q"), py::arg("kv"), py::arg("weights"),
@@ -423,7 +445,8 @@ static void register_apis(pybind11::module_& m) {
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);
py::arg("max_context_len"), py::arg("clean_logits") = false,
py::arg("indices") = std::nullopt);
#endif
}

55
csrc/apis/mega.hpp
View File

@@ -11,6 +11,10 @@
namespace deep_gemm::mega {
static int get_token_alignment_for_mega_moe() {
return layout::kLCMCandidateBlockM;
}
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,
@@ -20,8 +24,7 @@ get_symm_buffer_size_for_mega_moe(
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);
const auto workspace = layout::Workspace(nullptr, num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
// Layouts
const auto fp8_token_layout = layout::Data(hidden);
@@ -49,14 +52,20 @@ get_symm_buffer_size_for_mega_moe(
// 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);
int num_max_padded_sf_pool_tokens = 0;
for (int block_m: layout::kCandidateBlockM) {
num_max_padded_sf_pool_tokens = std::max(
num_max_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,
fp8_sf_layout, 1, num_max_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,
@@ -67,7 +76,7 @@ get_symm_buffer_size_for_mega_moe(
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,
fp8_intermediate_sf_layout, 1, num_max_padded_sf_pool_tokens,
l2_token_buffer.get_end_ptr());
// Combine input buffer: BF16 tokens for cross-rank combine
@@ -77,7 +86,7 @@ get_symm_buffer_size_for_mega_moe(
// 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);
DG_HOST_ASSERT(num_max_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
@@ -104,8 +113,8 @@ get_symm_buffer_size_for_mega_moe(
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},
{num_max_padded_sf_pool_tokens, hidden / 128},
{1, num_max_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)),
@@ -113,8 +122,8 @@ get_symm_buffer_size_for_mega_moe(
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},
{num_max_padded_sf_pool_tokens, intermediate_hidden / 128},
{1, num_max_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);
};
@@ -123,8 +132,9 @@ get_symm_buffer_size_for_mega_moe(
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 std::tuple<torch::Tensor, torch::Tensor>& l1_weights_tuple,
const std::tuple<torch::Tensor, torch::Tensor>& l2_weights_tuple,
const std::optional<torch::Tensor>& cumulative_local_expert_recv_stats,
const torch::Tensor& sym_buffer,
const std::vector<int64_t>& sym_buffer_ptrs, const int& rank_idx,
const int& num_max_tokens_per_rank,
@@ -132,9 +142,10 @@ static void fp8_fp4_mega_moe(
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_;
const bool& fast_math
) {
const auto [l1_weights, l1_weights_sf] = l1_weights_tuple;
const auto [l2_weights, l2_weights_sf] = l2_weights_tuple;
// Config checks
const auto num_tokens = static_cast<int>(y.size(0));
@@ -161,13 +172,20 @@ static void fp8_fp4_mega_moe(
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
// 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 stats counter
if (cumulative_local_expert_recv_stats.has_value()) {
DG_HOST_ASSERT(cumulative_local_expert_recv_stats->scalar_type() == torch::kInt);
DG_HOST_ASSERT(cumulative_local_expert_recv_stats->numel() == num_experts_per_rank);
DG_HOST_ASSERT(cumulative_local_expert_recv_stats->is_contiguous());
}
// Check buffer bytes
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts_ = num_experts_per_rank * num_ranks;
@@ -175,7 +193,7 @@ static void fp8_fp4_mega_moe(
num_ranks, num_experts,
num_max_tokens_per_rank, num_topk,
hidden, intermediate_hidden,
true, "swiglu");
true, activation);
DG_HOST_ASSERT(sym_buffer.nbytes() >= static_cast<size_t>(num_required_bytes));
DG_HOST_ASSERT(num_experts == num_experts_);
@@ -189,6 +207,7 @@ static void fp8_fp4_mega_moe(
l2_acts, l2_acts_sf,
l1_weights, l2_weights,
l1_weights_sf, l2_weights_sf,
cumulative_local_expert_recv_stats,
sym_buffer_ptrs,
rank_idx, num_max_tokens_per_rank,
num_experts_per_rank,
@@ -207,7 +226,7 @@ static void fp8_fp4_mega_moe(
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_token_alignment_for_mega_moe", &get_token_alignment_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

75
csrc/jit_kernels/heuristics/mega_moe.hpp
View File

@@ -55,38 +55,68 @@ struct MegaMoEConfig {
}
};
static int get_block_m_for_mega_moe(const int& num_ranks, const int& num_experts,
const int& num_max_tokens_per_rank, const int& num_topk) {
// TODO: compute based on configs
return 192;
static std::tuple<int, int, int, int> get_block_config_for_mega_moe(
const int& num_ranks, const int& num_experts,
const int& num_max_tokens_per_rank, const int& num_topk,
const int& num_tokens) {
const auto& [cluster_size, block_m, store_block_m, num_epilogue_warpgroups] = [&]() -> std::tuple<int, int, int, int> {
float num_expected_tokens_per_expert = static_cast<float>(num_tokens) * num_ranks * num_topk / num_experts;
if (num_expected_tokens_per_expert <= 8.5) {
// Really small token-per-expert (e.g. RL long-tail rollout), use the smallest block_m
return {2, 16, 8, 2};
} else if (num_expected_tokens_per_expert <= 16.5) {
// Small batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 128
return {2, 32, 16, 2};
} else if (num_expected_tokens_per_expert <= 32.5) {
// Medium batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 256
return {2, 64, 32, 1};
} else if (num_expected_tokens_per_expert <= 64.5) {
// Large batch size, small EP, decoding, e.g. 6/384 experts, EP8, bsz 512
return {2, 96, 16, 2};
} else if (num_expected_tokens_per_expert <= 96.5) {
// Medium batch size, Medium EP, decoding, e.g. 6/384 experts, EP16, bsz 256, or EP32, bsz128
return {2, 128, 32, 2};
} else {
// Prefill, or large EP decoding
return {2, 192, 32, 2};
}
}();
// Check whether our `block_m` lies in `kCandidateBlockM`
DG_HOST_ASSERT(std::any_of(
layout::kCandidateBlockM, layout::kCandidateBlockM + layout::kNumCandidateBlockMs,
[=](const auto& candidate) { return candidate == block_m; })
);
// Return configs
return {cluster_size, block_m, store_block_m, num_epilogue_warpgroups * 128};
}
static int get_num_experts_per_wave_for_mega_moe(
const int& num_experts_per_rank, const int& num_tokens, const int& num_topk,
const int& intermediate_hidden, const int& block_m, const int& block_n, const int& num_sms) {
float expected_tokens_per_expert = static_cast<float>(num_tokens) * num_topk / num_experts_per_rank;
if (expected_tokens_per_expert < 1) {
// Most experts don't have tokens, calculate all experts at once
return num_experts_per_rank;
}
// Reduce per-expert block count by this factor since uneven routing leaves some experts with fewer tokens
constexpr int kImbalanceFactor = 2;
// TODO: support num_experts_per_rank > 32
// Find the largest divisor of num_experts_per_rank that fits in 32 as the upper bound
int max_num_experts_per_wave = std::min(32, num_experts_per_rank);
while (max_num_experts_per_wave > 1 and num_experts_per_rank % max_num_experts_per_wave != 0)
-- max_num_experts_per_wave;
// Count L1 blocks per expert assuming tokens are evenly spread across experts
const int expected_tokens_per_expert =
num_tokens * num_topk / num_experts_per_rank + 1;
const int num_m_blocks = ceil_div(expected_tokens_per_expert, block_m);
const int num_n_blocks = intermediate_hidden / block_n;
const int num_m_blocks = ceil_div(static_cast<int>(std::ceil(expected_tokens_per_expert)), block_m);
const int num_n_blocks = (2 * intermediate_hidden) / block_n;
const int num_l1_blocks_per_expert = num_m_blocks * num_n_blocks;
// Pick the smallest value whose total blocks (after imbalance reduction) can keep all SMs busy
int num_experts_per_wave = num_l1_blocks_per_expert > 0
? ceil_div(kImbalanceFactor * num_sms, num_l1_blocks_per_expert) : 1;
num_experts_per_wave = std::min(num_experts_per_wave, max_num_experts_per_wave);
num_experts_per_wave = std::min(num_experts_per_wave, num_experts_per_rank);
// Round up to the nearest divisor of num_experts_per_rank so every wave processes the same count
while (num_experts_per_wave < max_num_experts_per_wave and num_experts_per_rank % num_experts_per_wave != 0)
while (num_experts_per_wave < num_experts_per_rank and num_experts_per_rank % num_experts_per_wave != 0)
++ num_experts_per_wave;
return num_experts_per_wave;
@@ -148,18 +178,18 @@ static std::pair<int, int> get_pipeline_config_for_mega_moe(
static MegaMoEConfig get_mega_moe_config(
const int& num_ranks, const int& num_experts, const int& num_experts_per_rank,
const int& num_max_tokens_per_rank, const int& num_tokens, const int& num_topk,
const int& hidden, const int& intermediate_hidden) {
// Block tiling
const int block_m = get_block_m_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk);
const int& hidden, const int& intermediate_hidden,
const int& num_padded_sf_pool_tokens) {
// Block config
const auto [cluster_size, block_m, store_block_m, num_epilogue_threads] =
get_block_config_for_mega_moe(num_ranks, num_experts, num_max_tokens_per_rank, num_topk, num_tokens);
const int block_n = 128;
const int block_k = 128;
const int load_block_m = block_m / 2;
const int load_block_n = block_n;
const int store_block_m = 32;
const auto [sf_block_m, sf_block_n] = SM100ArchSpec::get_sf_uttcp_aligned_block_sizes(block_m, block_n, MmaKind::MXFP8FP4);
const int num_max_pool_tokens = layout::get_num_max_pool_tokens(
num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank, block_m);
const int num_padded_sf_pool_tokens = layout::get_num_padded_sf_pool_tokens(num_max_pool_tokens, block_m);
num_ranks, num_max_tokens_per_rank, num_topk, num_experts_per_rank);
// NOTES: FP8 activations and FP4 weights (unpacked to 8-bit in smem) both use 128B swizzle
const int swizzle_acts_mode = 128;
const int swizzle_weights_mode = 128;
@@ -173,7 +203,6 @@ static MegaMoEConfig get_mega_moe_config(
// Thread layout
const int num_dispatch_threads = 128;
const int num_non_epilogue_threads = 128;
const int num_epilogue_threads = 256;
// Pipeline
const auto [num_stages, smem_size] = get_pipeline_config_for_mega_moe(

12
csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp
View File

@@ -29,6 +29,7 @@ public:
// Runtime arguments
void* y;
int* cumulative_local_expert_recv_stats;
int num_tokens;
layout::SymBuffer<> sym_buffer_ptrs;
@@ -91,6 +92,7 @@ static void __instantiate_kernel() {{
// TODO: optimize `args` copy
DG_CUDA_UNIFIED_CHECK(launch_kernel(kernel, config,
args.y,
args.cumulative_local_expert_recv_stats,
args.num_tokens,
args.sym_buffer_ptrs,
args.tensor_map_l1_acts,
@@ -112,6 +114,7 @@ static void sm100_fp8_fp4_mega_moe(
const torch::Tensor& l2_acts, const torch::Tensor& l2_acts_sf,
const torch::Tensor& l1_weights, const torch::Tensor& l2_weights,
const torch::Tensor& l1_weights_sf, const torch::Tensor& l2_weights_sf,
const std::optional<torch::Tensor> cumulative_local_expert_recv_stats,
const std::vector<int64_t>& sym_buffer_ptrs,
const int& rank_idx, const int& num_max_tokens_per_rank,
const int& num_experts_per_rank,
@@ -122,11 +125,12 @@ static void sm100_fp8_fp4_mega_moe(
) {
const auto num_ranks = static_cast<int>(sym_buffer_ptrs.size());
const auto num_experts = num_experts_per_rank * num_ranks;
const auto num_padded_sf_pool_tokens = static_cast<int>(l1_acts_sf.size(0));
// Heuristics
const auto config = get_mega_moe_config(
num_ranks, num_experts, num_experts_per_rank,
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden);
num_max_tokens_per_rank, num_tokens, num_topk, hidden, intermediate_hidden, num_padded_sf_pool_tokens);
// Make tensormap
constexpr int kGranK = 32;
@@ -175,6 +179,11 @@ static void sm100_fp8_fp4_mega_moe(
config.block_n, kGranK,
num_experts_per_rank, 0);
// Stats can be optional
int* cumulative_local_expert_recv_stats_ptr = nullptr;
if (cumulative_local_expert_recv_stats.has_value())
cumulative_local_expert_recv_stats_ptr = cumulative_local_expert_recv_stats->data_ptr<int>();
// Launch
const auto num_sms = device_runtime->get_num_sms();
const SM100FP8FP4MegaMoERuntime::Args args = {
@@ -186,6 +195,7 @@ static void sm100_fp8_fp4_mega_moe(
.fast_math = fast_math,
.config = config,
.y = y.data_ptr(),
.cumulative_local_expert_recv_stats = cumulative_local_expert_recv_stats_ptr,
.num_tokens = num_tokens,
.sym_buffer_ptrs = layout::SymBuffer<>(sym_buffer_ptrs, rank_idx),
.tensor_map_l1_acts = tensor_map_l1_acts,

48
csrc/jit_kernels/impls/smxx_fp8_fp4_paged_mqa_logits.hpp
View File

@@ -14,11 +14,13 @@ public:
int aligned_batch_size;
int split_kv;
int num_sms;
bool is_varlen;
int batch_size;
int next_n;
bool is_context_lens_2d;
int* context_lens;
int* indices;
int* schedule_metadata;
LaunchArgs launch_args;
@@ -32,10 +34,10 @@ using namespace deep_gemm;
static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sched::smxx_paged_mqa_logits_metadata<
{}, {}, {}
{}, {}, {}, {}
>);
}};
)", args.aligned_batch_size, args.split_kv, args.num_sms);
)", args.aligned_batch_size, args.split_kv, args.num_sms, args.is_varlen ? "true" : "false");
}
static void launch_impl(const KernelHandle& kernel, const LaunchConfigHandle& config, Args args) {
@@ -44,6 +46,7 @@ static void __instantiate_kernel() {{
args.next_n,
args.is_context_lens_2d,
args.context_lens,
args.indices,
args.schedule_metadata
));
}
@@ -53,14 +56,15 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens,
const torch::Tensor& schedule_metadata,
const int& batch_size, const int& next_n,
const int& block_kv, const int& num_sms,
const bool& is_context_lens_2d) {
const bool& is_context_lens_2d,
const bool& is_varlen, const int* indices_ptr) {
constexpr int split_kv = 256;
constexpr int num_threads = 32;
const int aligned_batch_size = align(batch_size, 32);
DG_HOST_ASSERT(split_kv % block_kv == 0);
// Calculate shared memory size
const int smem_size = aligned_batch_size * static_cast<int>(sizeof(int));
// Shared memory: prefix_sum[kAlignedBatchSize] + varlen_atom_token_start/context_len[kAlignedBatchSize] + varlen_num_atoms
const int smem_size = (3 * aligned_batch_size + 1) * static_cast<int>(sizeof(int));
DG_HOST_ASSERT(smem_size <= SM90ArchSpec::smem_capacity);
DG_HOST_ASSERT(smem_size <= SM100ArchSpec::smem_capacity);
@@ -69,10 +73,12 @@ static void smxx_paged_mqa_logits_metadata(const torch::Tensor& context_lens,
.aligned_batch_size = aligned_batch_size,
.split_kv = split_kv,
.num_sms = num_sms,
.is_varlen = is_varlen,
.batch_size = batch_size,
.next_n = next_n,
.is_context_lens_2d = is_context_lens_2d,
.context_lens = context_lens.data_ptr<int>(),
.indices = const_cast<int*>(indices_ptr),
.schedule_metadata = schedule_metadata.data_ptr<int>(),
.launch_args = LaunchArgs(1, num_threads, smem_size)
};
@@ -90,6 +96,7 @@ public:
int head_dim;
int block_kv;
bool is_context_lens_2d;
bool is_varlen;
int block_table_stride;
int logits_stride;
@@ -100,6 +107,7 @@ public:
int* context_lens;
void* logits;
int* block_table;
int* indices;
int* schedule_meta;
CUtensorMap tensor_map_q;
@@ -129,7 +137,7 @@ static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm{}_fp8_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{}, {},
{},
{}, {},
@@ -139,7 +147,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.is_context_lens_2d, args.is_varlen ? "true" : "false",
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads,
@@ -151,7 +159,7 @@ static void __instantiate_kernel() {{
args.batch_size,
args.logits_stride, args.block_table_stride,
args.context_lens, args.logits,
args.block_table, args.schedule_meta,
args.block_table, args.indices, args.schedule_meta,
args.tensor_map_q, args.tensor_map_kv,
args.tensor_map_kv_scales, args.tensor_map_weights
));
@@ -165,12 +173,14 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& context_lens,
const torch::Tensor& logits,
const torch::Tensor& block_table,
const torch::Tensor& indices,
const torch::Tensor& schedule_meta,
const at::ScalarType& logits_dtype,
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 bool& is_varlen,
const int& logits_stride,
const int& block_table_stride,
const int& num_sms,
@@ -183,7 +193,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
DG_HOST_ASSERT(split_kv % mma_m == 0 and logits_stride % split_kv == 0);
// Construct TMAs
const int next_n_atom = (next_n % 2 == 0) ? 2 : 1;
const int next_n_atom = (is_varlen or next_n >= 2) ? 2 : 1;
const auto tensor_map_q = make_tma_2d_desc(q, head_dim, batch_size * next_n * num_heads,
head_dim, next_n_atom * num_heads,
static_cast<int>(q.stride(2)),
@@ -245,6 +255,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.is_varlen = is_varlen,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
@@ -253,6 +264,7 @@ static void smxx_fp8_paged_mqa_logits(const torch::Tensor& q,
.context_lens = context_lens.data_ptr<int>(),
.logits = logits.data_ptr(),
.block_table = block_table.data_ptr<int>(),
.indices = is_varlen ? indices.data_ptr<int>() : nullptr,
.schedule_meta = schedule_meta.data_ptr<int>(),
.tensor_map_q = tensor_map_q,
.tensor_map_kv = tensor_map_kv,
@@ -279,6 +291,7 @@ public:
int head_dim;
int block_kv;
bool is_context_lens_2d;
bool is_varlen;
int block_table_stride;
int logits_stride;
@@ -289,6 +302,7 @@ public:
int* context_lens;
void* logits;
int* block_table;
int* indices;
int* schedule_meta;
CUtensorMap tensor_map_q;
@@ -314,7 +328,7 @@ static void __instantiate_kernel() {{
auto ptr = reinterpret_cast<void*>(&sm100_fp4_paged_mqa_logits<
{}, {},
{}, {},
{},
{}, {},
{}, {},
{},
{}, {},
@@ -323,7 +337,7 @@ static void __instantiate_kernel() {{
}};
)", args.next_n, args.num_heads,
args.head_dim, args.block_kv,
args.is_context_lens_2d,
args.is_context_lens_2d, args.is_varlen ? "true" : "false",
args.num_q_stages, args.num_kv_stages,
args.split_kv,
args.num_specialized_threads, args.num_math_threads,
@@ -335,7 +349,7 @@ static void __instantiate_kernel() {{
args.batch_size,
args.logits_stride, args.block_table_stride,
args.context_lens, args.logits,
args.block_table, args.schedule_meta,
args.block_table, args.indices, args.schedule_meta,
args.tensor_map_q, args.tensor_map_sf_q,
args.tensor_map_kv, args.tensor_map_sf_kv,
args.tensor_map_weights
@@ -351,12 +365,14 @@ static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q,
const torch::Tensor& context_lens,
const torch::Tensor& logits,
const torch::Tensor& block_table,
const torch::Tensor& indices,
const torch::Tensor& schedule_meta,
const at::ScalarType& logits_dtype,
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 bool& is_varlen,
const int& logits_stride,
const int& block_table_stride,
const int& num_sms,
@@ -366,8 +382,8 @@ static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q,
DG_HOST_ASSERT(split_kv == 256 and logits_stride % split_kv == 0);
// TODO: tuning num_stages
const int num_q_stages = 3, num_kv_stages = 6, num_tmem_stages = 3;
const int next_n_atom = (next_n % 2 == 0) ? 2 : 1;
const int num_q_stages = 3, num_kv_stages = 10, num_tmem_stages = 3;
const int next_n_atom = (is_varlen or next_n >= 2) ? 2 : 1;
// `head_dim` must be 128 for 64B swizzling
DG_HOST_ASSERT(head_dim == 128);
@@ -416,6 +432,7 @@ static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q,
.head_dim = head_dim,
.block_kv = block_kv,
.is_context_lens_2d = is_context_lens_2d,
.is_varlen = is_varlen,
.block_table_stride = block_table_stride,
.logits_stride = logits_stride,
.num_q_stages = num_q_stages,
@@ -424,6 +441,7 @@ static void sm100_fp4_paged_mqa_logits(const torch::Tensor& q,
.context_lens = context_lens.data_ptr<int>(),
.logits = logits.data_ptr(),
.block_table = block_table.data_ptr<int>(),
.indices = is_varlen ? indices.data_ptr<int>() : nullptr,
.schedule_meta = schedule_meta.data_ptr<int>(),
.tensor_map_q = tensor_map_q,
.tensor_map_sf_q = tensor_map_sf_q,