[PR #331] fix: use 64-bit mega moe sf offsets #369

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/331
Author: @netaddi
Created: 5/10/2026
Status: 🔄 Open

Base: main ← Head: fix/mega-moe-sf-offset-u64-pr

📝 Commits (1)

b0ba0cb fix: use 64-bit mega moe sf offsets

📊 Changes

3 files changed (+279 additions, -6 deletions)

View changed files

📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+9 -5)
➕ tests/test_mega_moe_capacity.py (+264 -0)
📝 tests/test_sanitizer.py (+6 -1)

📄 Description

Summary

Fix SM100 Mega MoE scale-factor pool address arithmetic for large
num_max_tokens_per_rank.

This PR:

uses 64-bit arithmetic for SM100 Mega MoE SF-pool offsets;
adds a manual SM100 capacity regression test;
excludes the heavy multi-process regression from sanitizer auto-discovery.

Bug Discovery

While validating Mega MoE in a serving-like SM100 setup, the output changed when
only the symmetric-memory buffer capacity changed.

The reproducer uses identical live tokens, inputs, routing, and weights, then
runs:

small capacity: num_max_tokens_per_rank=32
large capacity: num_max_tokens_per_rank=200000

The capacity should only affect allocation size. It should not affect the
mathematical output for the same live tokens.

Root Cause

In sm100_fp8_fp4_mega_moe_impl, some SF-pool offsets were computed with
32-bit intermediate values.

For large num_max_tokens_per_rank, kNumPaddedSFPoolTokens becomes large
enough that expressions like:

j * kNumPaddedSFPoolTokens + sf_pool_token_idx

and:

k_uint_idx * mn_stride + sf_pool_token_idx * sizeof(uint32_t) + byte_idx

can overflow 32-bit arithmetic before being used as addresses.

That makes the kernel write or read FP8/FP4 scale factors at the wrong SF-pool
locations, so the subsequent GEMMs consume incorrect scales and produce wrong
outputs.

Fix

Promote the SF-pool offset calculations to uint64_t in both affected paths:

dispatch-side local SF pool write;
L2 SF buffer address calculation.

This keeps the indexing valid for large Mega MoE capacities while preserving the
existing layout and kernel behavior for normal capacities.

Validation

Syntax / diff checks:

python3 -m py_compile tests/test_mega_moe_capacity.py tests/test_sanitizer.py
git diff --check

Manual SM100 regression test:

python tests/test_mega_moe_capacity.py \
--num-processes 8 \
--small-capacity 32 \
--large-capacity 200000

Expected result after the fix:

all ranks produce matching small/large output hashes;
max_abs_diff == 0.0;
mean_abs_diff == 0.0.

_{🔄 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/331 **Author:** [@netaddi](https://github.com/netaddi) **Created:** 5/10/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `fix/mega-moe-sf-offset-u64-pr` --- ### 📝 Commits (1) - [`b0ba0cb`](https://github.com/deepseek-ai/DeepGEMM/commit/b0ba0cbc7ec1ac2c7667033fef46d33a2b568c47) fix: use 64-bit mega moe sf offsets ### 📊 Changes **3 files changed** (+279 additions, -6 deletions) <details> <summary>View changed files</summary> 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+9 -5) ➕ `tests/test_mega_moe_capacity.py` (+264 -0) 📝 `tests/test_sanitizer.py` (+6 -1) </details> ### 📄 Description ## Summary Fix SM100 Mega MoE scale-factor pool address arithmetic for large `num_max_tokens_per_rank`. This PR: - uses 64-bit arithmetic for SM100 Mega MoE SF-pool offsets; - adds a manual SM100 capacity regression test; - excludes the heavy multi-process regression from sanitizer auto-discovery. ## Bug Discovery While validating Mega MoE in a serving-like SM100 setup, the output changed when only the symmetric-memory buffer capacity changed. The reproducer uses identical live tokens, inputs, routing, and weights, then runs: - small capacity: `num_max_tokens_per_rank=32` - large capacity: `num_max_tokens_per_rank=200000` The capacity should only affect allocation size. It should not affect the mathematical output for the same live tokens. ## Root Cause In `sm100_fp8_fp4_mega_moe_impl`, some SF-pool offsets were computed with 32-bit intermediate values. For large `num_max_tokens_per_rank`, `kNumPaddedSFPoolTokens` becomes large enough that expressions like: ```cpp j * kNumPaddedSFPoolTokens + sf_pool_token_idx ``` and: ```cpp k_uint_idx * mn_stride + sf_pool_token_idx * sizeof(uint32_t) + byte_idx ``` can overflow 32-bit arithmetic before being used as addresses. That makes the kernel write or read FP8/FP4 scale factors at the wrong SF-pool locations, so the subsequent GEMMs consume incorrect scales and produce wrong outputs. ## Fix Promote the SF-pool offset calculations to uint64_t in both affected paths: - dispatch-side local SF pool write; - L2 SF buffer address calculation. This keeps the indexing valid for large Mega MoE capacities while preserving the existing layout and kernel behavior for normal capacities. ## Validation Syntax / diff checks: ``` python3 -m py_compile tests/test_mega_moe_capacity.py tests/test_sanitizer.py git diff --check ``` Manual SM100 regression test: ``` python tests/test_mega_moe_capacity.py \ --num-processes 8 \ --small-capacity 32 \ --large-capacity 200000 ``` Expected result after the fix: - all ranks produce matching small/large output hashes; - max_abs_diff == 0.0; - mean_abs_diff == 0.0. --- <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:16:57 +00:00

bsakaguchi referenced this issue

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

[PR #369] [MERGED] support indexer_n_heads = 16 #394