[PR #31] [CLOSED] feat: support groupwise scaling for RHS #213

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/31
Author: @XG-zheng
Created: 3/1/2025
Status: ❌ Closed

Base: main ← Head: groupwise_gemm

📝 Commits (1)

6d0fd7d support groupwise scaling for b

📊 Changes

4 files changed (+445 additions, -26 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/fp8_gemm.cuh (+371 -0)
📝 deep_gemm/include/deep_gemm/mma_utils.cuh (+12 -0)
📝 deep_gemm/jit_kernels/gemm.py (+44 -15)
📝 tests/test_core.py (+18 -11)

📄 Description

This MR extends the existing FP8 GEMM to support 128x1 groupwise scaling on the RHS.
Just some minor modifications:

use TMA to load scales_b from global memory to shared memory
scales_b is loaded from shared memory to the register Before each stage of accumulation scaling

test on H800 && CUDA 12.4:

correctness: The shape in the readme can be passed.
perf:
- Small Shapes: Achieved 90+% performance of the original FP8 gemm.
- Large Shapes (e.g., 4096x4096x7168): Shows ~30% performance regression.Still need optimization. Tried some optimizations but no good results.

M	N	K	Groupwise FP8 Gemm-Computation	FP8 Gemm-Computation
64	2112	7168	153 TFLOPS	190 TFLOPS
64	24576	1536	262 TFLOPS	283 TFLOPS
64	32768	512	212 TFLOPS	207 TFLOPS
64	7168	16384	307 TFLOPS	326 TFLOPS
64	4096	7168	251 TFLOPS	268 TFLOPS
64	7168	2048	246 TFLOPS	262 TFLOPS
128	2112	7168	275 TFLOPS	337 TFLOPS
128	24576	1536	472 TFLOPS	526 TFLOPS
128	32768	512	328 TFLOPS	347 TFLOPS
128	7168	16384	595 TFLOPS	628 TFLOPS
128	4096	7168	456 TFLOPS	502 TFLOPS
128	7168	2048	424 TFLOPS	463 TFLOPS
4096	2112	7168	861 TFLOPS	1016 TFLOPS
4096	24576	1536	767 TFLOPS	973 TFLOPS
4096	32768	512	557 TFLOPS	592 TFLOPS
4096	7168	16384	964 TFLOPS	1313 TFLOPS
4096	4096	7168	925 TFLOPS	1278 TFLOPS
4096	7168	2048	844 TFLOPS	988 TFLOPS

_{🔄 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/31 **Author:** [@XG-zheng](https://github.com/XG-zheng) **Created:** 3/1/2025 **Status:** ❌ Closed **Base:** `main` ← **Head:** `groupwise_gemm` --- ### 📝 Commits (1) - [`6d0fd7d`](https://github.com/deepseek-ai/DeepGEMM/commit/6d0fd7de41ac37ce3ff5fe1c5da1df670a63169b) support groupwise scaling for b ### 📊 Changes **4 files changed** (+445 additions, -26 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/fp8_gemm.cuh` (+371 -0) 📝 `deep_gemm/include/deep_gemm/mma_utils.cuh` (+12 -0) 📝 `deep_gemm/jit_kernels/gemm.py` (+44 -15) 📝 `tests/test_core.py` (+18 -11) </details> ### 📄 Description This MR extends the existing FP8 GEMM to support 128x1 groupwise scaling on the RHS. Just some minor modifications: - use TMA to load scales_b from global memory to shared memory - scales_b is loaded from shared memory to the register Before each stage of accumulation scaling test on H800 && CUDA 12.4: 1. correctness: The shape in the readme can be passed. 2. perf: - Small Shapes: Achieved 90+% performance of the original FP8 gemm. - Large Shapes (e.g., 4096x4096x7168): Shows ~30% performance regression.Still need optimization. Tried some optimizations but no good results. M | N | K | Groupwise FP8 Gemm-Computation | FP8 Gemm-Computation -- | -- | -- | -- | -- 64 | 2112 | 7168 | 153 TFLOPS | 190 TFLOPS 64 | 24576 | 1536 | 262 TFLOPS | 283 TFLOPS 64 | 32768 | 512 | 212 TFLOPS | 207 TFLOPS 64 | 7168 | 16384 | 307 TFLOPS | 326 TFLOPS 64 | 4096 | 7168 | 251 TFLOPS | 268 TFLOPS 64 | 7168 | 2048 | 246 TFLOPS | 262 TFLOPS 128 | 2112 | 7168 | 275 TFLOPS | 337 TFLOPS 128 | 24576 | 1536 | 472 TFLOPS | 526 TFLOPS 128 | 32768 | 512 | 328 TFLOPS | 347 TFLOPS 128 | 7168 | 16384 | 595 TFLOPS | 628 TFLOPS 128 | 4096 | 7168 | 456 TFLOPS | 502 TFLOPS 128 | 7168 | 2048 | 424 TFLOPS | 463 TFLOPS 4096 | 2112 | 7168 | 861 TFLOPS | 1016 TFLOPS 4096 | 24576 | 1536 | 767 TFLOPS | 973 TFLOPS 4096 | 32768 | 512 | 557 TFLOPS | 592 TFLOPS 4096 | 7168 | 16384 | 964 TFLOPS | 1313 TFLOPS 4096 | 4096 | 7168 | 925 TFLOPS | 1278 TFLOPS 4096 | 7168 | 2048 | 844 TFLOPS | 988 TFLOPS --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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