deepseek-ai/DeepGEMM

Fork 0

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

[GH-ISSUE #476] [Bug] SM90 `fp8_einsum("bhr,hdr->bhd", a, b, d, c)` silently drops `c` and returns `a @ b.T` #199

New issue

Open

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

bsakaguchi commented

2026-10-26 21:14:39 +00:00

Owner

Originally created by @arxgy on GitHub (Oct 17, 2026).
Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/476

Description

On SM90, deep_gemm.fp8_einsum("bhr,hdr->bhd", a, b, d, c=c) returns d = a \cdot b^\top instead of d = c + a \cdot b^\top. No error is raised, and the output looks like a normal result. The same happens when d and c are the same tensor (in-place accumulation). On SM100 the result is correct.

The README's convention is D = C + A @ B ("fp8_gemm_nt will do a D = C + A @ B.T"), and fp8_einsum accepts c on every architecture. "bhr,hdr->bhd" is the only fp8_einsum expression that runs on SM90, so any caller that accumulates through it on Hopper gets a wrong result.

Environment

H100 80GB HBM3 (sm_90), driver 580.126.09, CUDA 13.0 (nvcc V13.0.88), torch 2.9.1+cu130.
DeepGEMM main 1fc9c66, CUTLASS f3fde58, built with ./develop.sh.

Root Cause

The host prepares d for an accumulating kernel, but on SM90 it launches a kernel that only stores:

fp8_einsum passes c to fp8_bmm (einsum.hpp L276), which calls gemm::early_return (L230). It copies c into d, expecting the kernel to add into d (gemm.hpp L49–L51):

// With accumulation, do copy before GEMM (assuming the GEMM kernel does not support different C/D)
if (c.has_value() and not is_cd_same)
    d.copy_(c.value());

On SM90, fp8_bmm calls sm90_fp8_bmm (einsum.hpp L247), which launches sm90_fp8_gemm_1d2d_impl. It sets with_accumulation = c.has_value() in the GemmDesc (sm90_fp8_gemm_1d2d.hpp L300), but the flag is not a template argument of the kernel (L36–L61), and the kernel writes D with a plain TMA store (sm90_fp8_gemm_1d2d.cuh L434):
```
if constexpr (kGemmType == GemmType::Batched) {
    cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr,
                                  n_idx, m_idx, scheduler.current_group_idx);
```
This overwrites the c that step 1 copied into d. For comparison, sm90_bf16_gemm passes with_accumulation to its kernel as kWithAccumulation and switches the store to a reduce-add.
The other launcher of the same kernel, sm90_fp8_gemm_1d2d (used by fp8_gemm_nt), rejects c (sm90_fp8_gemm_1d2d.hpp L81):
```
DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16);
```
sm90_fp8_bmm has no such check.

tests/test_einsum.py passes c to fp8_einsum only on SM100 (L183), so the SM90 path with c is untested.

Reproduction

repro.py:

import torch
import deep_gemm
from deep_gemm.utils.math import per_block_cast_to_fp8, per_token_cast_to_fp8

torch.manual_seed(0)
b, h, r, d = 128, 8, 4096, 1024
x = torch.randn(b, h, r, device='cuda', dtype=torch.bfloat16)
y = torch.randn(h, d, r, device='cuda', dtype=torch.bfloat16)
x_fp8, x_sf = per_token_cast_to_fp8(x.view(-1, r), use_ue8m0=False)
y_fp8, y_sf = zip(*(per_block_cast_to_fp8(y[i], use_ue8m0=False) for i in range(h)))
a_ = (x_fp8.view(b, h, r), x_sf.view(b, h, r // 128))
b_ = (torch.stack(y_fp8), torch.stack(y_sf))

ab = torch.einsum('bhr,hdr->bhd', x.float(), y.float())
# as large as a @ b.T, in the layout fp8_einsum requires for C: [h, b, d] viewed as [b, h, d]
c = (ab.std() * torch.randn(h, b, d, device='cuda')).bfloat16().permute(1, 0, 2)
d_out = torch.empty_like(c)
deep_gemm.fp8_einsum('bhr,hdr->bhd', a_, b_, d_out, c=c)

rel_err = lambda ref: ((d_out.float() - ref).norm() / ref.norm()).item()
print(f'relative error vs c + a @ b.T: {rel_err(c.float() + ab):.3f}')   # expected: small
print(f'relative error vs a @ b.T:     {rel_err(ab):.3f}')               # actual: small

Expected vs Actual

Expected: d_out ≈ c + a @ b.T, with a relative error at the FP8 level. If SM90 does not support c here, a host assertion is expected, as in sm90_fp8_gemm_1d2d.

Actual (python repro.py on the H100 above): d_out is a @ b.T up to the FP8 error, and c is missing.

relative error vs c + a @ b.T: 0.707
relative error vs a @ b.T:     0.037

Suggested Fix

Either of the following removes the silent wrong result. Which one fits the intended SM90 support?

Reject c in sm90_fp8_bmm, as sm90_fp8_gemm_1d2d does:

@@ -287,7 +287,7 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa,
                          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);

With this change, repro.py raises RuntimeError: Assertion error (csrc/apis/../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp:290): not c.has_value() and d.scalar_type() == torch::kBFloat16 on the same H100. As in sm90_fp8_gemm_1d2d, the check runs after early_return has copied c into d.

Support c on SM90: add an accumulation flag to sm90_fp8_gemm_1d2d_impl that writes D with a TMA reduce-add, as sm90_bf16_gemm does with kWithAccumulation, and pass with_accumulation to it from sm90_fp8_bmm.

I can open a PR for either option.

Originally created by @arxgy on GitHub (Oct 17, 2026). Original GitHub issue: https://github.com/deepseek-ai/DeepGEMM/issues/476 ## Description On SM90, `deep_gemm.fp8_einsum("bhr,hdr->bhd", a, b, d, c=c)` returns $d = a \cdot b^\top$ instead of $d = c + a \cdot b^\top$. No error is raised, and the output looks like a normal result. The same happens when `d` and `c` are the same tensor (in-place accumulation). On SM100 the result is correct. [The README's convention](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/README.md?plain=1#L63) is `D = C + A @ B` ("`fp8_gemm_nt` will do a `D = C + A @ B.T`"), and `fp8_einsum` accepts `c` on every architecture. `"bhr,hdr->bhd"` is the only `fp8_einsum` expression that runs on SM90, so any caller that accumulates through it on Hopper gets a wrong result. ## Environment - H100 80GB HBM3 (sm_90), driver 580.126.09, CUDA 13.0 (nvcc V13.0.88), torch 2.9.1+cu130. - DeepGEMM `main` 1fc9c66, CUTLASS f3fde58, built with `./develop.sh`. ## Root Cause The host prepares `d` for an accumulating kernel, but on SM90 it launches a kernel that only stores: 1. `fp8_einsum` passes `c` to `fp8_bmm` ([einsum.hpp L276](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/apis/einsum.hpp#L276)), which calls `gemm::early_return` ([L230](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/apis/einsum.hpp#L230)). It copies `c` into `d`, expecting the kernel to add into `d` ([gemm.hpp L49–L51](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/apis/gemm.hpp#L49-L51)): ```cpp // With accumulation, do copy before GEMM (assuming the GEMM kernel does not support different C/D) if (c.has_value() and not is_cd_same) d.copy_(c.value()); ``` 2. On SM90, `fp8_bmm` calls `sm90_fp8_bmm` ([einsum.hpp L247](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/apis/einsum.hpp#L247)), which launches `sm90_fp8_gemm_1d2d_impl`. It sets `with_accumulation = c.has_value()` in the `GemmDesc` ([sm90_fp8_gemm_1d2d.hpp L300](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L300)), but the flag is not a template argument of the kernel ([L36–L61](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L36-L61)), and the kernel 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)): ```cpp if constexpr (kGemmType == GemmType::Batched) { cute::SM90_TMA_STORE_3D::copy(&tensor_map_d, smem_ptr, n_idx, m_idx, scheduler.current_group_idx); ``` This overwrites the `c` that step 1 copied into `d`. For comparison, `sm90_bf16_gemm` passes `with_accumulation` to its kernel as `kWithAccumulation` and switches the store to a reduce-add. 3. The other launcher of the same kernel, `sm90_fp8_gemm_1d2d` (used by `fp8_gemm_nt`), rejects `c` ([sm90_fp8_gemm_1d2d.hpp L81](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp#L81)): ```cpp DG_HOST_ASSERT(not c.has_value() and d.scalar_type() == torch::kBFloat16); ``` `sm90_fp8_bmm` has no such check. `tests/test_einsum.py` passes `c` to `fp8_einsum` only on SM100 ([L183](https://github.com/deepseek-ai/DeepGEMM/blob/1fc9c66975ee322d628b3288534ca0e9ce12a64d/tests/test_einsum.py#L183)), so the SM90 path with `c` is untested. ## Reproduction `repro.py`: ```python import torch import deep_gemm from deep_gemm.utils.math import per_block_cast_to_fp8, per_token_cast_to_fp8 torch.manual_seed(0) b, h, r, d = 128, 8, 4096, 1024 x = torch.randn(b, h, r, device='cuda', dtype=torch.bfloat16) y = torch.randn(h, d, r, device='cuda', dtype=torch.bfloat16) x_fp8, x_sf = per_token_cast_to_fp8(x.view(-1, r), use_ue8m0=False) y_fp8, y_sf = zip(*(per_block_cast_to_fp8(y[i], use_ue8m0=False) for i in range(h))) a_ = (x_fp8.view(b, h, r), x_sf.view(b, h, r // 128)) b_ = (torch.stack(y_fp8), torch.stack(y_sf)) ab = torch.einsum('bhr,hdr->bhd', x.float(), y.float()) # as large as a @ b.T, in the layout fp8_einsum requires for C: [h, b, d] viewed as [b, h, d] c = (ab.std() * torch.randn(h, b, d, device='cuda')).bfloat16().permute(1, 0, 2) d_out = torch.empty_like(c) deep_gemm.fp8_einsum('bhr,hdr->bhd', a_, b_, d_out, c=c) rel_err = lambda ref: ((d_out.float() - ref).norm() / ref.norm()).item() print(f'relative error vs c + a @ b.T: {rel_err(c.float() + ab):.3f}') # expected: small print(f'relative error vs a @ b.T: {rel_err(ab):.3f}') # actual: small ``` ## Expected vs Actual Expected: `d_out` ≈ `c + a @ b.T`, with a relative error at the FP8 level. If SM90 does not support `c` here, a host assertion is expected, as in `sm90_fp8_gemm_1d2d`. Actual (`python repro.py` on the H100 above): `d_out` is `a @ b.T` up to the FP8 error, and `c` is missing. ``` relative error vs c + a @ b.T: 0.707 relative error vs a @ b.T: 0.037 ``` ## Suggested Fix Either of the following removes the silent wrong result. Which one fits the intended SM90 support? 1. **Reject `c` in `sm90_fp8_bmm`**, as `sm90_fp8_gemm_1d2d` does: ```diff @@ -287,7 +287,7 @@ static void sm90_fp8_bmm(const torch::Tensor& a, const torch::Tensor& sfa, 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); ``` With this change, `repro.py` raises `RuntimeError: Assertion error (csrc/apis/../jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp:290): not c.has_value() and d.scalar_type() == torch::kBFloat16` on the same H100. As in `sm90_fp8_gemm_1d2d`, the check runs after `early_return` has copied `c` into `d`. 2. **Support `c` on SM90**: add an accumulation flag to `sm90_fp8_gemm_1d2d_impl` that writes `D` with a TMA reduce-add, as `sm90_bf16_gemm` does with `kWithAccumulation`, and pass `with_accumulation` to it from `sm90_fp8_bmm`. I can open a PR for either option.

bsakaguchi referenced this issue

2026-10-26 21:16:24 +00:00

[PR #272] [CLOSED] Add FP8 per-tensor scaling to cuBLASLt for fair comparison #334

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#199

No description provided.

Rows
Columns

[GH-ISSUE #476] [Bug] SM90 fp8_einsum("bhr,hdr->bhd", a, b, d, c) silently drops c and returns a @ b.T #199