deepseek-ai/DeepGEMM

Fork 0

mirror of https://github.com/deepseek-ai/DeepGEMM.git synced 2026-11-05 00:47:26 +00:00

[PR #458] [CLOSED] Merge main (Public Release 26/09) into nv_dev, with sparse MQA-logits integration fixes #465

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/458
Author: @lfr-0531
Created: 9/23/2026
Status: ❌ Closed

Base: nv_dev ← Head: pr/nv-dev-merge-main-2609

📝 Commits (6)

66081d4 Public Release 26/09 (#432)
78b6900 Fix Mega MoE task info slot release ordering (#441)
0274ad7 Merge upstream main (Public Release 26/09 #432, #441) into nv_dev
1364fed Reuse a warmed sparse MQA logits workspace during graph capture
7a85767 Return a typed map from get_bf16_mega_gate_config
d0a7e77 Accept 16-byte aligned KV pages in the paged sparse MQA logits

📊 Changes

141 files changed (+13002 additions, -6247 deletions)

View changed files

📝 .gitmodules (+3 -3)
📝 CMakeLists.txt (+4 -2)
📝 README.md (+20 -22)
📝 csrc/apis/attention.hpp (+262 -45)
📝 csrc/apis/config.hpp (+17 -21)
📝 csrc/apis/einsum.hpp (+88 -43)
📝 csrc/apis/gemm.hpp (+242 -67)
📝 csrc/apis/hyperconnection.hpp (+1 -7)
📝 csrc/apis/layout.hpp (+6 -15)
➕ csrc/apis/mega_gate.hpp (+238 -0)
➕ csrc/apis/mega_mhc.hpp (+287 -0)
📝 csrc/apis/mega_moe.hpp (+13 -16)
📝 csrc/apis/sm90_mega.hpp (+5 -9)
📝 csrc/indexing/main.cu (+1 -0)
➖ csrc/jit/cache.hpp (+0 -31)
➖ csrc/jit/compiler.hpp (+0 -369)
➖ csrc/jit/device_runtime.hpp (+0 -153)
➖ csrc/jit/handle.hpp (+0 -222)
➖ csrc/jit/include_parser.hpp (+0 -80)
➖ csrc/jit/kernel_runtime.hpp (+0 -165)

...and 80 more files

📄 Description

Summary

Bring nv_dev onto the DeepJIT-based runtime of the 26/09 public release (main 78b6900 = #432 + #441) while keeping every nv_dev feature, plus three follow-up fixes found while integrating the sparse MQA-logits kernels into TensorRT-LLM.

Merge of `main` (26/09) into `nv_dev`

Policy of the merge: main's structure is the base and every nv_dev feature is re-applied on top of it.

Now on this branch from main:

DeepJIT submodule replaces csrc/jit/* and fmt (jit->compile / jit->launch, C++20, std::format); CUDA >= 12.9
Sparse MQA logits for the DeepSeek-V4.1 hierarchical sparse indexer (fp8_fp4_sparse_mqa_logits, fp8_fp4_paged_sparse_mqa_logits and their metadata builders)
MegaGate, Mega mHC, fp8xfp8 MegaMoE weights, MegaMoESignals workspace, L2 readiness mask, GEMM alpha / deterministic paths, scheduled MQA metadata
#441 task-info slot release ordering fence in the MegaMoE scheduler

nv_dev features re-applied on the new runtime:

SM120/SM121: FP8/FP4 GEMM (AB-swap, split-K), BF16 GEMM, bmk/bnk einsum, TF32 HC prenorm GEMM, MQA logits (contiguous + paged), per-arch contiguous M/K alignment, packed UE8M0 SF on SM120
NVFP4 (fp4xfp4) MegaMoE with BF16 shared experts (MmaKind::NVFP4 kept next to upstream's MmaKind::MXF4; bit-based buffer layout; legacy l2_full_count signal for the count protocol)
SiTU activation for the FP8xFP4 MegaMoE
SM90 FP8 MegaMoE
FP16-weights SM100 MQA logits (two-CTA FP16-accumulator kernel)
SM90 paged MQA logits: block_kv 32, varlen metadata, next_n=4 multicast clusters
Standalone smxx_clean_logits kept for the kernels that do not fuse cleaning (SM120 and FP16-weights paths)
Barrier timeout policy (Diagnostic / TrapOnly) and the configurable DG_JIT_BARRIER_TIMEOUT_SECONDS
Tests for all of the above merged into the upstream test files

Follow-up fixes

get_sparse_mqa_logits_workspace: frameworks capture CUDA graphs on a dedicated side stream but replay on the launching stream, so the first sparse-logits call on the capture stream used to trip the pre-capture assertion. While capturing, reuse a workspace already warmed up on the device.
get_bf16_mega_gate_config: return std::map<std::string, int> instead of py::dict so generated Python stubs are dict[str, int] (strict type checkers reject a bare dict); the Python-side result is unchanged.
Paged sparse MQA logits: accept any 16-byte aligned KV page stride. The kernel addresses pages directly and copies them with 16-byte cp.async chunks, so page-footer caches that are not padded to 512 bytes (for example 64 x (64 + 4) = 4352-byte pages) can be read in place.

Validation

GB300 (NGC PyTorch 26.08, CUDA 13.4), standalone test suites: test_attention (contiguous MQA logits full, paged sampled cases, sparse MQA logits 161/161 bitwise vs dense), test_fp8_fp4, test_bf16, test_einsum, test_layout, test_hyperconnection, test_mega_moe (4 ranks), test_nvfp4_mega_moe, test_mega_moe_situ. SM90 and SM120 paths compile but were not run. test_mega_gate / test_mega_mhc need tilelang, which was not available in that environment.

The branch is consumed by TensorRT-LLM's DeepSeek-V4.1 bring-up (CSA2 indexer using the sparse MQA-logits kernels for the candidate-consuming layers), where the DeepGEMM-side changes above were exercised end to end (unit suite, GSM8K, long-context throughput on GB300).

Note on DeepJIT

third-party/deep_jit stays at upstream e5bdee2. Its utils/exception.hpp includes <elfutils/libdwfl.h> for declarations only (libdw is dlopened at runtime), so build hosts need the elfutils development headers; making that include optional (__has_include with a symbol-only backtrace fallback) is proposed separately to DeepJIT.

_{🔄 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/458 **Author:** [@lfr-0531](https://github.com/lfr-0531) **Created:** 9/23/2026 **Status:** ❌ Closed **Base:** `nv_dev` ← **Head:** `pr/nv-dev-merge-main-2609` --- ### 📝 Commits (6) - [`66081d4`](https://github.com/deepseek-ai/DeepGEMM/commit/66081d4c9c7d7c44f13fea402e5b622aa0f409c2) Public Release 26/09 (#432) - [`78b6900`](https://github.com/deepseek-ai/DeepGEMM/commit/78b69000794d0937b47ae3387eff7663410264d1) Fix Mega MoE task info slot release ordering (#441) - [`0274ad7`](https://github.com/deepseek-ai/DeepGEMM/commit/0274ad7b911b99b4eb6f8d796403811cf1e18f3d) Merge upstream main (Public Release 26/09 #432, #441) into nv_dev - [`1364fed`](https://github.com/deepseek-ai/DeepGEMM/commit/1364fed0f23ef7b5c99a10bd0b25f846af5848c7) Reuse a warmed sparse MQA logits workspace during graph capture - [`7a85767`](https://github.com/deepseek-ai/DeepGEMM/commit/7a85767db7509d10e6348c00585179c414ab9f8b) Return a typed map from get_bf16_mega_gate_config - [`d0a7e77`](https://github.com/deepseek-ai/DeepGEMM/commit/d0a7e771f99fea7cc27513e46059de4f8761cc24) Accept 16-byte aligned KV pages in the paged sparse MQA logits ### 📊 Changes **141 files changed** (+13002 additions, -6247 deletions) <details> <summary>View changed files</summary> 📝 `.gitmodules` (+3 -3) 📝 `CMakeLists.txt` (+4 -2) 📝 `README.md` (+20 -22) 📝 `csrc/apis/attention.hpp` (+262 -45) 📝 `csrc/apis/config.hpp` (+17 -21) 📝 `csrc/apis/einsum.hpp` (+88 -43) 📝 `csrc/apis/gemm.hpp` (+242 -67) 📝 `csrc/apis/hyperconnection.hpp` (+1 -7) 📝 `csrc/apis/layout.hpp` (+6 -15) ➕ `csrc/apis/mega_gate.hpp` (+238 -0) ➕ `csrc/apis/mega_mhc.hpp` (+287 -0) 📝 `csrc/apis/mega_moe.hpp` (+13 -16) 📝 `csrc/apis/sm90_mega.hpp` (+5 -9) 📝 `csrc/indexing/main.cu` (+1 -0) ➖ `csrc/jit/cache.hpp` (+0 -31) ➖ `csrc/jit/compiler.hpp` (+0 -369) ➖ `csrc/jit/device_runtime.hpp` (+0 -153) ➖ `csrc/jit/handle.hpp` (+0 -222) ➖ `csrc/jit/include_parser.hpp` (+0 -80) ➖ `csrc/jit/kernel_runtime.hpp` (+0 -165) _...and 80 more files_ </details> ### 📄 Description ## Summary Bring `nv_dev` onto the DeepJIT-based runtime of the 26/09 public release (`main` 78b6900 = #432 + #441) while keeping every `nv_dev` feature, plus three follow-up fixes found while integrating the sparse MQA-logits kernels into TensorRT-LLM. ### Merge of `main` (26/09) into `nv_dev` Policy of the merge: `main`'s structure is the base and every `nv_dev` feature is re-applied on top of it. Now on this branch from `main`: - DeepJIT submodule replaces `csrc/jit/*` and `fmt` (`jit->compile` / `jit->launch`, C++20, `std::format`); CUDA >= 12.9 - Sparse MQA logits for the DeepSeek-V4.1 hierarchical sparse indexer (`fp8_fp4_sparse_mqa_logits`, `fp8_fp4_paged_sparse_mqa_logits` and their metadata builders) - MegaGate, Mega mHC, fp8xfp8 MegaMoE weights, `MegaMoESignals` workspace, L2 readiness mask, GEMM alpha / deterministic paths, scheduled MQA metadata - #441 task-info slot release ordering fence in the MegaMoE scheduler `nv_dev` features re-applied on the new runtime: - SM120/SM121: FP8/FP4 GEMM (AB-swap, split-K), BF16 GEMM, bmk/bnk einsum, TF32 HC prenorm GEMM, MQA logits (contiguous + paged), per-arch contiguous M/K alignment, packed UE8M0 SF on SM120 - NVFP4 (fp4xfp4) MegaMoE with BF16 shared experts (`MmaKind::NVFP4` kept next to upstream's `MmaKind::MXF4`; bit-based buffer layout; legacy `l2_full_count` signal for the count protocol) - SiTU activation for the FP8xFP4 MegaMoE - SM90 FP8 MegaMoE - FP16-weights SM100 MQA logits (two-CTA FP16-accumulator kernel) - SM90 paged MQA logits: block_kv 32, varlen metadata, next_n=4 multicast clusters - Standalone `smxx_clean_logits` kept for the kernels that do not fuse cleaning (SM120 and FP16-weights paths) - Barrier timeout policy (Diagnostic / TrapOnly) and the configurable `DG_JIT_BARRIER_TIMEOUT_SECONDS` - Tests for all of the above merged into the upstream test files ### Follow-up fixes - `get_sparse_mqa_logits_workspace`: frameworks capture CUDA graphs on a dedicated side stream but replay on the launching stream, so the first sparse-logits call on the capture stream used to trip the pre-capture assertion. While capturing, reuse a workspace already warmed up on the device. - `get_bf16_mega_gate_config`: return `std::map<std::string, int>` instead of `py::dict` so generated Python stubs are `dict[str, int]` (strict type checkers reject a bare `dict`); the Python-side result is unchanged. - Paged sparse MQA logits: accept any 16-byte aligned KV page stride. The kernel addresses pages directly and copies them with 16-byte `cp.async` chunks, so page-footer caches that are not padded to 512 bytes (for example 64 x (64 + 4) = 4352-byte pages) can be read in place. ## Validation GB300 (NGC PyTorch 26.08, CUDA 13.4), standalone test suites: `test_attention` (contiguous MQA logits full, paged sampled cases, sparse MQA logits 161/161 bitwise vs dense), `test_fp8_fp4`, `test_bf16`, `test_einsum`, `test_layout`, `test_hyperconnection`, `test_mega_moe` (4 ranks), `test_nvfp4_mega_moe`, `test_mega_moe_situ`. SM90 and SM120 paths compile but were not run. `test_mega_gate` / `test_mega_mhc` need `tilelang`, which was not available in that environment. The branch is consumed by TensorRT-LLM's DeepSeek-V4.1 bring-up (CSA2 indexer using the sparse MQA-logits kernels for the candidate-consuming layers), where the DeepGEMM-side changes above were exercised end to end (unit suite, GSM8K, long-context throughput on GB300). ## Note on DeepJIT `third-party/deep_jit` stays at upstream e5bdee2. Its `utils/exception.hpp` includes `<elfutils/libdwfl.h>` for declarations only (libdw is `dlopen`ed at runtime), so build hosts need the elfutils development headers; making that include optional (`__has_include` with a symbol-only backtrace fallback) is proposed separately to DeepJIT. --- <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:08 +00:00

closed this issue
added the
pull-request
label

bsakaguchi referenced this issue

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

[PR #465] Update copyright year in LICENSE file #469

No labels

pull-request

No milestone

No project

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#465

No description provided.

Rows
Columns