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

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/463
Author: @yuweih205
Created: 9/30/2026
Status: 🔄 Open

Base: main ← Head: perf/sm90-gathered-bf16

📝 Commits (1)

ca1291f Add SM90 BF16 grouped GEMM with gathered activation rows

📊 Changes

6 files changed (+258 additions, -11 deletions)

View changed files

📝 README.md (+2 -0)
📝 csrc/apis/gemm.hpp (+40 -0)
📝 csrc/jit_kernels/impls/sm90_bf16_gemm.hpp (+72 -3)
📝 deep_gemm/__init__.py (+1 -0)
📝 deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh (+64 -6)
📝 tests/test_bf16.py (+79 -2)

📄 Description

Routed BF16 forward paths currently materialize a.index_select(0, rows) before grouped GEMM. This adds m_grouped_bf16_gemm_nt_contiguous_gathered(a, b, d, grouped_layout, a_row_indices) on SM90, loading A through row indices directly into the existing WGMMA pipeline.

The first version supports K-major BF16, K divisible by 64, contiguous weights, 128-row expert alignment and an int64 CUDA row map. A may have a padded row stride. The A loader uses 16-byte cp.async copies with the existing shared-memory swizzle; B/output retain TMA. It does not add a fused SwiGLU epilogue or change default dispatch.

Current revision

Commit ca1291f7cd2dbe483af7d89f033b2a5a1e96ca1f, based on main 1fc9c66975ee322d628b3288534ca0e9ce12a64d, keeps the change in six existing files:

Public API and registration in csrc/apis/gemm.hpp; configuration, TMA descriptors and launch in the existing SM90 implementation header.
Integrated coverage in tests/test_bf16.py, rather than a separate gathered test file. It covers three generated configurations, six boundary cases and three rejected invalid inputs.
The CUDA kernel's executable tokens are unchanged from the previous published revision. Host checks are reorganized while retaining the dimensionality, dtype, device, row-stride, alignment and empty-input constraints.

The native extension was forcibly rebuilt from this exact source and the focused gathered test passed on the original H200/SM90 stack: PyTorch 2.13.0+cu132, CUDA runtime 13.2, seed 42, strict deterministic algorithms and uninitialized-memory filling disabled. All generated results are checked bitwise against the materialized path. Both paths use CUDA events in the integrated test, but those printed timings are not a new controlled performance result.

This is build/numerical/boundary validation, not a rerun of the historical full integration, sanitizer or performance studies below.

Earlier 2.8.1 measurements

The following results were collected against the 2.8.1 release, #462 (057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7), at PR commit 0ff99780929827f24f19d5e67a10c390c5f77592.

On eight H200s, against activation gather + the existing GEMM with prepacked gate/up weights + the same fused output-compaction/SwiGLU kernel:

D=3072, F=1280, 16 experts	Eager reduction, run 1 / repeat	CUDA Graph reduction
R=8192, M=16384, random routes	6.42% / 6.31%	—
R=108031, M=170003, random routes	6.07% / 5.83%	—
Source-ordered routes, nominal shape	6.18%	5.78%
Source-ordered routes, packet stride D+64	6.39%	6.16%

Eager numbers are median per-GPU ABBA/BAAB reductions; Graph is one GPU. The general eager sweep uses random routes with replacement, while Graph uses source-ordered unique routes; only the explicitly source-ordered rows compare matching routes across modes. Workers and process repeats reuse seed 42. An additional same-GPU, matching-row-map test at R=8192/M=16384 measures 6.25% eager / 6.01% Graph. The small R=1024, M=512, D=F=256 case measures 26.07% eager but regresses 2.31% under Graph, on matching routes. This is a fixed-route forward benchmark, not a full-model speedup.

All tested h/gate/up outputs match both existing-library baselines bitwise. API tests cover partial/empty M, empty experts, source/output stride, aligned offsets and misaligned pointer rejection. An additional 79-case regression covers existing BF16 GEMM/einsum layouts, output dtypes and accumulation modes using exact references and bytewise repeats. Independent FP32 reference checks pass; compute-sanitizer memcheck reports 0 errors. The full 13-phase 2.8.1 run passed, including live EP8 integration and receive-view lifetime checks; combined forward reductions are 16.99%–17.59% across two runs when both this change and DeepEP #769 are enabled.

2.8.1 reproducer, raw cycles, source hashes and limitations. These measurements retain their original test revision and predate the current reorganization; this update makes no additional performance claim. The earlier 2.8.0 measurements remain separately preserved in the same Gist.

Before this reorganization, the PR test scripts also passed a separate eight-H200 correctness rerun against the frozen production-equivalent runtimes. That DeepEP test revision checked an independent input/route oracle, raw deterministic order and two delayed consumer streams under both record-stream policies. Exact test revisions and retained test-development failures.

_{🔄 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/463 **Author:** [@yuweih205](https://github.com/yuweih205) **Created:** 9/30/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `perf/sm90-gathered-bf16` --- ### 📝 Commits (1) - [`ca1291f`](https://github.com/deepseek-ai/DeepGEMM/commit/ca1291f7cd2dbe483af7d89f033b2a5a1e96ca1f) Add SM90 BF16 grouped GEMM with gathered activation rows ### 📊 Changes **6 files changed** (+258 additions, -11 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+2 -0) 📝 `csrc/apis/gemm.hpp` (+40 -0) 📝 `csrc/jit_kernels/impls/sm90_bf16_gemm.hpp` (+72 -3) 📝 `deep_gemm/__init__.py` (+1 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh` (+64 -6) 📝 `tests/test_bf16.py` (+79 -2) </details> ### 📄 Description Routed BF16 forward paths currently materialize `a.index_select(0, rows)` before grouped GEMM. This adds `m_grouped_bf16_gemm_nt_contiguous_gathered(a, b, d, grouped_layout, a_row_indices)` on SM90, loading A through row indices directly into the existing WGMMA pipeline. The first version supports K-major BF16, K divisible by 64, contiguous weights, 128-row expert alignment and an int64 CUDA row map. A may have a padded row stride. The A loader uses 16-byte `cp.async` copies with the existing shared-memory swizzle; B/output retain TMA. It does not add a fused SwiGLU epilogue or change default dispatch. ## Current revision Commit `ca1291f7cd2dbe483af7d89f033b2a5a1e96ca1f`, based on main `1fc9c66975ee322d628b3288534ca0e9ce12a64d`, keeps the change in six existing files: - Public API and registration in `csrc/apis/gemm.hpp`; configuration, TMA descriptors and launch in the existing SM90 implementation header. - Integrated coverage in `tests/test_bf16.py`, rather than a separate gathered test file. It covers three generated configurations, six boundary cases and three rejected invalid inputs. - The CUDA kernel's executable tokens are unchanged from the previous published revision. Host checks are reorganized while retaining the dimensionality, dtype, device, row-stride, alignment and empty-input constraints. The native extension was forcibly rebuilt from this exact source and the focused gathered test passed on the original H200/SM90 stack: PyTorch `2.13.0+cu132`, CUDA runtime 13.2, seed 42, strict deterministic algorithms and uninitialized-memory filling disabled. All generated results are checked bitwise against the materialized path. Both paths use CUDA events in the integrated test, but those printed timings are not a new controlled performance result. This is build/numerical/boundary validation, not a rerun of the historical full integration, sanitizer or performance studies below. ## Earlier 2.8.1 measurements The following results were collected against the 2.8.1 release, #462 (`057ca5964aae0879ff2e0eb71ee05a3cb0ba3df7`), at PR commit `0ff99780929827f24f19d5e67a10c390c5f77592`. On eight H200s, against activation gather + the existing GEMM with **prepacked gate/up weights** + the same fused output-compaction/SwiGLU kernel: | D=3072, F=1280, 16 experts | Eager reduction, run 1 / repeat | CUDA Graph reduction | |---|---:|---:| | R=8192, M=16384, random routes | 6.42% / 6.31% | — | | R=108031, M=170003, random routes | 6.07% / 5.83% | — | | Source-ordered routes, nominal shape | 6.18% | 5.78% | | Source-ordered routes, packet stride D+64 | 6.39% | 6.16% | Eager numbers are median per-GPU ABBA/BAAB reductions; Graph is one GPU. The general eager sweep uses random routes with replacement, while Graph uses source-ordered unique routes; only the explicitly source-ordered rows compare matching routes across modes. Workers and process repeats reuse seed 42. An additional same-GPU, matching-row-map test at R=8192/M=16384 measures 6.25% eager / 6.01% Graph. The small R=1024, M=512, D=F=256 case measures 26.07% eager but **regresses 2.31% under Graph**, on matching routes. This is a fixed-route forward benchmark, not a full-model speedup. All tested h/gate/up outputs match both existing-library baselines bitwise. API tests cover partial/empty M, empty experts, source/output stride, aligned offsets and misaligned pointer rejection. An additional 79-case regression covers existing BF16 GEMM/einsum layouts, output dtypes and accumulation modes using exact references and bytewise repeats. Independent FP32 reference checks pass; compute-sanitizer memcheck reports 0 errors. The full 13-phase 2.8.1 run passed, including live EP8 integration and receive-view lifetime checks; combined forward reductions are 16.99%–17.59% across two runs when both this change and [DeepEP #769](https://github.com/deepseek-ai/DeepEP/pull/769) are enabled. [2.8.1 reproducer, raw cycles, source hashes and limitations](https://gist.github.com/yuweih205/d3fd1e4c3fcb97364f36ec8fda996cfe#file-dg281-infra-report-md). These measurements retain their original test revision and predate the current reorganization; this update makes no additional performance claim. The earlier 2.8.0 measurements remain separately preserved in the same Gist. Before this reorganization, the PR test scripts also passed a separate eight-H200 correctness rerun against the frozen production-equivalent runtimes. That DeepEP test revision checked an independent input/route oracle, raw deterministic order and two delayed consumer streams under both record-stream policies. [Exact test revisions and retained test-development failures](https://gist.github.com/yuweih205/d3fd1e4c3fcb97364f36ec8fda996cfe#file-review-gates-report-md). --- <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:09 +00:00

bsakaguchi referenced this issue

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

[PR #466] Make pack_ue8m0_to_int CUDA graph capture safe #470