[PR #471] [MERGED] Public release 10/14 #473

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/471
Author: @LyricZhao
Created: 10/14/2026
Status: ✅ Merged
Merged: 10/14/2026
Merged by: @LyricZhao

Base: main ← Head: public-sync-20261014

📝 Commits (1)

d810967 Update FP8 epilogues and paged sparse MQA alignment

📊 Changes

18 files changed (+380 additions, -263 deletions)

View changed files

📝 csrc/apis/attention.hpp (+1 -1)
📝 csrc/apis/einsum.hpp (+2 -2)
📝 csrc/apis/epilogue_class.hpp (+9 -3)
📝 csrc/jit_kernels/impls/epilogue_class.hpp (+31 -5)
📝 deep_gemm/__init__.py (+1 -1)
📝 deep_gemm/include/deep_gemm/common/math.cuh (+10 -4)
📝 deep_gemm/include/deep_gemm/common/types.cuh (+5 -2)
📝 deep_gemm/include/deep_gemm/epilogue/operators.cuh (+157 -29)
📝 deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh (+30 -61)
📝 deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh (+60 -109)
📝 deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh (+15 -13)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh (+16 -14)
📝 deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh (+1 -1)
📝 deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh (+2 -2)
📝 deep_gemm/include/deep_gemm/ptx/ld_st.cuh (+10 -0)
📝 tests/test_attention.py (+1 -1)
📝 tests/test_einsum.py (+27 -14)
📝 tests/test_mega_moe.py (+2 -1)

📄 Description

DeepGEMM 2.8.2

This update adds fused FP8 clamping and statistics, improves GEMM epilogue conversion, and expands paged Sparse Indexer layout support.

New features

Fused FP8 cast with optional clamping: deep_gemm.epilogue.FP8CastWithClamp converts BF16-rounded GEMM output to FP8, with optional symmetric clamping before scaling.
Optional clamped-element counts: Pass a single-element CUDA torch.uint32 tensor as clamped_count to collect the total asynchronously. Each invocation overwrites the previous count.
More flexible paged Sparse Indexer layouts: Paged sparse KV caches now accept page strides aligned to 128 bytes, relaxing the previous 512-byte requirement.

Performance improvements

Use packed FP8 conversion and a packed BF16 clamp instruction to improve fused conversion efficiency.
Consolidate epilogue conversion paths for FP8 and BF16 outputs, including stochastic rounding.
Improve swap-AB epilogue scheduling by preparing scaling factors for both TMEM fragments before FP8 conversion.
Reduce clamped-element counts within each warp before updating the global counter.

Fixes

Keep the Mega MoE weight localizer alive across correctness iterations and destroy it during final test teardown, avoiding repeated MLOPart and MPS restarts.
Update paged Sparse Indexer validation and tests to use the relaxed 128-byte alignment requirement.
Extend einsum correctness coverage for clamped FP8 output, scaling factors, and counter overwrite behavior.

Breaking changes

deep_gemm.epilogue.FP8Quantization has been replaced by FP8CastWithClamp. Existing callers should update the class name. Omitting clamp_value preserves unclamped conversion.
The FP8 (output, sfd) output-pair shorthand remains supported and uses unclamped FP8CastWithClamp.

epilogue = deep_gemm.epilogue.FP8CastWithClamp(
    sfd=sfd,
    clamp_value=clamp_value,      # Optional; None disables clamping
    clamped_count=clamped_count,  # Optional CUDA uint32 tensor with one element
)

Contributors

Fused FP8 cast, clamping, statistics, and epilogue optimizations: @xyf007
Mega MoE localizer lifecycle fix: @yukuai26
Paged Sparse Indexer alignment and release version update: @zheanxu

_{🔄 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/471 **Author:** [@LyricZhao](https://github.com/LyricZhao) **Created:** 10/14/2026 **Status:** ✅ Merged **Merged:** 10/14/2026 **Merged by:** [@LyricZhao](https://github.com/LyricZhao) **Base:** `main` ← **Head:** `public-sync-20261014` --- ### 📝 Commits (1) - [`d810967`](https://github.com/deepseek-ai/DeepGEMM/commit/d810967a1158d5fbfd4a5bf8145e75223a831e41) Update FP8 epilogues and paged sparse MQA alignment ### 📊 Changes **18 files changed** (+380 additions, -263 deletions) <details> <summary>View changed files</summary> 📝 `csrc/apis/attention.hpp` (+1 -1) 📝 `csrc/apis/einsum.hpp` (+2 -2) 📝 `csrc/apis/epilogue_class.hpp` (+9 -3) 📝 `csrc/jit_kernels/impls/epilogue_class.hpp` (+31 -5) 📝 `deep_gemm/__init__.py` (+1 -1) 📝 `deep_gemm/include/deep_gemm/common/math.cuh` (+10 -4) 📝 `deep_gemm/include/deep_gemm/common/types.cuh` (+5 -2) 📝 `deep_gemm/include/deep_gemm/epilogue/operators.cuh` (+157 -29) 📝 `deep_gemm/include/deep_gemm/epilogue/sm100_store_cd.cuh` (+30 -61) 📝 `deep_gemm/include/deep_gemm/epilogue/sm100_store_cd_swap_ab.cuh` (+60 -109) 📝 `deep_gemm/include/deep_gemm/impls/sm100_bf16_gemm.cuh` (+15 -13) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_gemm_1d1d.cuh` (+16 -14) 📝 `deep_gemm/include/deep_gemm/impls/sm100_sparse_mqa_logits.cuh` (+1 -1) 📝 `deep_gemm/include/deep_gemm/impls/sm90_bf16_gemm.cuh` (+2 -2) 📝 `deep_gemm/include/deep_gemm/ptx/ld_st.cuh` (+10 -0) 📝 `tests/test_attention.py` (+1 -1) 📝 `tests/test_einsum.py` (+27 -14) 📝 `tests/test_mega_moe.py` (+2 -1) </details> ### 📄 Description # DeepGEMM 2.8.2 This update adds fused FP8 clamping and statistics, improves GEMM epilogue conversion, and expands paged Sparse Indexer layout support. ## New features - Fused FP8 cast with optional clamping: deep_gemm.epilogue.FP8CastWithClamp converts BF16-rounded GEMM output to FP8, with optional symmetric clamping before scaling. - Optional clamped-element counts: Pass a single-element CUDA torch.uint32 tensor as clamped_count to collect the total asynchronously. Each invocation overwrites the previous count. - More flexible paged Sparse Indexer layouts: Paged sparse KV caches now accept page strides aligned to 128 bytes, relaxing the previous 512-byte requirement. ## Performance improvements - Use packed FP8 conversion and a packed BF16 clamp instruction to improve fused conversion efficiency. - Consolidate epilogue conversion paths for FP8 and BF16 outputs, including stochastic rounding. - Improve swap-AB epilogue scheduling by preparing scaling factors for both TMEM fragments before FP8 conversion. - Reduce clamped-element counts within each warp before updating the global counter. ## Fixes - Keep the Mega MoE weight localizer alive across correctness iterations and destroy it during final test teardown, avoiding repeated MLOPart and MPS restarts. - Update paged Sparse Indexer validation and tests to use the relaxed 128-byte alignment requirement. - Extend einsum correctness coverage for clamped FP8 output, scaling factors, and counter overwrite behavior. ## Breaking changes - deep_gemm.epilogue.FP8Quantization has been replaced by FP8CastWithClamp. Existing callers should update the class name. Omitting clamp_value preserves unclamped conversion. - The FP8 (output, sfd) output-pair shorthand remains supported and uses unclamped FP8CastWithClamp. ``` epilogue = deep_gemm.epilogue.FP8CastWithClamp( sfd=sfd, clamp_value=clamp_value, # Optional; None disables clamping clamped_count=clamped_count, # Optional CUDA uint32 tensor with one element ) ``` ## Contributors - Fused FP8 cast, clamping, statistics, and epilogue optimizations: @xyf007 - Mega MoE localizer lifecycle fix: @yukuai26 - Paged Sparse Indexer alignment and release version update: @zheanxu --- <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:17 +00:00