[PR #474] [Feat]Add native FP8 Tensor Core GEMM support for SM89 (Ada Lovelace) #476

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/474
Author: @Ammar-Alnagar
Created: 10/14/2026
Status: 🔄 Open

Base: main ← Head: main

📝 Commits (3)

4d2e2fc Add SM80 and SM89 architecture support and fallbacks (#1)
fc5729f Remove SM80 and SM89 hardware validation test suite
fdcb2af Add support for FP8 Tensor Cores on SM89 architecture

📊 Changes

20 files changed (+1479 additions, -38 deletions)

View changed files

📝 README.md (+41 -13)
📝 csrc/apis/attention.hpp (+22 -6)
📝 csrc/apis/config.hpp (+3 -0)
📝 csrc/apis/gemm.hpp (+47 -5)
📝 csrc/apis/layout.hpp (+8 -4)
➕ csrc/jit_kernels/heuristics/sm89.hpp (+170 -0)
➕ csrc/jit_kernels/impls/sm89_fp8_gemm.hpp (+143 -0)
📝 csrc/jit_kernels/impls/smxx_cublaslt.hpp (+5 -1)
📝 csrc/utils/exception.hpp (+11 -0)
📝 csrc/utils/layout.hpp (+13 -2)
📝 deep_gemm/include/deep_gemm/common/types.cuh (+34 -0)
➕ deep_gemm/include/deep_gemm/impls/sm89_fp8_gemm.cuh (+417 -0)
📝 deep_gemm/include/deep_gemm/impls/smxx_layout.cuh (+3 -3)
➕ deep_gemm/include/deep_gemm/mma/sm89.cuh (+40 -0)
📝 deep_gemm/include/deep_gemm/ptx/ld_st.cuh (+19 -0)
📝 deep_gemm/include/deep_gemm/ptx/utils.cuh (+9 -0)
📝 deep_gemm/testing/utils.py (+14 -0)
➕ docs/architecture_support.md (+158 -0)
📝 tests/generators.py (+5 -4)
➕ tests/test_sm89.py (+317 -0)

📄 Description

Title
Add native FP8 Tensor Core GEMM support for SM89 (Ada Lovelace)

Body

Summary

This change adds a native FP8 Tensor Core GEMM for SM89 (Ada Lovelace). SM89 is the first architecture below Hopper that has FP8 Tensor Cores.

Previously, SM80 and SM89 used the same path. The path required SM90 or higher for FP8. This was correct for SM80 because Ampere has no FP8 Tensor Cores. It was wrong for SM89. SM89 implements the instruction mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32 (QMMA.16832.F32.E4M3.E4M3). The change separates the two architectures and gives SM89 a real kernel.

Why

SM89 has FP8 Tensor Cores. It does not have the SM90 features (no TMA, no mbarrier.try_wait.parity, no WGMMA, no clusters). It could not reuse the SM90 or SM100 kernels. It fell back to BF16 through cuBLASLt.
FP8 users on Ada had to dequantize and run BF16. They lost about 2× of Tensor Core throughput.
The shared SM80/SM89 path also gave wrong error messages. An SM89 user saw a message that did not apply to their GPU.

Design

SM89 has no TMA and no asynchronous barriers. The kernel uses the pre-Hopper primitives.

It uses mma.sync (wrapped in deep_gemm/mma/sm89.cuh, with A/B/C fragment layouts documented) and ldmatrix.x4 for operand staging.
It uses a multi-stage cp.async pipeline with a 128-byte XOR swizzle.
It uses a persistent block scheduler (one CTA per SM). Tiles walk in a loop.
It uses direct global stores because SM89 has no TMA store. It has an alignment-safe scalar tail. An MMA lane owns a pair of adjacent columns. For odd N the element offset of a row can be odd. A 2-element vector store is then illegal. This caused the original misaligned address failure.
Block K is fixed at 128 because of the per-128-channel scaling granularity. Pipeline depth and tile shapes come from a new cost model in csrc/jit_kernels/heuristics/sm89.hpp. The warp split uses one shared function get_mma_sync_warps_n(). Host heuristics and device code both call it so they cannot disagree.

Both SM90 SF recipes are supported:
(1, 1, 128) (per-token A and B scales) and
(1, 128, 128) (per-token A, one B scale per 128-column N block).
The second recipe is the default when FP32 scales are used. SM89 uses the same FP32 SF layouts as SM90. The function transform_sf_into_required_layout was extended. It was not duplicated.

Scope

Supported features:
dense FP8 GEMM, K-major (NT) operands only, FP32 and BF16 outputs, D = C + A @ B, ragged M/N, k % 16 == 0, both aliases (fp8_fp4_gemm_nt, fp8_gemm_nt, fp4_gemm_nt), and shape specialization.

Still rejected (with an explicit message that names the requirement):
grouped and K-grouped FP8, MN-major FP8 operands, FP4, custom epilogues, alpha scaling, PDL, and MQA/attention.
BF16 keeps its cuBLASLt fallback.

Incidental fixes

set_pdl() now panics below SM90. It no longer sets a flag that griddepcontrol cannot honour.
cudaGridDependencySynchronize() is wrapped in a no-op below __CUDA_ARCH__ >= 900 so the SF-layout kernels compile for SM89.
call_cublaslt_api used CUBLAS_COMPUTE_32F_FAST_TF32 for non-FP32 operands. That mode only down-casts FP32 inputs to TF32. BF16/FP8 operands ran at TF32 precision by mistake. Now only FP32 inputs use that mode.
A new DG_PANIC macro throws always and carries a message.

Bug found during validation

The persistent loop primed its cp.async pipeline only in the prologue before the tile loop. It relied on a cross-tile tail load to prime the next tile. A tile’s tail can only reach the stages that are already free of the current tile — prologue stages p ∈ [0, kNumStages-2-num_k_blocks]. The next tile’s K loop consumes p = 0..num_k_blocks-1. When num_k_blocks < kNumStages-1 (example: k=512, 4 blocks, 8 stages) the high stages were never reloaded. The next tile read the previous tile’s shared memory. Symptom: non-deterministic nan and values near 1e30. The problem appeared only for small M spread over many N tiles.

The fix extracts the prologue into prime_pipeline() and re-primes unconditionally after the epilogue. The re-prime is guarded by cp_async_wait<0>() and a __syncthreads(). cp.async completion is per-thread, so both are required before the stages can be reused. The change removes a prologue/epilogue overlap that was never correct. Measured cost is about 0–3.5 % (4096³: 2518 µs vs 2507 µs, within noise; 4096×4096×512: 420 µs vs 405 µs).

tests/test_sm89.py added the shapes (1, 4096, 512), (64, 4096, 512) and (128, 4096, 128). Reverting the fix makes these tests fail (34 failures). They are real regression tests.

Testing

Hardware: RTX 4060 Laptop (SM89, 24 SMs, 7.6 GiB), CUDA 13.0, PyTorch 2.14.1+cu130.

tests/test_sm89.py — 17 tests, OK from a cold JIT cache (~130 s).
Covers both recipes × both output dtypes × accumulate × 12 shapes, explicit recipe=(1,128,128), recipe_a/recipe_b pairs, entry-point aliases, compiled-dim specialization, bitwise determinism across repeats, accumulation-versus-zero-C equivalence, BF16 GEMM and all its layouts, cuBLASLt FP8/BF16, SF layout kernels, and one rejection case per unsupported feature.
tests/test_layout.py — OK.
Upstream tests/test_fp8_fp4.py::test_gemm — OK.
The upstream shape grid reaches m=129280 / n=7168 and OOMs on a 7.6 GiB card. Shapes were clamped to m≤4096, n≤4096, k≤1280 by monkey-patching enumerate_normal in a throw-away driver. The test body, recipes and assertions were not changed.
Upstream tests/test_bf16.py::test_gemm — correctness OK, same clamping.
The perf-print block was stubbed because bench_kineto matches on the kernel name bf16_gemm, which does not exist on SM89 (BF16 goes through the cuBLASLt fallback). All diff < 1e-5 assertions ran without change.
Other suites fail only for environmental reasons and are unchanged by this PR:
test_einsum raises the intended “Unsupported architecture”,
test_lazy_init and test_locality_domain need multiple GPUs,
test_bf16 / test_legacy OOM at upstream shapes.

Throughput, single GPU, FP8 in / BF16 out:

m	n	k	time	TFLOPS
1	4096	4096	33.6 µs	1.0
128	4096	4096	168.4 µs	25.5
4096	4096	4096	2517.8 µs	54.6
4096	4096	512	419.7 µs	40.9

AI usage Disclaimer : Ai was used in the documentation and creation of this PR message

_{🔄 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/474 **Author:** [@Ammar-Alnagar](https://github.com/Ammar-Alnagar) **Created:** 10/14/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `main` --- ### 📝 Commits (3) - [`4d2e2fc`](https://github.com/deepseek-ai/DeepGEMM/commit/4d2e2fcda640ccbe4b903d3be61cc2cc24b45d7a) Add SM80 and SM89 architecture support and fallbacks (#1) - [`fc5729f`](https://github.com/deepseek-ai/DeepGEMM/commit/fc5729ffe65c57a3f831f8cf0a518a4583df9b54) Remove SM80 and SM89 hardware validation test suite - [`fdcb2af`](https://github.com/deepseek-ai/DeepGEMM/commit/fdcb2af766549332dbdd57a1ed119aecda93c264) Add support for FP8 Tensor Cores on SM89 architecture ### 📊 Changes **20 files changed** (+1479 additions, -38 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+41 -13) 📝 `csrc/apis/attention.hpp` (+22 -6) 📝 `csrc/apis/config.hpp` (+3 -0) 📝 `csrc/apis/gemm.hpp` (+47 -5) 📝 `csrc/apis/layout.hpp` (+8 -4) ➕ `csrc/jit_kernels/heuristics/sm89.hpp` (+170 -0) ➕ `csrc/jit_kernels/impls/sm89_fp8_gemm.hpp` (+143 -0) 📝 `csrc/jit_kernels/impls/smxx_cublaslt.hpp` (+5 -1) 📝 `csrc/utils/exception.hpp` (+11 -0) 📝 `csrc/utils/layout.hpp` (+13 -2) 📝 `deep_gemm/include/deep_gemm/common/types.cuh` (+34 -0) ➕ `deep_gemm/include/deep_gemm/impls/sm89_fp8_gemm.cuh` (+417 -0) 📝 `deep_gemm/include/deep_gemm/impls/smxx_layout.cuh` (+3 -3) ➕ `deep_gemm/include/deep_gemm/mma/sm89.cuh` (+40 -0) 📝 `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` (+19 -0) 📝 `deep_gemm/include/deep_gemm/ptx/utils.cuh` (+9 -0) 📝 `deep_gemm/testing/utils.py` (+14 -0) ➕ `docs/architecture_support.md` (+158 -0) 📝 `tests/generators.py` (+5 -4) ➕ `tests/test_sm89.py` (+317 -0) </details> ### 📄 Description **Title** Add native FP8 Tensor Core GEMM support for SM89 (Ada Lovelace) **Body** ## Summary This change adds a native FP8 Tensor Core GEMM for SM89 (Ada Lovelace). SM89 is the first architecture below Hopper that has FP8 Tensor Cores. Previously, SM80 and SM89 used the same path. The path required SM90 or higher for FP8. This was correct for SM80 because Ampere has no FP8 Tensor Cores. It was wrong for SM89. SM89 implements the instruction `mma.sync.aligned.m16n8k32.row.col.f32.e4m3.e4m3.f32` (`QMMA.16832.F32.E4M3.E4M3`). The change separates the two architectures and gives SM89 a real kernel. ## Why - SM89 has FP8 Tensor Cores. It does not have the SM90 features (no TMA, no `mbarrier.try_wait.parity`, no WGMMA, no clusters). It could not reuse the SM90 or SM100 kernels. It fell back to BF16 through cuBLASLt. - FP8 users on Ada had to dequantize and run BF16. They lost about 2× of Tensor Core throughput. - The shared SM80/SM89 path also gave wrong error messages. An SM89 user saw a message that did not apply to their GPU. ## Design SM89 has no TMA and no asynchronous barriers. The kernel uses the pre-Hopper primitives. - It uses `mma.sync` (wrapped in `deep_gemm/mma/sm89.cuh`, with A/B/C fragment layouts documented) and `ldmatrix.x4` for operand staging. - It uses a multi-stage `cp.async` pipeline with a 128-byte XOR swizzle. - It uses a persistent block scheduler (one CTA per SM). Tiles walk in a loop. - It uses direct global stores because SM89 has no TMA store. It has an alignment-safe scalar tail. An MMA lane owns a pair of adjacent columns. For odd N the element offset of a row can be odd. A 2-element vector store is then illegal. This caused the original misaligned address failure. - Block K is fixed at 128 because of the per-128-channel scaling granularity. Pipeline depth and tile shapes come from a new cost model in `csrc/jit_kernels/heuristics/sm89.hpp`. The warp split uses one shared function `get_mma_sync_warps_n()`. Host heuristics and device code both call it so they cannot disagree. Both SM90 SF recipes are supported: `(1, 1, 128)` (per-token A and B scales) and `(1, 128, 128)` (per-token A, one B scale per 128-column N block). The second recipe is the default when FP32 scales are used. SM89 uses the same FP32 SF layouts as SM90. The function `transform_sf_into_required_layout` was extended. It was not duplicated. ## Scope Supported features: dense FP8 GEMM, K-major (NT) operands only, FP32 and BF16 outputs, `D = C + A @ B`, ragged M/N, `k % 16 == 0`, both aliases (`fp8_fp4_gemm_nt`, `fp8_gemm_nt`, `fp4_gemm_nt`), and shape specialization. Still rejected (with an explicit message that names the requirement): grouped and K-grouped FP8, MN-major FP8 operands, FP4, custom epilogues, `alpha` scaling, PDL, and MQA/attention. BF16 keeps its cuBLASLt fallback. ## Incidental fixes - `set_pdl()` now panics below SM90. It no longer sets a flag that `griddepcontrol` cannot honour. `cudaGridDependencySynchronize()` is wrapped in a no-op below `__CUDA_ARCH__ >= 900` so the SF-layout kernels compile for SM89. - `call_cublaslt_api` used `CUBLAS_COMPUTE_32F_FAST_TF32` for non-FP32 operands. That mode only down-casts FP32 inputs to TF32. BF16/FP8 operands ran at TF32 precision by mistake. Now only FP32 inputs use that mode. - A new `DG_PANIC` macro throws always and carries a message. ## Bug found during validation The persistent loop primed its `cp.async` pipeline only in the prologue before the tile loop. It relied on a cross-tile tail load to prime the next tile. A tile’s tail can only reach the stages that are already free of the current tile — prologue stages `p ∈ [0, kNumStages-2-num_k_blocks]`. The next tile’s K loop consumes `p = 0..num_k_blocks-1`. When `num_k_blocks < kNumStages-1` (example: `k=512`, 4 blocks, 8 stages) the high stages were never reloaded. The next tile read the previous tile’s shared memory. Symptom: non-deterministic `nan` and values near 1e30. The problem appeared only for small M spread over many N tiles. The fix extracts the prologue into `prime_pipeline()` and re-primes unconditionally after the epilogue. The re-prime is guarded by `cp_async_wait<0>()` and a `__syncthreads()`. `cp.async` completion is per-thread, so both are required before the stages can be reused. The change removes a prologue/epilogue overlap that was never correct. Measured cost is about 0–3.5 % (4096³: 2518 µs vs 2507 µs, within noise; 4096×4096×512: 420 µs vs 405 µs). `tests/test_sm89.py` added the shapes `(1, 4096, 512)`, `(64, 4096, 512)` and `(128, 4096, 128)`. Reverting the fix makes these tests fail (34 failures). They are real regression tests. ## Testing Hardware: RTX 4060 Laptop (SM89, 24 SMs, 7.6 GiB), CUDA 13.0, PyTorch 2.14.1+cu130. - `tests/test_sm89.py` — 17 tests, OK from a cold JIT cache (~130 s). Covers both recipes × both output dtypes × accumulate × 12 shapes, explicit `recipe=(1,128,128)`, `recipe_a`/`recipe_b` pairs, entry-point aliases, compiled-dim specialization, bitwise determinism across repeats, accumulation-versus-zero-C equivalence, BF16 GEMM and all its layouts, cuBLASLt FP8/BF16, SF layout kernels, and one rejection case per unsupported feature. - `tests/test_layout.py` — OK. - Upstream `tests/test_fp8_fp4.py::test_gemm` — OK. The upstream shape grid reaches m=129280 / n=7168 and OOMs on a 7.6 GiB card. Shapes were clamped to m≤4096, n≤4096, k≤1280 by monkey-patching `enumerate_normal` in a throw-away driver. The test body, recipes and assertions were not changed. - Upstream `tests/test_bf16.py::test_gemm` — correctness OK, same clamping. The perf-print block was stubbed because `bench_kineto` matches on the kernel name `bf16_gemm`, which does not exist on SM89 (BF16 goes through the cuBLASLt fallback). All `diff < 1e-5` assertions ran without change. - Other suites fail only for environmental reasons and are unchanged by this PR: `test_einsum` raises the intended “Unsupported architecture”, `test_lazy_init` and `test_locality_domain` need multiple GPUs, `test_bf16` / `test_legacy` OOM at upstream shapes. Throughput, single GPU, FP8 in / BF16 out: | m | n | k | time | TFLOPS | |------|------|------|-----------|--------| | 1 | 4096 | 4096 | 33.6 µs | 1.0 | | 128 | 4096 | 4096 | 168.4 µs | 25.5 | | 4096 | 4096 | 4096 | 2517.8 µs | 54.6 | | 4096 | 4096 | 512 | 419.7 µs | 40.9 | AI usage Disclaimer : Ai was used in the documentation and creation of this PR message --- <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:19 +00:00

bsakaguchi referenced this issue

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

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