[PR #477] [Fix][Einsum] Reject `c` in SM90 `fp8_einsum("bhr,hdr->bhd")` to prevent silent data corruption in `sm90_fp8_bmm` #480

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/477
Author: @arxgy
Created: 10/17/2026
Status: 🔄 Open

Base: main ← Head: fix/sm90-fp8-einsum-c

📝 Commits (1)

7035f91 Reject C in SM90 FP8 BMM instead of silently dropping it

📊 Changes

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

View changed files

📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp (+1 -1)
📝 tests/test_einsum.py (+9 -0)

📄 Description

This PR fixes #476 .

Summary

On SM90, deep_gemm.fp8_einsum("bhr,hdr->bhd", a, b, d, c) returns a @ b.T instead of c + a @ b.T: c is silently dropped.

The host first copies c into d, expecting the kernel to add a @ b.T onto d (gemm.hpp L49–L51).
On SM90, sm90_fp8_bmm launches sm90_fp8_gemm_1d2d_impl, which cannot add onto d. It writes d with a plain TMA store (sm90_fp8_gemm_1d2d.cuh L434), which overwrites c. sm90_fp8_bmm sets with_accumulation (L300), but never passes it to the kernel (L36–L61).
The other launcher of this kernel, sm90_fp8_gemm_1d2d, rejects c (L81). sm90_fp8_bmm does not (L290).

"bhr,hdr->bhd" is the only fp8_einsum expression that runs on SM90, and tests/test_einsum.py passes c to it only on SM100, which accumulates correctly.

Reproduction: See Issue #476 .

Fix

Add the missing not c.has_value() check to sm90_fp8_bmm, as in sm90_fp8_gemm_1d2d. Passing c on SM90 now raises an error instead of returning a wrong result. As with sm90_fp8_gemm_1d2d, d has already been overwritten with c when the error is raised.

--- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
+++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp
@@ -287,7 +287,7 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
                          const int& batch_size, const int& m, const int& n, const int& k,
                          const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb,
                          const std::string& compiled_dims) {
-    DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16);
+    DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
     DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K);
 
     const auto desc = GemmDesc {

Tests

test_fp8_bhr_hdr_bhd already checks c on SM100. This PR adds the SM90 counterpart: a call with c must raise. The accumulator has the layout of the SM100 check ([h, b, d] viewed as [b, h, d]), in BF16, the only d type sm90_fp8_bmm accepts.

                         assert_stochastic_bf16_matches_fp32_accumulation(
                             z_sr, launch_accumulated, f'{b=}, {h=}, {r=}, {d=}')
+                    else:
+                        # SM90's FP8 BMM kernel cannot add C into D, so C must be rejected
+                        acc = torch.zeros((h, b, d), device='cuda', dtype=torch.bfloat16).permute(1, 0, 2)
+                        try:
+                            deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, acc, c=acc)
+                        except RuntimeError:
+                            pass
+                        else:
+                            raise AssertionError(f'{b=}, {h=}, {r=}, {d=}: C must be rejected on SM90')

Tested on an H100 80GB HBM3 (driver 580.126.09, CUDA 13.0, torch 2.9.1+cu130), on main (1fc9c66) and on this branch, both built with ./develop.sh. The test functions run in the order of __main__, after one warm-up bench_kineto call. Without the warm-up, python tests/test_einsum.py stops at its first benchmark with a ZeroDivisionError on both trees, because bench_kineto returns 0 on its first call in a process (#426).

main with the new test: the call with c does not raise, so the test fails.

Testing "bmk, bnk -> mn":
Testing "bhr, hdr -> bhd":
Testing "bhd, hdr -> bhr":
Testing "bhd, bhr -> hdr":
Testing FP8 "bhr, hdr -> bhd":
...
AssertionError: b=4, h=8, r=4096, d=1024: C must be rejected on SM90

This PR:

Testing "bmk, bnk -> mn":
Testing "bhr, hdr -> bhd":
Testing "bhd, hdr -> bhr":
Testing "bhd, bhr -> hdr":
Testing FP8 "bhr, hdr -> bhd":
test_fp8_bhd_hdr_bhr:
 > Filtered by <function <lambda> at 0x73748c055760>
test_fp8_bhd_bhr_hdr:
 > Filtered by <function <lambda> at 0x73748c055800>
ALL TESTS PASSED

_{🔄 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/477 **Author:** [@arxgy](https://github.com/arxgy) **Created:** 10/17/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/sm90-fp8-einsum-c` --- ### 📝 Commits (1) - [`7035f91`](https://github.com/deepseek-ai/DeepGEMM/commit/7035f910f76393e21762768ee5381166e92c3ace) Reject C in SM90 FP8 BMM instead of silently dropping it ### 📊 Changes **2 files changed** (+10 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp` (+1 -1) 📝 `tests/test_einsum.py` (+9 -0) </details> ### 📄 Description This PR fixes #476 . ## Summary On SM90, `deep_gemm.fp8_einsum("bhr,hdr->bhd", a, b, d, c)` returns `a @ b.T` instead of `c + a @ b.T`: `c` is silently dropped. - The host first copies `c` into `d`, expecting the kernel to add `a @ b.T` onto `d` ([gemm.hpp L49–L51](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/apis/gemm.hpp#L49-L51)). - On SM90, `sm90_fp8_bmm` launches `sm90_fp8_gemm_1d2d_impl`, which cannot add onto `d`. It writes `d` with a plain TMA store ([sm90_fp8_gemm_1d2d.cuh L434](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh#L434)), which overwrites `c`. `sm90_fp8_bmm` sets `with_accumulation` ([L300](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L300)), but never passes it to the kernel ([L36–L61](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L36-L61)). - The other launcher of this kernel, `sm90_fp8_gemm_1d2d`, rejects `c` ([L81](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L81)). `sm90_fp8_bmm` does not ([L290](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L290)). `"bhr,hdr->bhd"` is the only `fp8_einsum` expression that runs on SM90, and `tests/test_einsum.py` passes `c` to it only on SM100, which accumulates correctly. Reproduction: See Issue #476 . ## Fix Add the missing `not c.has_value()` check to `sm90_fp8_bmm`, as in `sm90_fp8_gemm_1d2d`. Passing `c` on SM90 now raises an error instead of returning a wrong result. As with `sm90_fp8_gemm_1d2d`, `d` has already been overwritten with `c` when the error is raised. ```diff --- a/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp +++ b/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp @@ -287,7 +287,7 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, const int& batch_size, const int& m, const int& n, const int& k, const cute::UMMA::Major& major_a, const cute::UMMA::Major& major_b, const cute::UMMA::Major& major_sfb, const std::string& compiled_dims) { - DG_HOST_ASSERT(d.scalar_type() == torch::kBFloat16); + DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16); DG_HOST_ASSERT(major_a == cute::UMMA::Major::K and major_b == cute::UMMA::Major::K); const auto desc = GemmDesc { ``` ## Tests `test_fp8_bhr_hdr_bhd` already checks `c` on SM100. This PR adds the SM90 counterpart: a call with `c` must raise. The accumulator has the layout of the SM100 check (`[h, b, d]` viewed as `[b, h, d]`), in BF16, the only `d` type `sm90_fp8_bmm` accepts. ```diff assert_stochastic_bf16_matches_fp32_accumulation( z_sr, launch_accumulated, f'{b=}, {h=}, {r=}, {d=}') + else: + # SM90's FP8 BMM kernel cannot add C into D, so C must be rejected + acc = torch.zeros((h, b, d), device='cuda', dtype=torch.bfloat16).permute(1, 0, 2) + try: + deep_gemm.fp8_einsum('bhr,hdr->bhd', x_fp8, y_fp8, acc, c=acc) + except RuntimeError: + pass + else: + raise AssertionError(f'{b=}, {h=}, {r=}, {d=}: C must be rejected on SM90') ``` Tested on an H100 80GB HBM3 (driver 580.126.09, CUDA 13.0, torch 2.9.1+cu130), on `main` (1fc9c66) and on this branch, both built with `./develop.sh`. The test functions run in the order of `__main__`, after one warm-up `bench_kineto` call. Without the warm-up, `python tests/test_einsum.py` stops at its first benchmark with a `ZeroDivisionError` on both trees, because `bench_kineto` returns 0 on its first call in a process (#426). `main` with the new test: the call with `c` does not raise, so the test fails. ``` Testing "bmk, bnk -> mn": Testing "bhr, hdr -> bhd": Testing "bhd, hdr -> bhr": Testing "bhd, bhr -> hdr": Testing FP8 "bhr, hdr -> bhd": ... AssertionError: b=4, h=8, r=4096, d=1024: C must be rejected on SM90 ``` This PR: ``` Testing "bmk, bnk -> mn": Testing "bhr, hdr -> bhd": Testing "bhd, hdr -> bhr": Testing "bhd, bhr -> hdr": Testing FP8 "bhr, hdr -> bhd": test_fp8_bhd_hdr_bhr: > Filtered by <function <lambda> at 0x73748c055760> test_fp8_bhd_bhr_hdr: > Filtered by <function <lambda> at 0x73748c055800> ALL TESTS PASSED ``` --- <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:21 +00:00

bsakaguchi referenced this issue

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

[PR #480] Use the allocated TMEM base in the SM100 HC prenorm GEMM #482