[PR #325] [CLOSED] feat: add epilogue warpgroup to SM90 FP8 MQA logits kernel #364

New issue

Closed

opened 2026-10-26 21:16:50 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:16:50 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/325
Author: @leavelet
Created: 5/1/2026
Status: ❌ Closed

Base: main ← Head: ws-upstream

📝 Commits (6)

d39b2f3 feat: warp-specialized SM90 FP8 MQA logits kernel
0c41e5d fix: cleanup SM90 FP8 MQA logits kernel for upstream
c7218cd refactor: derive kNumEpiStages from kNumKVStages, restore single-path dispatcher
e4c8bf0 rename: is_ws -> has_epi_wg for clarity
441c493 refactor: replace DO_PARTIAL_REDUCE macro with lambda
8425593 style: minor cleanups — use named constants, remove extra blank line

📊 Changes

2 files changed (+278 additions, -127 deletions)

View changed files

📝 csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp (+10 -2)
📝 deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh (+268 -125)

📄 Description

Add a dedicated epilogue warpgroup to the SM90 FP8 MQA logits kernel, overlapping post-processing (reduce + STG) with the next KV-block's WGMMA compute.

Architecture: 4 TC WGs (512) + 1 Epi WG (128) + 1 TMA WG (128) = 768 threads. TC does dual-bank WGMMA (N=64 split, warpgroup_wait<1> overlap) + partial reduce → STS to epi_buf in SMEM. Epi does LDS → horizontal sum → STG to global. TMA is unchanged.

Key design decisions:

TC@104 regs (0 spill): scale_kv deferred to between epi-empty wait and warpgroup_wait<1>
epi_buf stride = BLOCK_Q*4+1 for bank conflict reduction
kNumEpiStages derived from kNumKVStages (no extra template parameter)

Performance (342 configs, H200 SXM):

Speedup vs production: mean 1.065x, min 1.027x, max 1.111x
Regressions: 0 / 342
Correctness: 295/295 full test configs pass

_{🔄 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/325 **Author:** [@leavelet](https://github.com/leavelet) **Created:** 5/1/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `ws-upstream` --- ### 📝 Commits (6) - [`d39b2f3`](https://github.com/deepseek-ai/DeepGEMM/commit/d39b2f3fe23c2b81583bfadd1f7ace6b1a671080) feat: warp-specialized SM90 FP8 MQA logits kernel - [`0c41e5d`](https://github.com/deepseek-ai/DeepGEMM/commit/0c41e5d1d9a8fb9221a9ea668379c3161b92cce8) fix: cleanup SM90 FP8 MQA logits kernel for upstream - [`c7218cd`](https://github.com/deepseek-ai/DeepGEMM/commit/c7218cd04e0fff5d8425a805dffdf8be0a2bc0b4) refactor: derive kNumEpiStages from kNumKVStages, restore single-path dispatcher - [`e4c8bf0`](https://github.com/deepseek-ai/DeepGEMM/commit/e4c8bf04636cf0656acf8c7c8a6268f3516ba280) rename: is_ws -> has_epi_wg for clarity - [`441c493`](https://github.com/deepseek-ai/DeepGEMM/commit/441c493c5d04ecff9be2cef507f83adb5505de4c) refactor: replace DO_PARTIAL_REDUCE macro with lambda - [`8425593`](https://github.com/deepseek-ai/DeepGEMM/commit/8425593ba618b7b98b5b569c0579b670eadde78a) style: minor cleanups — use named constants, remove extra blank line ### 📊 Changes **2 files changed** (+278 additions, -127 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/smxx_fp8_fp4_mqa_logits.hpp` (+10 -2) 📝 `deep_gemm/include/deep_gemm/impls/sm90_fp8_mqa_logits.cuh` (+268 -125) </details> ### 📄 Description Add a dedicated epilogue warpgroup to the SM90 FP8 MQA logits kernel, overlapping post-processing (reduce + STG) with the next KV-block's WGMMA compute. Architecture: 4 TC WGs (512) + 1 Epi WG (128) + 1 TMA WG (128) = 768 threads. TC does dual-bank WGMMA (N=64 split, warpgroup_wait<1> overlap) + partial reduce → STS to epi_buf in SMEM. Epi does LDS → horizontal sum → STG to global. TMA is unchanged. Key design decisions: - TC@104 regs (0 spill): scale_kv deferred to between epi-empty wait and warpgroup_wait<1> - epi_buf stride = BLOCK_Q*4+1 for bank conflict reduction - kNumEpiStages derived from kNumKVStages (no extra template parameter) Performance (342 configs, H200 SXM): - Speedup vs production: mean 1.065x, min 1.027x, max 1.111x - Regressions: 0 / 342 - Correctness: 295/295 full test configs pass --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

2026-10-26 21:16:50 +00:00