deepseek-ai/DeepGEMM

Fork 0

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

[PR #389] fix: fence combine buffer reuse in mega moe #408

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/389
Author: @Bruce-Lee-LY
Created: 7/22/2026
Status: 🔄 Open

Base: main ← Head: feat/fix_mega

📝 Commits (1)

bbbf8e4 fix: fence combine buffer reuse in mega moe

📊 Changes

3 files changed (+8 additions, -0 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh (+2 -0)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+2 -0)
📝 deep_gemm/include/deep_gemm/ptx/utils.cuh (+4 -0)

📄 Description

Summary

This PR fixes a shared-memory reuse race in the final combine stage of the SM100 Mega MoE kernels:

fp8_fp4_mega_moe
bf16_mega_moe

Both kernels use two shared-memory load stages. A stage is filled by TMA through the async proxy, consumed by the warp through generic shared-memory loads, and then reused by a later TMA load. The old code did not order the generic reads against the next async-proxy write to the same stage.

The fix adds fence.proxy.async.shared::cta after every lane has consumed the current stage and before that stage can be selected for reuse.

Problem

The combine loop overlaps loading and reduction with a two-stage ping-pong buffer:

sequenceDiagram
    participant TMA as TMA async proxy
    participant S0 as Shared stage 0
    participant S1 as Shared stage 1
    participant W as Epilogue warp

    TMA->>S0: Load top-k contribution 0
    TMA->>S1: Prefetch contribution 1
    W->>S0: Wait on mbarrier and read stage 0
    Note over W,S0: Generic shared-memory reads
    TMA->>S0: Reuse stage 0 for contribution 2
    Note over TMA,S0: Async-proxy write may race with prior generic reads

The mbarrier wait makes the previous TMA load visible before the warp reads the stage. It does not, by itself, order those generic reads before a later TMA async-proxy write that reuses the same shared-memory addresses. Under unlucky timing, the new TMA load can overwrite data while lanes are still consuming the old contribution, corrupting the final reduction.

Fix

flowchart LR
    A[TMA load into stage] --> B[mbarrier wait]
    B --> C[All lanes read and accumulate]
    C --> D[fence.proxy.async.shared::cta]
    D --> E[Rotate load stage]
    E --> F[Later TMA load may reuse stage]

The new fence establishes the required generic-to-async proxy ordering at CTA scope:

// Order generic shared-memory reads before the next TMA load reuses this stage.
ptx::fence_proxy_async_shared_cta();
combine_phase ^= load_stage_idx;
load_stage_idx ^= 1;

The PTX instruction is exposed through a small helper in deep_gemm/ptx/utils.cuh and used by both affected kernels. There are no API, layout, scheduling, or numerical changes.

_{🔄 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/389 **Author:** [@Bruce-Lee-LY](https://github.com/Bruce-Lee-LY) **Created:** 7/22/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `feat/fix_mega` --- ### 📝 Commits (1) - [`bbbf8e4`](https://github.com/deepseek-ai/DeepGEMM/commit/bbbf8e4535a42f40d68e113e628add0317c4fa9b) fix: fence combine buffer reuse in mega moe ### 📊 Changes **3 files changed** (+8 additions, -0 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm100_bf16_mega_moe.cuh` (+2 -0) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+2 -0) 📝 `deep_gemm/include/deep_gemm/ptx/utils.cuh` (+4 -0) </details> ### 📄 Description ## Summary This PR fixes a shared-memory reuse race in the final combine stage of the SM100 Mega MoE kernels: - `fp8_fp4_mega_moe` - `bf16_mega_moe` Both kernels use two shared-memory load stages. A stage is filled by TMA through the async proxy, consumed by the warp through generic shared-memory loads, and then reused by a later TMA load. The old code did not order the generic reads against the next async-proxy write to the same stage. The fix adds `fence.proxy.async.shared::cta` after every lane has consumed the current stage and before that stage can be selected for reuse. ## Problem The combine loop overlaps loading and reduction with a two-stage ping-pong buffer: ```mermaid sequenceDiagram participant TMA as TMA async proxy participant S0 as Shared stage 0 participant S1 as Shared stage 1 participant W as Epilogue warp TMA->>S0: Load top-k contribution 0 TMA->>S1: Prefetch contribution 1 W->>S0: Wait on mbarrier and read stage 0 Note over W,S0: Generic shared-memory reads TMA->>S0: Reuse stage 0 for contribution 2 Note over TMA,S0: Async-proxy write may race with prior generic reads ``` The mbarrier wait makes the previous TMA load visible before the warp reads the stage. It does not, by itself, order those generic reads before a later TMA async-proxy write that reuses the same shared-memory addresses. Under unlucky timing, the new TMA load can overwrite data while lanes are still consuming the old contribution, corrupting the final reduction. ## Fix ```mermaid flowchart LR A[TMA load into stage] --> B[mbarrier wait] B --> C[All lanes read and accumulate] C --> D[fence.proxy.async.shared::cta] D --> E[Rotate load stage] E --> F[Later TMA load may reuse stage] ``` The new fence establishes the required generic-to-async proxy ordering at CTA scope: ```cpp // Order generic shared-memory reads before the next TMA load reuses this stage. ptx::fence_proxy_async_shared_cta(); combine_phase ^= load_stage_idx; load_stage_idx ^= 1; ``` The PTX instruction is exposed through a small helper in `deep_gemm/ptx/utils.cuh` and used by both affected kernels. There are no API, layout, scheduling, or numerical changes. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

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

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

No description provided.

Rows
Columns