Public release 26/09/30 by LyricZhao · Pull Request #462 · deepseek-ai/DeepGEMM

LyricZhao · 2026-09-30T01:36:28Z

This update introduces new features, performance improvements, and important API changes:

New features

Locality-aware Mega MoE: optional weight localization and domain-aware scheduling, with helpers for localized allocation (via MLOPart) and SM locality discovery. Saving 100W power and boosting 5-10% GPU frequency for heavy workload.
New deep_gemm.epilogue classes: Identity, Alpha, FP8Quantization, and opt-in BF16StochasticRounding for direct BF16 output on SM100.
BF16 batched einsums now support accumulation and FP32 output for bhr,hdr->bhd and bhd,hdr->bhr.
Sparse Indexer now supports head counts divisible by 4, from 4 through 32, for both contiguous and paged MXFP4/MXFP8 inputs.
Support for padded K strides in packed UE8M0 scaling factors.

Performance improvements

Improved Lightning Indexer kernels for V4.1, with better scheduling, KV pipelines, and head reduction for prefill and paged decoding.
Mega MoE epilogue prefetch and shared memory optimizations enable deeper pipelines for FP8xFP4 prefill.

Fixes

Fixed Sparse Indexer memory bounds and strengthened paged-cache and metadata validation.
Fixed Mega MoE combine reduction order for bitwise agreement with the shared-expert reference, and improved locality probing and fallback behavior.
Fixed runtime shutdown ordering for cuBLASLt resources and a missing FP8 header in Mega mHC.
Removed the Mega mHC test baseline's dependency on external tile-kernels code.

Breaking changes

The non-paged fp8_fp4_mqa_logits API now requires a positive max_seqlen_k and returns compressed logits. The dense non-paged and paged MQA APIs no longer accept clean_logits or logits_dtype; callers must mask entries beyond each row's valid KV span.
Dense MQA on SM100 now requires MXFP8/MXFP4 inputs with packed UE8M0 scales and BF16 weights, and returns BF16 logits. SM90 retains non-MX FP8 inputs with FP32 weights and logits. Weight rows must be 16-byte aligned.
Paged MQA on SM100 now requires the variable-length interface: pass indices and use next_n=1.
Removed the legacy fp8_mqa_logits, fp8_paged_mqa_logits, and fp8_gemm_nt_skip_head_mid APIs. Use the fp8_fp4_*mqa_logits APIs for MQA workloads.
Mega Gate now supports only scoring_func="sqrtsoftplus", which is also the default. get_bf16_mega_gate_config has been removed.

Related release

DeepGEMM Ascend is now available in a separate repository.

Contributors

Locality-aware Mega MoE and SM probing: @Yongqi-Zhuo @LyricZhao
Lightning Indexer, Sparse Indexer, and Mega MoE correctness: @zheanxu
GEMM epilogue classes, stochastic rounding, and BF16 einsum extensions: @xyf007
Mega MoE epilogue prefetch and shared-memory optimizations: @THUTwenTy
Packed scaling-factor strides: @xay5421

Add locality-aware Mega MoE execution, GEMM epilogue classes with BF16 stochastic rounding, wider sparse MQA head support, and packed SF strides. Improve Indexer and Mega MoE pipelines, update attention and Mega Gate interfaces, and include runtime and correctness fixes. Export the current source through the OSS filter while retaining public release history and the Mega MoE task-info release-ordering fix. Co-authored-by: Anyi Xu <41318296+xay5421@users.noreply.github.com> Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com> Co-authored-by: huahuaquq <289927700+huahuaquq@users.noreply.github.com> Co-authored-by: TwenTy_ <47315699+THUTwenTy@users.noreply.github.com> Co-authored-by: xyf007 <41227294+xyf007@users.noreply.github.com> Co-authored-by: Yongqi Zhuo <Yongqi-Zhuo@users.noreply.github.com> Co-authored-by: yukuai26 <1010238651@qq.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com>

GSM2017PMK-OSV · 2026-10-02T17:03:05Z

Ое

…stograms, next-n 1..6 Port the BF16 paged MQA producer of deepseek-ai#462 (merge 057ca59: MXFP4 Q/KV, BF16 head weights, BF16 head reduction and BF16 logits) as an isolated `mqa_bf16` kernel with its own split-384 cost-balanced schedule, and let it emit the same coarse histogram as the FP32 path. - APIs: `get_paged_mqa_logits_bf16_metadata(context_lens, 128, num_sms, indices=, tokens_per_request=)` and `fp4_paged_mqa_logits_bf16(..., indices=, histogram=, tokens_per_request=)` (pybind, TVM-FFI, deep_gemm / sgl_deep_gemm). - Adjacent equal request IDs share KV loads; the CPU hint selects Q4 / 3 TMEM stages (1..4) or Q6 / 2 stages (5..6; 6 also XOR-swizzles the histogram). Q4 keeps deepseek-ai#462's static TMEM stage mapping (its no-histogram kernel is SASS-identical to deepseek-ai#462); Q6 rotates a two-stage TMEM ring. - Exact histogram: every live non-NaN stored BF16 score counted once per row (16-bit shared counters for a window of rows, packed fast path for 2^-14 <= |x| < 16, published with packed 64-bit global atomics). The histogram kernels' scheduler caps a task's splits so the 16-bit overflow guard runs once per task; Q6 counts a 12-row window with 6 KV stages, so a CTA moving to the next request does not publish mid-kernel. Logits are bit-identical with and without the histogram. - Shared with the FP32 path: the coarse-bin encoding (epilogue/ coarse_histogram.cuh, plus a BF16 overload) and NoHistogram; dev's packing, ring pipeline (new parity-phase option for steps longer than a 2-stage ring), math / PTX / UMMA descriptor helpers. Only the kernel core, its shared-memory layout, its scheduler and the histogram accumulator are BF16-specific. - Test: sgl_deep_gemm/tests/test_mqa_bf16.py (n = 1..6, PDL off / on, exact counts, accumulation, CUDA-graph replay with changing lengths and request IDs). Replaces sgl-project#96.

Bring in locality-aware Mega MoE, the GEMM epilogue classes with BF16 stochastic rounding, the reworked Lightning/Sparse Indexer kernels, and upstream's MQA API changes. Adapt the new upstream APIs to the TORCH_LIBRARY bindings: - Epilogue classes are a `deep_gemm::Epilogue` custom class built by factory ops; `deep_gemm.epilogue.X(...)` returns one and every GEMM and einsum op takes it as `epilogue=None`. - Locality domain helpers are `locality_domain_*` ops, exposed as `deep_gemm._C.locality_domain` for `deep_gemm.locality_domain`. - MLOPart calls into Python through the CPython limited API instead of pybind11, taking the GIL itself. - `shutdown` is a config op. - Python errors from MLOPart surface as `RuntimeError` through the dispatcher, which `test_locality_domain` now accepts. Adopt upstream's MQA API: drop `fp8_mqa_logits`, `fp8_paged_mqa_logits`, `fp8_gemm_nt_skip_head_mid`, `clean_logits` and `logits_dtype`, and always return compressed logits. SM120 keeps its kernels unchanged: FP8 inputs with per-token float KV scales or MXFP4, FP32 weights and FP32 logits, fixed-length or varlen paged queries. SM90 keeps 32-row pages and next_n=4. Keep dev's other changes on top of upstream: - SM120 GEMMs move to `epilogue::operators` and lose the head-split index remap along with `fp8_gemm_nt_skip_head_mid`. - The SiTU Mega MoE variant is re-derived from upstream's kernel, and the NVFP4 Mega MoE uses the new scheduler with one locality domain. - SwiGLU alpha/beta is re-applied on top of upstream's prefetching epilogue. - Mega Gate keeps the removal of the finite-score device trap. - `get_bf16_mega_gate_config` is removed, following upstream. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

Add locality-aware Mega MoE execution, GEMM epilogue classes with BF16 stochastic rounding, wider sparse MQA head support, and packed SF strides. Improve Indexer and Mega MoE pipelines, update attention and Mega Gate interfaces, and include runtime and correctness fixes. Export the current source through the OSS filter while retaining public release history and the Mega MoE task-info release-ordering fix. Co-authored-by: Chenggang Zhao <chenggangz@deepseek.com> Co-authored-by: Anyi Xu <41318296+xay5421@users.noreply.github.com> Co-authored-by: huahuaquq <289927700+huahuaquq@users.noreply.github.com> Co-authored-by: TwenTy_ <47315699+THUTwenTy@users.noreply.github.com> Co-authored-by: xyf007 <41227294+xyf007@users.noreply.github.com> Co-authored-by: Yongqi Zhuo <Yongqi-Zhuo@users.noreply.github.com> Co-authored-by: yukuai26 <1010238651@qq.com> Co-authored-by: Zhean Xu <94977922+zheanxu@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>

…stograms, next-n 1..6 Port the BF16 paged MQA producer of deepseek-ai#462 (merge 057ca59: MXFP4 Q/KV, BF16 head weights, BF16 head reduction and BF16 logits) as an isolated `mqa_bf16` kernel with its own split-384 cost-balanced schedule, and let it emit the same coarse histogram as the FP32 path. - APIs: `get_paged_mqa_logits_bf16_metadata(context_lens, 128, num_sms, indices=, tokens_per_request=)` and `fp4_paged_mqa_logits_bf16(..., indices=, histogram=, tokens_per_request=)`: TORCH_LIBRARY ops in csrc/torch_library_api.cpp (`Tensor(a!)? histogram`, no schema defaults), deep_gemm / sgl_deep_gemm wrappers that own the defaults, and pybind. - Adjacent equal request IDs share KV loads; the CPU hint selects Q4 / 3 TMEM stages (1..4) or Q6 / 2 stages (5..6; 6 also XOR-swizzles the histogram). Q4 keeps deepseek-ai#462's static TMEM stage mapping (its no-histogram kernel is SASS-identical to deepseek-ai#462); Q6 rotates a two-stage TMEM ring. Each stage count is the deepest ring that fits (TMEM columns, SM100 shared memory with the histogram window); split 384 and the 3 x 128 math threads are named `Config` constants. The schedule (one CTA, two int32 per row in shared memory) checks its documented limit of 16384 rows. - Exact histogram: every live non-NaN stored BF16 score counted once per row (16-bit shared counters for a window of rows, packed fast path for 2^-14 <= |x| < 16, published with packed 64-bit global atomics). The histogram kernels' scheduler caps a task's splits so the 16-bit overflow guard runs once per task; Q6 counts a 12-row window with 6 KV stages, so a CTA moving to the next request does not publish mid-kernel. Logits are bit-identical with and without the histogram. - Shared with the FP32 path: the coarse-bin encoding (epilogue/ coarse_histogram.cuh, plus a BF16 overload) and NoHistogram; dev's packing, ring pipeline (new parity-phase option for steps longer than a 2-stage ring), math / PTX / UMMA descriptor helpers. Only the kernel core, its shared-memory layout, its scheduler and the histogram accumulator are BF16-specific. - Test: sgl_deep_gemm/tests/test_mqa_bf16.py, a plain script that run_tests.sh runs on SM100 (n = 1..6, PDL off / on, exact counts, accumulation, CUDA-graph replay with changing lengths and request IDs; the schedule's row limit). Replaces sgl-project#96. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

…-ai#471) into dev Two-parent merge of upstream 1fc9c66 (deepseek-ai#441, deepseek-ai#462, deepseek-ai#471) into sgl dev b296ce3, keeping the SGL wheel, the TORCH_LIBRARY binding and the downstream kernels. - Mega MoE: upstream locality-domain scheduling, 4D weight TMA, L1/L2 store blocks and the TMEM prefetch epilogue, with NVFP4/MXFP4, TRT-LLM weights, alphas, SiTU and FP8 combine carried over. TRT weights stay unlocalized. - GEMM epilogue classes threaded through every dispatch; SM90/SM120 assert that no custom epilogue is set. SM120 kernels drop the head-split epilogue parameter, as upstream did for SM90. - Attention takes upstream's API as-is: no clean_logits/logits_dtype, no legacy fp8_mqa_logits/fp8_paged_mqa_logits, no fp8_gemm_nt_skip_head_mid. The SM120 MQA kernels that only served the old API are removed. - Mega gate is sqrtsoftplus-only; bf16_mega_gate defaults to it and get_bf16_mega_gate_config is gone. - TORCH_LIBRARY binding: epilogue classes (Identity, Alpha, FP8CastWithClamp, BF16StochasticRounding) cross as (kind, scalar, sfd, count) into every GEMM/einsum op; shutdown and the locality-domain ops are registered; MLOPart reaches Python through ctypes callbacks installed by deep_gemm/locality_domain.py, which re-raises the callback's exception from the allocating call. - sgl_deep_gemm tests pick up upstream's test changes and the new test_locality_domain.py. - Dev's per-expert post-norm (#104) sits on upstream's restructured epilogue and combine. Its store block of 48 at block_m 192 becomes the L1 store block, upstream's FP8xFP4 L1/L2 blocks at 240 are kept, and DG_MEGA_MOE_TILES sets both L1 and L2. - deep_jit stays at dev's 3732a3b, a descendant of upstream's 2efdab4. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

LyricZhao merged commit 057ca59 into main Sep 30, 2026
2 of 3 checks passed

yuweih205 mentioned this pull request Sep 30, 2026

Add SM90 BF16 grouped GEMM with gathered activation rows #463

Open

XFDG mentioned this pull request Oct 1, 2026

[Bugfix][Layout] Add cuda_fp8.h include to fix __nv_fp8_e4m3 undefined #439

Closed

Oasis-Git mentioned this pull request Oct 2, 2026

[Roadmap] DeepSeek V4.1 Optimization sgl-project/sglang#42170

Open

17 of 56 tasks

This was referenced Oct 3, 2026

[SM100] Add exact BF16 paged MQA histograms for next-n 1–6 sgl-project/DeepGEMM#96

Closed

[DSv4.1] Add opt-in LiteTopK decode top-512 path for SM100 sgl-project/sglang#42351

Closed

Heisenberg-Yin mentioned this pull request Oct 9, 2026

[SM100] Paged MQA logits: optional exact coarse score histogram for LiteTopK decode sgl-project/DeepGEMM#94

Open

zyongye mentioned this pull request Oct 10, 2026

Merge upstream public release 26/09/30 (deepseek-ai/DeepGEMM#462) vllm-project/DeepGEMM#30

Merged

zyongye mentioned this pull request Oct 12, 2026

[Kernel][Build] Bump DeepGEMM to the 26/09/30 release and adapt the sparse indexer vllm-project/vllm#61149

Merged

6 tasks done

zcnrex mentioned this pull request Oct 14, 2026

Merge upstream main (26/09/30 release + #471) into dev sgl-project/DeepGEMM#103

Merged

1 of 2 tasks

Heisenberg-Yin mentioned this pull request Oct 16, 2026

[Perf] [NVIDIA] attention: Add opt-in LiteTopK decode top-k for DSA and DeepSeek-V4.1 on SM100 sgl-project/sglang#40764

Open

linmuchuiyang mentioned this pull request Oct 19, 2026

Support DeepSeek V4.1 Flash NVIDIA-NeMo/RL#4246

Open

YLGH mentioned this pull request Oct 24, 2026

Sync upstream through PR #478: public release 26/09/30, FP8 epilogues, mHC epoch barriers fw-ai/DeepGEMM#22

Merged

Navigation Menu

Repository navigation

Public release 26/09/30 - #462

Public release 26/09/30#462
LyricZhao merged 1 commit into
mainfrom
0930-release

LyricZhao commented Sep 30, 2026 •

edited

Loading

Uh oh!

Uh oh!

GSM2017PMK-OSV commented Oct 2, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

LyricZhao commented Sep 30, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

New features

Performance improvements

Fixes

Breaking changes

Related release

Contributors

Uh oh!

Uh oh!

GSM2017PMK-OSV commented Oct 2, 2026

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

LyricZhao commented Sep 30, 2026 •

edited

Loading