Skip to content

Merge upstream public release 26/09/30 (deepseek-ai/DeepGEMM#462) #30

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
zyongye merged 2 commits into from
Oct 10, 2026
Merged

Merge upstream public release 26/09/30 (deepseek-ai/DeepGEMM#462) #30

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
057ca59
Public release 26/09/30 (#462)
LyricZhao Sep 30, 2026
63aa62a
Merge upstream public release 26/09/30 (deepseek-ai/DeepGEMM#462)
zyongye Oct 10, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 17 additions & 9 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,9 @@ Despite its lightweight design, DeepGEMM's performance matches or exceeds expert

## News

- 2026.09.30
- DeepGEMM Ascend is available! Check [DeepGEMM-Ascend](https://github.com/deepseek-ai/DeepGEMM-Ascend/) for more details
- Add more optimizations, including locality domain features, check [#462](https://github.com/deepseek-ai/DeepGEMM/pull/462) for more details
- 2026.09.10: Sparse Indexer, Mega Gate, Mega mHC, DeepJIT, MoE and Indexer optimizations and more.
- Please see [#432](https://github.com/deepseek-ai/DeepGEMM/pull/432) for more details.
- 2026.04.16: Mega MoE, FP8xFP4 GEMM, FP4 Indexer, PDL, faster JIT compilation and more.
Expand Down Expand Up @@ -85,17 +88,17 @@ Use `m_grouped_fp8_gemm_nt_masked` for this purpose and consult the relevant doc
#### V3.2 MQA kernels for the indexer

The kernel family has two versions, non-paged (for prefilling) and paged (for decoding).
Take the non-paged version `fp8_mqa_logits` as an example. It has 6 inputs:
Take the non-paged version `fp8_fp4_mqa_logits` as an example. Its main inputs are:

- `q`, E4M3 tensor with shape `[seq_len, num_heads, head_dim]`
- `kv`, E4M3 tensor (shaped as `[seq_len_kv, head_dim]`) with float SF (shaped as `[seq_len_kv]`)
- `weights`, float tensor with shape `[seq_len, num_heads]`
- `q`, a `(q_data, q_sf)` tuple; SM100 accepts MXFP4/MXFP8 data with packed UE8M0 scales
- `kv`, a `(kv_data, kv_sf)` tuple with shape `[seq_len_kv, head_dim]` logically
- `weights`, tensor with shape `[seq_len, num_heads]` (BF16 on SM100)
- `cu_seq_len_k_start` and `cu_seq_len_k_end`, int tensor with shape `[seq_len]`
- `clean_logits`, whether to clean the unfilled logits into `-inf`
- `max_seqlen_k`, the maximum valid KV span of any query row

The output tensor is shaped as `[seq_len, seq_len_kv]`, indicating token-to-token logits.
The output is compressed to `[seq_len, max_seqlen_k]`; row `i` stores its valid KV span starting at column zero.
For each token `i` in `q`, it will iterate all tokens `j` from `[cu_seq_len_k_start[i], cu_seq_len_k_end[i])`,
and calculate the logit `out[i, j]` as:
and calculate the corresponding compressed logit as:

```python
kv_j = kv[0][j, :] * kv[1][j].unsqueeze(1) # [head_dim]
Expand All @@ -104,7 +107,7 @@ out_ij = out_ij.relu() * weights[i, :] # [num_heads]
out_ij = out_ij.sum() # Scalar
```

For more details and the paged version `fp8_paged_mqa_logits`, please refer to `tests/test_attention.py`.
For more details and the paged version `fp8_fp4_paged_mqa_logits`, please refer to `tests/test_attention.py`.

#### Mega MoE

Expand All @@ -118,9 +121,14 @@ buffer = deep_gemm.get_symm_buffer_for_mega_moe(
mma_type='fp8xfp4', # Use 'fp8xfp8' for FP8 routed-expert weights
)

# Transform routed weights (FP4 or FP8 with UE8M0 SF) into the required layout
# Transform weights (FP4 or FP8 with UE8M0 SF) into the required layout
transformed_l1, transformed_l2 = deep_gemm.transform_weights_for_mega_moe(l1_weights, l2_weights)

# (Optional) Localize weights into locality domains
transformed_l1 = (deep_gemm.localize(transformed_l1[0]), transformed_l1[1])
transformed_l2 = (deep_gemm.localize(transformed_l2[0]), transformed_l2[1])
deep_gemm.destroy_localizer()

# Copy inputs into the buffer before each call
# You may fuse these into previous kernels
buffer.x[:num_tokens].copy_(x_fp8)
Expand Down
725 changes: 263 additions & 462 deletions csrc/apis/attention.hpp
View file
Open in desktop

Large diffs are not rendered by default.

3 changes: 3 additions & 0 deletions csrc/apis/config.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10,6 +10,9 @@ TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
m.def("init(str library_root_path) -> ()", [&](const std::string& library_root_path) {
init_jit(library_root_path);
});
m.def("shutdown() -> ()", []() {
runtime = deep_jit::LazyInit<Runtime>(nullptr);
});
m.def("set_num_sms(int new_num_sms) -> ()", [&](const int64_t& new_num_sms) {
runtime->set_num_sms(new_num_sms);
});
Expand Down
120 changes: 85 additions & 35 deletions csrc/apis/einsum.hpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#include "../utils/exception.hpp"
#include "../utils/layout.hpp"
#include "../utils/compatibility.hpp"
#include "epilogue_class.hpp"
#include "gemm.hpp"

#include "../jit_kernels/impls/sm90_bmk_bnk_mn.hpp"
Expand Down Expand Up @@ -78,51 +79,79 @@ static void bmk_bnk_mn(const torch::Tensor& a, const torch::Tensor& b, const tor
}
}

static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D) {
static void bhr_hdr_bhd(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D,
const std::optional<torch::Tensor>& C,
const std::shared_ptr<EpilogueClass>& epilogue_class) {
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);
DG_HOST_ASSERT(B.scalar_type() == torch::kBFloat16 and B.stride(2) == 1);
DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1);
DG_HOST_ASSERT(D.stride(2) == 1);

// Dispatch implementation
// Early return for trivial cases
if (h == 0 or gemm::early_return(b, d, r, D, C))
return;

const auto resolved_epilogue_class = resolve_epilogue_class(epilogue_class, C, D);
const auto epilogue_alpha = resolved_epilogue_class->get_alpha();

// Dispatch implementation; cuBLASLt only covers a plain store, optionally alpha-scaled
const auto arch_major = jit->device.get_arch_major();
if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()) {
cublaslt_bhr_hdr_bhd(A, B, D, b, h, r, d);
if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()
and epilogue_alpha.has_value()) {
cublaslt_bhr_hdr_bhd(A, B, D, b, h, r, d, C.has_value(), epilogue_alpha.value());
} else if (arch_major == 9) {
sm90_bf16_bhr_hdr_bhd(A, B, D, b, h, r, d);
DG_HOST_ASSERT(epilogue_class == nullptr and "Custom epilogues require SM100");
sm90_bf16_bhr_hdr_bhd(A, B, C, D, b, h, r, d);
} else if (arch_major == 12) {
DG_HOST_ASSERT(epilogue_class == nullptr and "Custom epilogues require SM100");
DG_HOST_ASSERT(not C.has_value() and D.scalar_type() == torch::kBFloat16 and
"SM120 supports BF16 output without accumulation only");
sm120_bf16_bhr_hdr_bhd(A, B, D, b, h, r, d);
} else if (arch_major == 10) {
sm100_bf16_bhr_hdr_bhd(A, B, D, b, h, r, d);
sm100_bf16_bhr_hdr_bhd(A, B, C, D, b, h, r, d, "nk", resolved_epilogue_class);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
}

static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D) {
static void bhd_hdr_bhr(const torch::Tensor& A, const torch::Tensor& B, const torch::Tensor& D,
const std::optional<torch::Tensor>& C,
const std::shared_ptr<EpilogueClass>& epilogue_class) {
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);
DG_HOST_ASSERT(B.scalar_type() == torch::kBFloat16 and B.stride(2) == 1);
DG_HOST_ASSERT(D.scalar_type() == torch::kBFloat16 and D.stride(2) == 1);
DG_HOST_ASSERT(D.stride(2) == 1);

// Dispatch implementation
// Early return for trivial cases
if (h == 0 or gemm::early_return(b, r, d, D, C))
return;

const auto resolved_epilogue_class = resolve_epilogue_class(epilogue_class, C, D);
const auto epilogue_alpha = resolved_epilogue_class->get_alpha();

// Dispatch implementation; cuBLASLt only covers a plain store, optionally alpha-scaled
const auto arch_major = jit->device.get_arch_major();
if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()) {
cublaslt_bhd_hdr_bhr(A, B, D, b, h, r, d);
if (not heuristics_runtime->get_deterministic_algorithms() and runtime->is_cublaslt_available()
and epilogue_alpha.has_value()) {
cublaslt_bhd_hdr_bhr(A, B, D, b, h, r, d, C.has_value(), epilogue_alpha.value());
} else if (arch_major == 9) {
sm90_bf16_bhd_hdr_bhr(A, B, D, b, h, r, d);
DG_HOST_ASSERT(epilogue_class == nullptr and "Custom epilogues require SM100");
sm90_bf16_bhd_hdr_bhr(A, B, C, D, b, h, r, d);
} else if (arch_major == 12) {
DG_HOST_ASSERT(epilogue_class == nullptr and "Custom epilogues require SM100");
DG_HOST_ASSERT(not C.has_value() and D.scalar_type() == torch::kBFloat16 and
"SM120 supports BF16 output without accumulation only");
sm120_bf16_bhd_hdr_bhr(A, B, D, b, h, r, d);
} else if (arch_major == 10) {
sm100_bf16_bhd_hdr_bhr(A, B, D, b, h, r, d);
sm100_bf16_bhd_hdr_bhr(A, B, C, D, b, h, r, d, "nk", resolved_epilogue_class);
} else {
DG_HOST_UNREACHABLE("Unsupported architecture");
}
Expand Down Expand Up @@ -154,7 +183,8 @@ static void einsum(const std::string& expr,
const torch::Tensor& a,
const torch::Tensor& b,
const torch::Tensor& d,
const std::optional<torch::Tensor>& c) {
const std::optional<torch::Tensor>& c,
const std::shared_ptr<EpilogueClass>& epilogue_class) {
DG_HOST_ASSERT(a.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(b.scalar_type() == torch::kBFloat16);
DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16 or d.scalar_type() == torch::kFloat);
Expand All @@ -166,14 +196,14 @@ static void einsum(const std::string& expr,
// TODO: support any expression
// TODO: canonicalize expression
if (expr == "bmk,bnk->mn") {
DG_HOST_ASSERT(epilogue_class == nullptr and "Custom epilogues are unsupported for this expression");
bmk_bnk_mn(a, b, d, c);
} else if (expr == "bhr,hdr->bhd") {
DG_HOST_ASSERT(not c.has_value());
bhr_hdr_bhd(a, b, d);
bhr_hdr_bhd(a, b, d, c, epilogue_class);
} else if (expr == "bhd,hdr->bhr") {
DG_HOST_ASSERT(not c.has_value());
bhd_hdr_bhr(a, b, d);
bhd_hdr_bhr(a, b, d, c, epilogue_class);
} else if (expr == "bhd,bhr->hdr") {
DG_HOST_ASSERT(epilogue_class == nullptr and "Custom epilogues are unsupported for this expression");
bhd_bhr_hdr(a, b, d, c);
} else {
DG_HOST_UNREACHABLE(std::format("Unsupported einsum expression: {}", expr));
Expand All @@ -187,10 +217,20 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
const std::variant<torch::Tensor, std::pair<torch::Tensor, torch::Tensor>>& d,
const std::optional<torch::Tensor>& c,
std::optional<std::tuple<int, int, int>> recipe,
const std::string& compiled_dims) {
const std::string& compiled_dims,
const std::shared_ptr<EpilogueClass>& epilogue_class) {
// The FP8 output pair and a user `FP8Quantization` are two spellings of the same epilogue
const auto d_fp8 = std::get_if<std::pair<torch::Tensor, torch::Tensor>>(&d);
const auto d_tensor = d_fp8 != nullptr ? d_fp8->first : std::get<torch::Tensor>(d);
const auto sfd = d_fp8 != nullptr ? std::make_optional(d_fp8->second) : std::nullopt;
auto merged_epilogue_class = epilogue_class;
std::optional<torch::Tensor> sfd;
if (d_fp8 != nullptr) {
DG_HOST_ASSERT(epilogue_class == nullptr and "The FP8 output pair supports no other epilogue");
sfd = d_fp8->second;
merged_epilogue_class = std::make_shared<FP8QuantizationEpilogue>(d_fp8->second);
} else {
sfd = merged_epilogue_class != nullptr ? merged_epilogue_class->output_sfd() : std::nullopt;
}

// 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;
Expand All @@ -213,12 +253,9 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
// - `d.stride(0) == n`: the batches are contiguous along the columns of the view
// - `n % 128 == 0`: each batch is quantized independently, so neither an SF group
// nor a packed 4-SF word may cross a batch boundary
DG_HOST_ASSERT(jit->device.get_arch_major() == 10 and not c.has_value());
DG_HOST_ASSERT(d_tensor.scalar_type() == torch::kFloat8_e4m3fn);
DG_HOST_ASSERT(jit->device.get_arch_major() == 10);
DG_HOST_ASSERT(d_tensor.stride(0) == n and n % 128 == 0);
check_sf_layout(sfd.value(), m, batch_size * n, 1, 32, std::nullopt, true, false, torch::kInt);
} else {
DG_HOST_ASSERT(d_tensor.scalar_type() == torch::kBFloat16 or d_tensor.scalar_type() == torch::kFloat);
}

// Early return for trivial cases
Expand All @@ -239,11 +276,15 @@ static void fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
// Dispatch implementation
const auto arch_major = jit->device.get_arch_major();
if (arch_major == 12) {
DG_HOST_ASSERT(merged_epilogue_class == nullptr and "Custom epilogues require SM100");
sm120_fp8_fp4_bmm(a, transformed_sfa, b, transformed_sfb, c, d_tensor, batch_size, m, n, k,
gran_k_a, gran_k_b, major_a, major_b, compiled_dims);
} else if (arch_major == 10) {
sm100_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d_tensor, batch_size, m, n, k, gran_k_a, gran_k_b, major_a, major_b, compiled_dims, sfd);
sm100_fp8_bmm(a, transformed_sfa, b, transformed_sfb, c, d_tensor, batch_size, m, n, k,
gran_k_a, gran_k_b, major_a, major_b, compiled_dims,
resolve_epilogue_class(merged_epilogue_class, c, d_tensor));
} else {
DG_HOST_ASSERT(merged_epilogue_class == nullptr and "Custom epilogues require SM100");
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_tensor, batch_size, m, n, k, major_a, major_b, major_sfb, compiled_dims);
Expand All @@ -255,7 +296,8 @@ static void fp8_einsum(const std::string& expr,
const std::pair<torch::Tensor, torch::Tensor>& b,
const std::variant<torch::Tensor, std::pair<torch::Tensor, torch::Tensor>>& d,
const std::optional<torch::Tensor>& c,
const std::tuple<int, int, int>& recipe) {
const std::tuple<int, int, int>& recipe,
const std::shared_ptr<EpilogueClass>& epilogue_class) {
// Some hardcoded Einstein sum kernels
// NOTES: only `bhr,hdr->bhd` accepts an FP8 `(d, sfd)` output pair; the other
// expressions take a plain BF16/FP32 D
Expand All @@ -274,7 +316,7 @@ static void fp8_einsum(const std::string& expr,
else
std::get<torch::Tensor>(perm_d) = std::get<torch::Tensor>(perm_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");
fp8_bmm(perm_a, perm_sfa, b.first, b.second, perm_d, perm_c, recipe, "nk", epilogue_class);
} else if (expr == "bhd,hdr->bhr" and (arch_major == 10 or arch_major == 12)) {
// (batch_size, m, n, k): (h, b, r, d)
DG_HOST_ASSERT(std::holds_alternative<torch::Tensor>(d));
Expand All @@ -285,7 +327,7 @@ static void fp8_einsum(const std::string& expr,
const auto perm_sfb = b.second.permute({0, 2, 1});
const auto perm_d = std::get<torch::Tensor>(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");
fp8_bmm(perm_a, perm_sfa, perm_b, perm_sfb, perm_d, perm_c, recipe, "nk", epilogue_class);
} else if (expr == "bhd,bhr->hdr" and (arch_major == 10 or arch_major == 12)) {
// (batch_size, m, n, k): (h, d, r, b)
DG_HOST_ASSERT(std::holds_alternative<torch::Tensor>(d));
Expand All @@ -294,7 +336,7 @@ static void fp8_einsum(const std::string& expr,
const auto perm_sfa = a.second.permute({1, 2, 0});
const auto perm_b = arch_major == 12 ? b.first.permute({1, 2, 0}).contiguous() : 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");
fp8_bmm(perm_a, perm_sfa, perm_b, perm_sfb, d, c, recipe, "mn", epilogue_class);
} else {
DG_HOST_UNREACHABLE(std::format("Unsupported einsum expression: {}", expr));
}
Expand All @@ -304,28 +346,36 @@ static void fp8_einsum(const std::string& expr,
namespace deep_gemm::torch_registration {
using namespace deep_gemm::torch_utils;

static void einsum(const std::string& expr,
const torch::Tensor& a, const torch::Tensor& b,
const torch::Tensor& d, const c10::optional<torch::Tensor>& c,
const epilogue_class::EpilogueArg& epilogue) {
einsum::einsum(expr, a, b, d, c, epilogue_class::unwrap(epilogue));
}

static void fp8_einsum(const std::string& expr,
const torch::Tensor& a, const torch::Tensor& sfa,
const torch::Tensor& b, const torch::Tensor& sfb,
const torch::Tensor& d, const c10::optional<torch::Tensor>& c,
const std::vector<int64_t>& recipe,
const std::optional<torch::Tensor>& sfd) {
const std::optional<torch::Tensor>& sfd,
const epilogue_class::EpilogueArg& epilogue) {
std::variant<torch::Tensor, std::pair<torch::Tensor, torch::Tensor>> output = d;
if (sfd.has_value()) output = std::make_pair(d, *sfd);
einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, output, c, list_to_tuple3(recipe));
einsum::fp8_einsum(expr, {a, sfa}, {b, sfb}, output, c, list_to_tuple3(recipe), epilogue_class::unwrap(epilogue));
}
} // namespace deep_gemm::torch_registration

TORCH_LIBRARY_FRAGMENT(deep_gemm, m) {
m.def(
"einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None) -> ()");
"einsum(str expr, Tensor a, Tensor b, Tensor(d!) d, Tensor? c=None, __torch__.torch.classes.deep_gemm.Epilogue? epilogue=None) -> ()");
m.def(
"fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3] recipe, Tensor(sfd!)? sfd=None) -> ()");
"fp8_einsum(str expr, Tensor a, Tensor sfa, Tensor b, Tensor sfb, Tensor(d!) d, Tensor? c=None, int[3] recipe, Tensor(sfd!)? sfd=None, __torch__.torch.classes.deep_gemm.Epilogue? epilogue=None) -> ()");
}

TORCH_LIBRARY_IMPL(deep_gemm, CUDA, m) {
using namespace deep_gemm::torch_registration;

m.impl("einsum", TORCH_FN(deep_gemm::einsum::einsum));
m.impl("einsum", TORCH_FN(einsum));
m.impl("fp8_einsum", TORCH_FN(fp8_einsum));
}
Loading