deepseek-ai/DeepGEMM
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00
Code Issues 162 Projects Packages Wiki Activity

[PR #94] [MERGED] Refactor JIT compilation (+NVRTC support) #243

New issue
Closed
opened 2026-10-26 21:15:16 +00:00 by bsakaguchi · 0 comments
bsakaguchi commented 2026-10-26 21:15:16 +00:00
Owner
Copy link

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/94
Author: @lucifer1004
Created: 5/6/2025
Status: ✅ Merged
Merged: 5/7/2025
Merged by: @LyricZhao

Base: main ← Head: nvrtc

📝 Commits (10+)

  • 27cd276 [wip] refactor: compile to .cubin
  • c14cad0 refactor: compile to .cubin and add NVRTC option
  • 78c7fa3 fix: compiler version
  • 767793b feat: compat for old drivers
  • a3210ac feat: save kernel name to file
  • 40c09fb feat: fix win compat
  • 2d8c4f2 fix: windows compat
  • 6f0a17c Merge pull request #1 from lucifer1004/nvrtc-compat
  • 6c98279 Merge remote-tracking branch 'upstream/main' into nvrtc
  • 46762b6 feat: make API more general

📊 Changes

19 files changed (+917 additions, -668 deletions)

View changed files

📝 README.md (+22 -12)
📝 deep_gemm/include/deep_gemm/fp8_gemm.cuh (+6 -118)
📝 deep_gemm/include/deep_gemm/mma_utils.cuh (+2 -0)
➕ deep_gemm/include/deep_gemm/nvrtc_std.cuh (+103 -0)
📝 deep_gemm/include/deep_gemm/scheduler.cuh (+3 -2)
📝 deep_gemm/include/deep_gemm/tma_utils.cuh (+1 -76)
📝 deep_gemm/include/deep_gemm/utils.cuh (+4 -23)
📝 deep_gemm/jit/__init__.py (+1 -2)
📝 deep_gemm/jit/compiler.py (+209 -76)
📝 deep_gemm/jit/interleave_ffma.py (+3 -3)
📝 deep_gemm/jit/runtime.py (+56 -32)
➖ deep_gemm/jit/template.py (+0 -114)
📝 deep_gemm/jit_kernels/gemm.py (+51 -50)
📝 deep_gemm/jit_kernels/m_grouped_gemm.py (+84 -69)
➕ deep_gemm/jit_kernels/runtime.py (+254 -0)
📝 deep_gemm/jit_kernels/tuner.py (+26 -24)
📝 deep_gemm/utils.py (+1 -18)
📝 tests/test_core.py (+5 -0)
📝 tests/test_jit.py (+86 -49)

📄 Description

With this PR, the JIT compilation time is reduced by ~60% when using NVCC, ~80% when using NVRTC w/o PCH, and ~90% when using NVRTC w/ PCH.

Benchmark on Xeon(R) Platinum 8480C + H100 HBM:

m n k compilation time (s) TFLOPS
baseline NVCC NVRTC NVRTC+PCH baseline NVCC NVRTC NVRTC+PCH
64 2112 7168 5.95 2.12 0.96 0.35 154 154 153 153
64 24576 1536 6.43 2.55 1.37 0.78 230 230 231 230
64 32768 512 6.06 2.20 0.99 0.40 180 180 180 180
64 7168 16384 6.11 2.25 1.01 0.43 289 290 289 289
64 4096 7168 5.98 2.15 0.97 0.37 207 207 206 206
64 7168 2048 6.03 2.19 1.01 0.42 188 189 187 188
128 2112 7168 5.98 2.18 0.96 0.37 283 280 281 282
128 24576 1536 6.38 2.56 1.39 0.78 440 435 434 434
128 32768 512 6.00 2.18 0.99 0.40 334 336 330 332
128 7168 16384 6.24 2.40 1.18 0.59 534 534 534 533
128 4096 7168 6.04 2.21 1.03 0.43 370 371 370 371
128 7168 2048 6.19 2.38 1.19 0.62 332 334 331 332
4096 2112 7168 6.09 2.30 1.08 0.49 1053 1058 1113 1114
4096 24576 1536 6.14 2.33 1.10 0.50 1288 1286 1290 1289
4096 32768 512 6.23 2.41 1.21 0.61 913 912 912 911
4096 7168 16384 6.31 2.51 1.27 0.68 1524 1524 1459 1458
4096 4096 7168 6.28 2.47 1.27 0.68 1446 1445 1394 1396
4096 7168 2048 6.18 2.40 1.21 0.61 1240 1243 1234 1237

Note that there is some perf drop when using NVRTC due to a known bug of NVRTC which leads to extra instructions (but in the m=4096,n=2112,k=7168 case, NVRTC version was faster, which was a bit strange). So NVCC is kept as the default compiler for now, while NVRTC can be enabled with extra env var DG_JIT_USE_NVRTC.

🔄 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/94 **Author:** [@lucifer1004](https://github.com/lucifer1004) **Created:** 5/6/2025 **Status:** ✅ Merged **Merged:** 5/7/2025 **Merged by:** [@LyricZhao](https://github.com/LyricZhao) **Base:** `main` ← **Head:** `nvrtc` --- ### 📝 Commits (10+) - [`27cd276`](https://github.com/deepseek-ai/DeepGEMM/commit/27cd276e19515ceffa7231d3768cbfc01d5cc874) [wip] refactor: compile to .cubin - [`c14cad0`](https://github.com/deepseek-ai/DeepGEMM/commit/c14cad0c067c76beb5f14a62d9e1abfa511fa2f8) refactor: compile to .cubin and add NVRTC option - [`78c7fa3`](https://github.com/deepseek-ai/DeepGEMM/commit/78c7fa347e58dd1737ca67a7f0ec0b03ad2f9809) fix: compiler version - [`767793b`](https://github.com/deepseek-ai/DeepGEMM/commit/767793bf95c1c91b38509f107defebf42a824ddd) feat: compat for old drivers - [`a3210ac`](https://github.com/deepseek-ai/DeepGEMM/commit/a3210ac850fccc17774be0cb43e1dcdc8197b72a) feat: save kernel name to file - [`40c09fb`](https://github.com/deepseek-ai/DeepGEMM/commit/40c09fb88312a4f2a7684194255af8bb7da6855d) feat: fix win compat - [`2d8c4f2`](https://github.com/deepseek-ai/DeepGEMM/commit/2d8c4f22d5d8857b760f3e92e5a749dbf1b4ffcc) fix: windows compat - [`6f0a17c`](https://github.com/deepseek-ai/DeepGEMM/commit/6f0a17cb108d5ebb24e24b8d4e061906c2db5e08) Merge pull request #1 from lucifer1004/nvrtc-compat - [`6c98279`](https://github.com/deepseek-ai/DeepGEMM/commit/6c982791eb51af9b133dc5d1aa2f48ce0cfb1c77) Merge remote-tracking branch 'upstream/main' into nvrtc - [`46762b6`](https://github.com/deepseek-ai/DeepGEMM/commit/46762b690355eb2c7026c571b9f1d2e58d89b293) feat: make API more general ### 📊 Changes **19 files changed** (+917 additions, -668 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+22 -12) 📝 `deep_gemm/include/deep_gemm/fp8_gemm.cuh` (+6 -118) 📝 `deep_gemm/include/deep_gemm/mma_utils.cuh` (+2 -0) ➕ `deep_gemm/include/deep_gemm/nvrtc_std.cuh` (+103 -0) 📝 `deep_gemm/include/deep_gemm/scheduler.cuh` (+3 -2) 📝 `deep_gemm/include/deep_gemm/tma_utils.cuh` (+1 -76) 📝 `deep_gemm/include/deep_gemm/utils.cuh` (+4 -23) 📝 `deep_gemm/jit/__init__.py` (+1 -2) 📝 `deep_gemm/jit/compiler.py` (+209 -76) 📝 `deep_gemm/jit/interleave_ffma.py` (+3 -3) 📝 `deep_gemm/jit/runtime.py` (+56 -32) ➖ `deep_gemm/jit/template.py` (+0 -114) 📝 `deep_gemm/jit_kernels/gemm.py` (+51 -50) 📝 `deep_gemm/jit_kernels/m_grouped_gemm.py` (+84 -69) ➕ `deep_gemm/jit_kernels/runtime.py` (+254 -0) 📝 `deep_gemm/jit_kernels/tuner.py` (+26 -24) 📝 `deep_gemm/utils.py` (+1 -18) 📝 `tests/test_core.py` (+5 -0) 📝 `tests/test_jit.py` (+86 -49) </details> ### 📄 Description With this PR, the JIT compilation time is reduced by ~60% when using NVCC, ~80% when using NVRTC w/o PCH, and ~90% when using NVRTC w/ PCH. Benchmark on Xeon(R) Platinum 8480C + H100 HBM: | m | n | k | compilation time (s) | | | | TFLOPS | | | | |----|----|----|----|----|----|----|----|----|----|----| | | | | baseline | NVCC | NVRTC | NVRTC+PCH | baseline | NVCC | NVRTC | NVRTC+PCH | | 64 | 2112 | 7168 | 5.95 | 2.12 | 0.96 | 0.35 | 154 | 154 | 153 | 153 | | 64 | 24576 | 1536 | 6.43 | 2.55 | 1.37 | 0.78 | 230 | 230 | 231 | 230 | | 64 | 32768 | 512 | 6.06 | 2.20 | 0.99 | 0.40 | 180 | 180 | 180 | 180 | | 64 | 7168 | 16384 | 6.11 | 2.25 | 1.01 | 0.43 | 289 | 290 | 289 | 289 | | 64 | 4096 | 7168 | 5.98 | 2.15 | 0.97 | 0.37 | 207 | 207 | 206 | 206 | | 64 | 7168 | 2048 | 6.03 | 2.19 | 1.01 | 0.42 | 188 | 189 | 187 | 188 | | 128 | 2112 | 7168 | 5.98 | 2.18 | 0.96 | 0.37 | 283 | 280 | 281 | 282 | | 128 | 24576 | 1536 | 6.38 | 2.56 | 1.39 | 0.78 | 440 | 435 | 434 | 434 | | 128 | 32768 | 512 | 6.00 | 2.18 | 0.99 | 0.40 | 334 | 336 | 330 | 332 | | 128 | 7168 | 16384 | 6.24 | 2.40 | 1.18 | 0.59 | 534 | 534 | 534 | 533 | | 128 | 4096 | 7168 | 6.04 | 2.21 | 1.03 | 0.43 | 370 | 371 | 370 | 371 | | 128 | 7168 | 2048 | 6.19 | 2.38 | 1.19 | 0.62 | 332 | 334 | 331 | 332 | | 4096 | 2112 | 7168 | 6.09 | 2.30 | 1.08 | 0.49 | 1053 | 1058 | 1113 | 1114 | | 4096 | 24576 | 1536 | 6.14 | 2.33 | 1.10 | 0.50 | 1288 | 1286 | 1290 | 1289 | | 4096 | 32768 | 512 | 6.23 | 2.41 | 1.21 | 0.61 | 913 | 912 | 912 | 911 | | 4096 | 7168 | 16384 | 6.31 | 2.51 | 1.27 | 0.68 | 1524 | 1524 | 1459 | 1458 | | 4096 | 4096 | 7168 | 6.28 | 2.47 | 1.27 | 0.68 | 1446 | 1445 | 1394 | 1396 | | 4096 | 7168 | 2048 | 6.18 | 2.40 | 1.21 | 0.61 | 1240 | 1243 | 1234 | 1237 | Note that there is some perf drop when using NVRTC due to a known bug of NVRTC which leads to extra instructions (but in the `m=4096,n=2112,k=7168` case, NVRTC version was faster, which was a bit strange). So NVCC is kept as the default compiler for now, while NVRTC can be enabled with extra env var `DG_JIT_USE_NVRTC`. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
bsakaguchi 2026-10-26 21:15:16 +00:00
Sign in to join this conversation.
No Branch/Tag specified
Branches Tags
main
public-sync-20261014
0930-release
nv_dev
fix-mega-moe-fence
public-release-2609
public-release-2607
public-release-2606
mega-update
public-release-260416
ko3n1g/ci/build-wheels
nv_dev_f8e8fb5
nv_dev_950d31a
nv_dev_67fc648
nv_dev_c491439
nv_dev_4ff3f54
v2.1.1.post3
v2.1.1.post2
v2.1.1.post1
v2.1.1
v2.1.0
Labels
Clear labels   
pull-request
Mirrored from GitHub Pull Request

No labels
pull-request
Milestone
Clear milestone
No items
No milestone
Projects
Clear projects
No items
No project
Assignees
Clear assignees
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#243
No description provided.