deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #456] [Fix] Reject FP32 bmk,bnk->mn einsum without C instead of crashing #461

New issue

Open

opened 2026-10-26 21:18:06 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:06 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/456
Author: @LiRunGuo
Created: 9/22/2026
Status: 🔄 Open

Base: main ← Head: fix-einsum-fp32-without-c

📝 Commits (1)

c139db6 Reject FP32 bmk,bnk->mn einsum without C instead of crashing

📊 Changes

2 files changed (+18 additions, -1 deletions)

View changed files

📝 csrc/apis/einsum.hpp (+1 -1)
📝 tests/test_einsum.py (+17 -0)

📄 Description

Problem

deep_gemm.einsum('bmk,bnk->mn', a, b, d) with an FP32 d and no c crashes the whole Python process with a segmentation fault, instead of raising an error.

For an FP32 D, bmk_bnk_mn only supports the accumulated expression, so c must be d itself. The check that enforces this dereferences the optional c without testing it first:

// csrc/apis/einsum.hpp
if (d.scalar_type() == torch::kFloat) {
    DG_HOST_ASSERT(c->data_ptr() == d.data_ptr() and ...);   // `c` may be empty here

When c is None, this is undefined behavior on an empty std::optional. On an H200 with DeepGEMM main (78b6900) the process exits with SIGSEGV:

Thread 1 "python" received signal SIGSEGV, Segmentation fault.
#1  c10::TensorImpl::data_impl<...>
#2  c10::TensorImpl::mutable_data
#5  deep_gemm::einsum::bmk_bnk_mn (a=..., b=..., d=..., c=std::optional [no contained value]) at csrc/apis/einsum.hpp:26

The check runs before architecture dispatch, so SM100 is affected as well. I found this with a randomized shape fuzzer that calls the public einsum API.

Fix

Check c.has_value() first, so the call fails the host assertion like any other unsupported argument. All other -> dereferences of optional tensors under csrc/apis/ are already guarded by a has_value() check.

Test

tests/test_einsum.py gets test_bmk_bnk_mn_fp32_requires_c, which expects a RuntimeError for an FP32 d without c.

Results on an H200 (CUDA 13.0, PyTorch 2.11.0+cu130):

Build	`test_bmk_bnk_mn_fp32_requires_c`
`main` (`78b6900`)	process killed by SIGSEGV (exit code 139)
this PR	`RuntimeError` from the host assertion, test passes

The full python tests/test_einsum.py passes on this branch.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/456 **Author:** [@LiRunGuo](https://github.com/LiRunGuo) **Created:** 9/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix-einsum-fp32-without-c` --- ### 📝 Commits (1) - [`c139db6`](https://github.com/deepseek-ai/DeepGEMM/commit/c139db6fd8518b358ff628d31f2a2aeb12dee98b) Reject FP32 bmk,bnk->mn einsum without C instead of crashing ### 📊 Changes **2 files changed** (+18 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/einsum.hpp` (+1 -1) 📝 `tests/test_einsum.py` (+17 -0) </details> ### 📄 Description ## Problem `deep_gemm.einsum('bmk,bnk->mn', a, b, d)` with an FP32 `d` and no `c` crashes the whole Python process with a segmentation fault, instead of raising an error. For an FP32 D, `bmk_bnk_mn` only supports the accumulated expression, so `c` must be `d` itself. The check that enforces this dereferences the optional `c` without testing it first: ```cpp // csrc/apis/einsum.hpp if (d.scalar_type() == torch::kFloat) { DG_HOST_ASSERT(c->data_ptr() == d.data_ptr() and ...); // `c` may be empty here ``` When `c` is `None`, this is undefined behavior on an empty `std::optional`. On an H200 with DeepGEMM `main` (78b6900) the process exits with SIGSEGV: ``` Thread 1 "python" received signal SIGSEGV, Segmentation fault. #1 c10::TensorImpl::data_impl<...> #2 c10::TensorImpl::mutable_data #5 deep_gemm::einsum::bmk_bnk_mn (a=..., b=..., d=..., c=std::optional [no contained value]) at csrc/apis/einsum.hpp:26 ``` The check runs before architecture dispatch, so SM100 is affected as well. I found this with a randomized shape fuzzer that calls the public einsum API. ## Fix Check `c.has_value()` first, so the call fails the host assertion like any other unsupported argument. All other `->` dereferences of optional tensors under `csrc/apis/` are already guarded by a `has_value()` check. ## Test `tests/test_einsum.py` gets `test_bmk_bnk_mn_fp32_requires_c`, which expects a `RuntimeError` for an FP32 `d` without `c`. Results on an H200 (CUDA 13.0, PyTorch 2.11.0+cu130): | Build | `test_bmk_bnk_mn_fp32_requires_c` | |---|---| | `main` (78b6900) | process killed by SIGSEGV (exit code 139) | | this PR | `RuntimeError` from the host assertion, test passes | The full `python tests/test_einsum.py` passes on this branch. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:18:06 +00:00

No labels

pull-request

No milestone

No project

No assignees

1 participant

Notifications

Due date

The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference

deepseek-ai/DeepGEMM#461

No description provided.

Rows
Columns