[PR #469] [CLOSED] Fix cross-expert output overwrites in SM90 masked GEMM #472

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/469
Author: @Sunt-ing
Created: 10/12/2026
Status: ❌ Closed

Base: main ← Head: codex/sm90-masked-output-boundaries

📝 Commits (1)

2086593 Fix SM90 masked GEMM output boundaries

📊 Changes

7 files changed (+62 additions, -13 deletions)

View changed files

📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+5 -3)
📝 csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp (+7 -5)
📝 deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh (+2 -2)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh (+2 -3)
📝 tests/test_bf16.py (+7 -0)
📝 tests/test_fp8_fp4.py (+7 -0)
📝 tests/utils.py (+32 -0)

📄 Description

SM90 masked GEMM flattens the output into a 2D TMA descriptor. When a store tile extends past an expert's capacity, it can overwrite valid rows of the next expert. This affects both BF16 and FP8: capacity 32 reproduces it, and capacity 96 also fails when experts are nearly full.

Keep the expert dimension in the output descriptor and reuse the existing 3D TMA store path. This preserves tile selection and the compute pipeline without adding buffers or synchronization. Boundary regressions are integrated into the existing BF16 and FP8 test suites.

This revisits the limitation discussed in #15 and the small-capacity failure reported in #268. #15 previously fixed small capacities by shrinking the store tile and rejected larger capacities not divisible by the tile size; it also suggested 3D TMA but deferred support for those cases. This change handles both small and non-tile-aligned capacities through the same output path.

On H200, the median paired latency changes across six rounds were between −0.17% and +0.16% for the larger tested shapes. The smallest shape was about 1.9% slower for BF16 and 1.5% slower for FP8, roughly 0.06–0.08 µs. The uniform 3D path avoids maintaining separate aligned and unaligned implementations, at that small measured cost.

Reproduction and validation

Requires an SM90 GPU and a source build of DeepGEMM. Save the following as repro.py outside the source checkout and run python repro.py against the installed baseline or patched package:

import torch
import deep_gemm
from deep_gemm.testing import calc_diff

torch.manual_seed(29)
a = torch.randn((8, 32, 256), device='cuda', dtype=torch.bfloat16)
b = torch.randn((8, 256, 256), device='cuda', dtype=torch.bfloat16)
d = torch.empty((8, 32, 256), device='cuda', dtype=torch.bfloat16)
counts = torch.tensor([0, 1, 16, 31, 32, 1, 32, 32],
                      device='cuda', dtype=torch.int)
valid = torch.arange(32, device='cuda')[None, :] < counts[:, None]
reference = torch.bmm(a.double(), b.double().mT).bfloat16()

for _ in range(16):
    deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, counts, 32)
    diff = calc_diff(d[valid], reference[valid])
    assert diff < 1e-5, f'{diff=}'

Running this exact script against the baseline failed with AssertionError: diff=tensor(0.1456, device='cuda:0', dtype=torch.float64); the patched wheel completed all 16 checks. Only valid rows are checked; padding values are unspecified.

Run the included regressions from the checkout after building and installing the patched package:

cd tests
python - <<'PY'
import torch
import test_bf16
import test_fp8_fp4

torch.manual_seed(17)
test_bf16.test_m_grouped_gemm_masked_boundaries()
test_fp8_fp4.test_m_grouped_gemm_masked_boundaries()
PY

Both passed with the final built wheel, including a fresh-JIT-cache run. Both fail at capacity 16 when run against the baseline package. A separate timed run with the JIT cache populated took 0.23 s for the BF16 test body and 0.09 s for FP8; this excludes process startup and imports, and does not estimate cold compilation time.

Supplemental validation, using separate local harnesses not included in this PR:

Base: 057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7; candidate: 20865936a058d55c4c583cc2b0c2bcf3bc319a2f. H200, PyTorch 2.13.0+cu130, CUDA 13.0.
34 configurations per arm, covering BF16 and FP8, capacities from 1 to 257, single-expert and all-empty cases; 16 eager and 16 Graph calls each. Baseline and revert each had 24 failing configurations; the fix had none. Two configurations were all-empty and only checked completion. The revert reused the original package in a fresh process/cache.
Independent FP64 matmul references, rounded to BF16; FP8 references used the actual quantized inputs and scales. No run-to-run changes were observed in the fixed panel.
37 adjacent-path checks passed, covering ordinary GEMM layouts, output dtypes and accumulation, batched GEMM, and K-grouped GEMM.
Native python setup.py bdist_wheel passed with setuptools 75.8.0.

Performance used six paired rounds with alternating arm order. Each measurement took the median of seven CUDA-event samples, each containing 1,000 Graph calls. Input generation and quantization were outside timing; FP8 API scale conversion was included. Positive values below mean slower; values are medians of paired latency changes.

Groups	Capacity	N	K	Expected M	BF16	FP8
8	128	256	256	64	+1.931%	+1.526%
32	4096	1536	2048	20	+0.070%	−0.143%
6	4096	4096	4096	1024	+0.155%	+0.023%
32	4096	6144	7168	192	−0.167%	−0.056%

Clocks were not locked; sampled SM clocks ranged from 1305 to 1980 MHz. Small differences on the larger shapes should not be interpreted as established speedups. CPU launch latency, model-level impact, other GPU architectures, and the full build matrix were not evaluated.

_{🔄 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/469 **Author:** [@Sunt-ing](https://github.com/Sunt-ing) **Created:** 10/12/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `codex/sm90-masked-output-boundaries` --- ### 📝 Commits (1) - [`2086593`](https://github.com/deepseek-ai/DeepGEMM/commit/20865936a058d55c4c583cc2b0c2bcf3bc319a2f) Fix SM90 masked GEMM output boundaries ### 📊 Changes **7 files changed** (+62 additions, -13 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+5 -3) 📝 `csrc/jit_kernels/impls/sm90_fp8_gemm_1d2d.hpp` (+7 -5) 📝 `deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh` (+2 -2) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_gemm_1d2d.cuh` (+2 -3) 📝 `tests/test_bf16.py` (+7 -0) 📝 `tests/test_fp8_fp4.py` (+7 -0) 📝 `tests/utils.py` (+32 -0) </details> ### 📄 Description SM90 masked GEMM flattens the output into a 2D TMA descriptor. When a store tile extends past an expert's capacity, it can overwrite valid rows of the next expert. This affects both BF16 and FP8: capacity 32 reproduces it, and capacity 96 also fails when experts are nearly full. Keep the expert dimension in the output descriptor and reuse the existing 3D TMA store path. This preserves tile selection and the compute pipeline without adding buffers or synchronization. Boundary regressions are integrated into the existing BF16 and FP8 test suites. This revisits the limitation discussed in #15 and the small-capacity failure reported in #268. #15 previously fixed small capacities by shrinking the store tile and rejected larger capacities not divisible by the tile size; it also suggested 3D TMA but deferred support for those cases. This change handles both small and non-tile-aligned capacities through the same output path. On H200, the median paired latency changes across six rounds were between −0.17% and +0.16% for the larger tested shapes. The smallest shape was about 1.9% slower for BF16 and 1.5% slower for FP8, roughly 0.06–0.08 µs. The uniform 3D path avoids maintaining separate aligned and unaligned implementations, at that small measured cost. <details> <summary>Reproduction and validation</summary> Requires an SM90 GPU and a source build of DeepGEMM. Save the following as `repro.py` outside the source checkout and run `python repro.py` against the installed baseline or patched package: ```python import torch import deep_gemm from deep_gemm.testing import calc_diff torch.manual_seed(29) a = torch.randn((8, 32, 256), device='cuda', dtype=torch.bfloat16) b = torch.randn((8, 256, 256), device='cuda', dtype=torch.bfloat16) d = torch.empty((8, 32, 256), device='cuda', dtype=torch.bfloat16) counts = torch.tensor([0, 1, 16, 31, 32, 1, 32, 32], device='cuda', dtype=torch.int) valid = torch.arange(32, device='cuda')[None, :] < counts[:, None] reference = torch.bmm(a.double(), b.double().mT).bfloat16() for _ in range(16): deep_gemm.m_grouped_bf16_gemm_nt_masked(a, b, d, counts, 32) diff = calc_diff(d[valid], reference[valid]) assert diff < 1e-5, f'{diff=}' ``` Running this exact script against the baseline failed with `AssertionError: diff=tensor(0.1456, device='cuda:0', dtype=torch.float64)`; the patched wheel completed all 16 checks. Only valid rows are checked; padding values are unspecified. Run the included regressions from the checkout after building and installing the patched package: ```bash cd tests python - <<'PY' import torch import test_bf16 import test_fp8_fp4 torch.manual_seed(17) test_bf16.test_m_grouped_gemm_masked_boundaries() test_fp8_fp4.test_m_grouped_gemm_masked_boundaries() PY ``` Both passed with the final built wheel, including a fresh-JIT-cache run. Both fail at capacity 16 when run against the baseline package. A separate timed run with the JIT cache populated took 0.23 s for the BF16 test body and 0.09 s for FP8; this excludes process startup and imports, and does not estimate cold compilation time. Supplemental validation, using separate local harnesses not included in this PR: - Base: `057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7`; candidate: `20865936a058d55c4c583cc2b0c2bcf3bc319a2f`. H200, PyTorch `2.13.0+cu130`, CUDA 13.0. - 34 configurations per arm, covering BF16 and FP8, capacities from 1 to 257, single-expert and all-empty cases; 16 eager and 16 Graph calls each. Baseline and revert each had 24 failing configurations; the fix had none. Two configurations were all-empty and only checked completion. The revert reused the original package in a fresh process/cache. - Independent FP64 matmul references, rounded to BF16; FP8 references used the actual quantized inputs and scales. No run-to-run changes were observed in the fixed panel. - 37 adjacent-path checks passed, covering ordinary GEMM layouts, output dtypes and accumulation, batched GEMM, and K-grouped GEMM. - Native `python setup.py bdist_wheel` passed with setuptools `75.8.0`. Performance used six paired rounds with alternating arm order. Each measurement took the median of seven CUDA-event samples, each containing 1,000 Graph calls. Input generation and quantization were outside timing; FP8 API scale conversion was included. Positive values below mean slower; values are medians of paired latency changes. | Groups | Capacity | N | K | Expected M | BF16 | FP8 | |---:|---:|---:|---:|---:|---:|---:| | 8 | 128 | 256 | 256 | 64 | +1.931% | +1.526% | | 32 | 4096 | 1536 | 2048 | 20 | +0.070% | −0.143% | | 6 | 4096 | 4096 | 4096 | 1024 | +0.155% | +0.023% | | 32 | 4096 | 6144 | 7168 | 192 | −0.167% | −0.056% | Clocks were not locked; sampled SM clocks ranged from 1305 to 1980 MHz. Small differences on the larger shapes should not be interpreted as established speedups. CPU launch latency, model-level impact, other GPU architectures, and the full build matrix were not evaluated. </details> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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