[PR #345] [CLOSED] Fix Mega MoE kernel shared memory calculation error #378

New issue

Closed

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/345
Author: @Oxygen56
Created: 5/31/2026
Status: ❌ Closed

Base: main ← Head: fix/issue-339

📝 Commits (1)

36abc9c Fix Mega MoE kernel shared memory calculation error

📊 Changes

1 file changed (+1 additions, -1 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+1 -1)

📄 Description

Description

Fixes a bug in the Mega MoE kernel where SMEM_BEFORE_BARRIER_SIZE was missing SMEM_SFA_SIZE_PER_STAGE and SMEM_SFB_SIZE_PER_STAGE in its per-stage shared memory calculation.

Root Cause

SMEM_BEFORE_BARRIER_SIZE only accounted for SMEM_A_SIZE_PER_STAGE and SMEM_B_SIZE_PER_STAGE per stage, but the scale factor (SF) buffers (SMEM_SFA_SIZE_PER_STAGE and SMEM_SFB_SIZE_PER_STAGE) were omitted. This caused the shared memory allocation before the barrier to be too small, triggering assertion failures for larger hidden sizes (e.g., hidden=8192, num_tokens=1024).

Fix

Add the missing SF size constants to the per-stage calculation:

 constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE =
-     SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE);
+     SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE + SMEM_SFB_SIZE_PER_STAGE);

Fixes #339

Co-Authored-By: Claude Opus 4.8 noreply@anthropic.com

_{🔄 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/345 **Author:** [@Oxygen56](https://github.com/Oxygen56) **Created:** 5/31/2026 **Status:** ❌ Closed **Base:** `main` ← **Head:** `fix/issue-339` --- ### 📝 Commits (1) - [`36abc9c`](https://github.com/deepseek-ai/DeepGEMM/commit/36abc9c3dd95c9b7412613c72c271b6bb6ad577b) Fix Mega MoE kernel shared memory calculation error ### 📊 Changes **1 file changed** (+1 additions, -1 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+1 -1) </details> ### 📄 Description ## Description Fixes a bug in the Mega MoE kernel where `SMEM_BEFORE_BARRIER_SIZE` was missing `SMEM_SFA_SIZE_PER_STAGE` and `SMEM_SFB_SIZE_PER_STAGE` in its per-stage shared memory calculation. ## Root Cause `SMEM_BEFORE_BARRIER_SIZE` only accounted for `SMEM_A_SIZE_PER_STAGE` and `SMEM_B_SIZE_PER_STAGE` per stage, but the scale factor (SF) buffers (`SMEM_SFA_SIZE_PER_STAGE` and `SMEM_SFB_SIZE_PER_STAGE`) were omitted. This caused the shared memory allocation before the barrier to be too small, triggering assertion failures for larger hidden sizes (e.g., hidden=8192, num_tokens=1024). ## Fix Add the missing SF size constants to the per-stage calculation: ```diff constexpr uint32_t SMEM_BEFORE_BARRIER_SIZE = - SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE); + SMEM_EXPERT_COUNT_SIZE + SMEM_SEND_BUFFER_SIZE + SMEM_CD_SIZE + kNumStages * (SMEM_A_SIZE_PER_STAGE + SMEM_B_SIZE_PER_STAGE + SMEM_SFA_SIZE_PER_STAGE + SMEM_SFB_SIZE_PER_STAGE); ``` Fixes #339 Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi

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