[PR #457] [Fix] Honor deterministic algorithms in SM90 bmk,bnk->mn einsum #462

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/457
Author: @LiRunGuo
Created: 9/22/2026
Status: 🔄 Open

Base: main ← Head: deterministic-bmk-bnk-mn

📝 Commits (2)

56d6e1c Honor deterministic algorithms in SM90 bmk,bnk->mn einsum
ad03662 Assert an FP32 D in the SM90 bmk,bnk->mn launcher

📊 Changes

3 files changed (+50 additions, -5 deletions)

View changed files

📝 csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp (+14 -4)
📝 deep_gemm/include/deep_gemm/impls/sm90_bmk_bnk_mn.cuh (+6 -1)
📝 tests/test_einsum.py (+30 -0)

📄 Description

Problem

deep_gemm.einsum('bmk,bnk->mn', ...) returns a different result on every call for the same inputs, even with deep_gemm.use_deterministic_algorithms(True).

On SM90, sm90_bmn_bnk_mn_gemm splits the s * k reduction across CTAs (split_factor) and adds each CTA's partial sum into D with float2 atomicAdd. The order of those additions changes between runs. With DeepGEMM main (78b6900) on an H200, 50 identical calls give:

Shape (s, m, n, k)	Default	`use_deterministic_algorithms(True)`
121, 128, 1024, 64	50 distinct outputs	50 distinct outputs
4096, 128, 384, 128	50 distinct outputs	49 distinct outputs

The flag is already honored by the other einsum paths, the BF16 GEMM, mega_mhc and mega_gate, but not by this one.

Fix (SM90)

In deterministic mode, and only when there is more than one split:

The launcher allocates a zeroed [num_splits, m, n] FP32 workspace. Its size is at most about num_sms * 128 * 128 floats, since num_splits * num_mn_blocks is bounded by the SM count.
A new kUseSplitWorkspace template flag makes each split accumulate into its own slice. Every element then has a single writer, so the atomic additions no longer depend on ordering.
The slices are summed into D in a fixed order with split_workspace.sum(0).

With the flag off, or with a single split, the generated kernel and launch are the same as before.

The SM100 kernel reduces the splits with TMA reduce-add in the same way, and I left it unchanged because I have no SM100 machine to test on. The same approach should apply there.

Testing (H200, CUDA 13.0, PyTorch 2.11.0+cu130)

The new test_bmk_bnk_mn_deterministic (SM90 only) fails on main and passes with this change. It checks 10 identical calls for bitwise equality and compares them with the reference.
Distinct outputs over 50 identical calls, with this change:

Shape (s, m, n, k)	Default	`use_deterministic_algorithms(True)`
121, 128, 1024, 64	50	1
4096, 128, 384, 128	50	1

python tests/test_einsum.py passes.
End-to-end time per call, in µs, including the workspace zeroing and the final reduction. The default mode is unchanged:

s, m, n, k (FP32 D)	`main`	this PR, default	this PR, deterministic
129, 128, 384, 128	9.1	9.1	23.9
129, 384, 128, 384	17.8	17.5	35.2
4096, 128, 384, 128	134.8	133.7	149.2
4096, 256, 256, 256	254.1	254.2	281.2
8192, 384, 128, 384	777.4	774.0	806.9

The deterministic mode costs a roughly fixed 12–15 µs for the extra zeroing, reduction and addition. That is about 2x on the smallest problems and 2–11% on the larger ones.

_{🔄 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/457 **Author:** [@LiRunGuo](https://github.com/LiRunGuo) **Created:** 9/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `deterministic-bmk-bnk-mn` --- ### 📝 Commits (2) - [`56d6e1c`](https://github.com/deepseek-ai/DeepGEMM/commit/56d6e1c95d2694f4b1ffaed325b4799f4059a603) Honor deterministic algorithms in SM90 bmk,bnk->mn einsum - [`ad03662`](https://github.com/deepseek-ai/DeepGEMM/commit/ad036627012183c8fe698ffe2b36418b313f435f) Assert an FP32 D in the SM90 bmk,bnk->mn launcher ### 📊 Changes **3 files changed** (+50 additions, -5 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/sm90_bmk_bnk_mn.hpp` (+14 -4) 📝 `deep_gemm/include/deep_gemm/impls/sm90_bmk_bnk_mn.cuh` (+6 -1) 📝 `tests/test_einsum.py` (+30 -0) </details> ### 📄 Description ## Problem `deep_gemm.einsum('bmk,bnk->mn', ...)` returns a different result on every call for the same inputs, even with `deep_gemm.use_deterministic_algorithms(True)`. On SM90, `sm90_bmn_bnk_mn_gemm` splits the `s * k` reduction across CTAs (`split_factor`) and adds each CTA's partial sum into D with `float2` `atomicAdd`. The order of those additions changes between runs. With DeepGEMM `main` (78b6900) on an H200, 50 identical calls give: | Shape (s, m, n, k) | Default | `use_deterministic_algorithms(True)` | |---|---|---| | 121, 128, 1024, 64 | 50 distinct outputs | 50 distinct outputs | | 4096, 128, 384, 128 | 50 distinct outputs | 49 distinct outputs | The flag is already honored by the other einsum paths, the BF16 GEMM, `mega_mhc` and `mega_gate`, but not by this one. ## Fix (SM90) In deterministic mode, and only when there is more than one split: - The launcher allocates a zeroed `[num_splits, m, n]` FP32 workspace. Its size is at most about `num_sms * 128 * 128` floats, since `num_splits * num_mn_blocks` is bounded by the SM count. - A new `kUseSplitWorkspace` template flag makes each split accumulate into its own slice. Every element then has a single writer, so the atomic additions no longer depend on ordering. - The slices are summed into D in a fixed order with `split_workspace.sum(0)`. With the flag off, or with a single split, the generated kernel and launch are the same as before. The SM100 kernel reduces the splits with TMA reduce-add in the same way, and I left it unchanged because I have no SM100 machine to test on. The same approach should apply there. ## Testing (H200, CUDA 13.0, PyTorch 2.11.0+cu130) - The new `test_bmk_bnk_mn_deterministic` (SM90 only) fails on `main` and passes with this change. It checks 10 identical calls for bitwise equality and compares them with the reference. - Distinct outputs over 50 identical calls, with this change: | Shape (s, m, n, k) | Default | `use_deterministic_algorithms(True)` | |---|---|---| | 121, 128, 1024, 64 | 50 | 1 | | 4096, 128, 384, 128 | 50 | 1 | - `python tests/test_einsum.py` passes. - End-to-end time per call, in µs, including the workspace zeroing and the final reduction. The default mode is unchanged: | s, m, n, k (FP32 D) | `main` | this PR, default | this PR, deterministic | |---|---|---|---| | 129, 128, 384, 128 | 9.1 | 9.1 | 23.9 | | 129, 384, 128, 384 | 17.8 | 17.5 | 35.2 | | 4096, 128, 384, 128 | 134.8 | 133.7 | 149.2 | | 4096, 256, 256, 256 | 254.1 | 254.2 | 281.2 | | 8192, 384, 128, 384 | 777.4 | 774.0 | 806.9 | The deterministic mode costs a roughly fixed 12–15 µs for the extra zeroing, reduction and addition. That is about 2x on the smallest problems and 2–11% on the larger ones. --- <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

bsakaguchi referenced this issue

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

[PR #463] Add SM90 BF16 grouped GEMM with gathered activation rows #466

bsakaguchi referenced this issue

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

[PR #462] [MERGED] Public release 26/09/30 #467