[PR #404] feat: add opt-in MegaMoE EVICT_FIRST weight loads #422

New issue

Open

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

bsakaguchi commented

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

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/404
Author: @freeliuzc
Created: 8/12/2026
Status: 🔄 Open

Base: main ← Head: pr/mega-moe-evict-first

📝 Commits (4)

9a2a4dd feat: add opt-in MegaMoE EVICT_FIRST weight loads
d149ce1 fix: gate MegaMoE EVICT_FIRST at 512 tokens
f03a59e fix: pass MegaMoE cache hint through JIT template
12aa4d5 fix: constrain supported TMA cache hints

📊 Changes

4 files changed (+29 additions, -11 deletions)

View changed files

📝 README.md (+2 -0)
📝 csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp (+9 -2)
📝 deep_gemm/include/deep_gemm/common/tma_copy.cuh (+14 -8)
📝 deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh (+4 -1)

📄 Description

Add an opt-in EVICT_FIRST hint for small-batch MegaMoE

Summary

This change adds an opt-in SM100 FP8xFP4 MegaMoE cache-policy tuning switch.
When DG_MEGA_MOE_WEIGHT_EVICT_FIRST=1, routed expert-weight demand TMA loads
use the GPU L2 EVICT_FIRST replacement hint. The normal path remains the
default (0), and shared-expert weights and scale-factor loads are unchanged.

The hint is deliberately independent of speculative previous-active-expert
prefetch. Prefetch remains excluded from this PR because its benefit is
routing-sensitive.

The switch is enabled only when num_tokens <= 512; larger batches always use
the normal cache policy.

Performance

Measured on 8 SM100 GPU ranks with the DeepSeek-V4-Flash-FP8 configuration,
FP8xFP4 MegaMoE, random routing with sequence length 8, one shared expert,
hidden size 4096, intermediate size 2048, 256 total experts, and top-k 6.
These measurements use the existing upstream/origin MegaMoE heuristic; no
v3, v6, tiling, cluster, or prefetch experiment is included in this PR.
Each value is the mean across 8 ranks from a correctness-plus-performance run.

Tokens per rank	`EVICT_NORMAL`	`EVICT_FIRST`	Improvement
8	102.304 us	96.851 us	+5.33%
16	110.812 us	106.937 us	+3.50%
32	113.374 us	109.755 us	+3.19%
512	140.852 us	138.810 us	+1.45%
1024	211.897 us	211.544 us	+0.17%

The 1024-token result is expected to be effectively unchanged because the
runtime gate selects EVICT_NORMAL for that case.

Repro

Build the branch in a CUDA/SM100 environment, use a fresh JIT cache for
each variant, and run the two variants in alternating order on an otherwise
idle 8-GPU node. num_shared_experts=1 and
num_max_tokens_per_rank=2048 intentionally match the measurements above.

for tokens in 8 16 32 512 1024; do
  for hint in 0 1; do
    export DG_JIT_CACHE_DIR="/tmp/deepgemm-evict-first-${tokens}-${hint}"
    DG_MEGA_MOE_WEIGHT_EVICT_FIRST=$hint \
    CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \
    python tests/test_mega_moe.py \
      --num-processes 8 \
      --num-max-tokens-per-rank 2048 \
      --num-tokens "$tokens" \
      --hidden 4096 \
      --intermediate-hidden 2048 \
      --num-shared-experts 1 \
      --num-experts 256 \
      --num-topk 6 \
      --mma-type fp8xfp4 \
      --num-correctness-tests 1
  done
done

The commands above use only arguments available in the existing upstream test
script and reproduce correctness plus a basic A/B measurement. The reported
performance table was collected with a separate experiment harness that cycled
a deterministic sequence of eight random routings during benchmarking. Collect
fused-kernel latency and report rank means. The default path (hint=0) must
retain EVICT_NORMAL behavior.

Scope

Add a compile-time cache-hint parameter to the existing TMA copy helper.
Thread one boolean through MegaMoE runtime code generation.
Apply the hint only to routed expert-weight demand copies.
Document the opt-in environment variable.

No routing, scheduler, tiling, cluster, synchronization, numerical, or public
Python API behavior changes are included.

Testing

git diff --check
python -m py_compile tests/test_mega_moe.py
Focused SM100 A/B benchmark and correctness runs described above
Default path remains EVICT_NORMAL when the variable is unset or 0

Known Limitations

This is an experimental, opt-in cache-policy hint and is not enabled by
default.
The measured benefit is workload-dependent and decreases with larger token
counts.
Speculative previous-active-expert prefetch is deliberately excluded from
this PR and remains in the experiment archive for a separate proposal.

_{🔄 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/404 **Author:** [@freeliuzc](https://github.com/freeliuzc) **Created:** 8/12/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `pr/mega-moe-evict-first` --- ### 📝 Commits (4) - [`9a2a4dd`](https://github.com/deepseek-ai/DeepGEMM/commit/9a2a4dd7f005a827cc145487ceed8b586ecbb2f2) feat: add opt-in MegaMoE EVICT_FIRST weight loads - [`d149ce1`](https://github.com/deepseek-ai/DeepGEMM/commit/d149ce151fddf9b6c60038b61883e5e2c05c0402) fix: gate MegaMoE EVICT_FIRST at 512 tokens - [`f03a59e`](https://github.com/deepseek-ai/DeepGEMM/commit/f03a59e5cbd422ebff311053c67639218e8b282a) fix: pass MegaMoE cache hint through JIT template - [`12aa4d5`](https://github.com/deepseek-ai/DeepGEMM/commit/12aa4d5f48bfae04839fdbfdee942028b3419675) fix: constrain supported TMA cache hints ### 📊 Changes **4 files changed** (+29 additions, -11 deletions) <details> <summary>View changed files</summary> 📝 `README.md` (+2 -0) 📝 `csrc/jit_kernels/impls/sm100_fp8_fp4_mega_moe.hpp` (+9 -2) 📝 `deep_gemm/include/deep_gemm/common/tma_copy.cuh` (+14 -8) 📝 `deep_gemm/include/deep_gemm/impls/sm100_fp8_fp4_mega_moe.cuh` (+4 -1) </details> ### 📄 Description # Add an opt-in EVICT_FIRST hint for small-batch MegaMoE ## Summary This change adds an opt-in SM100 FP8xFP4 MegaMoE cache-policy tuning switch. When `DG_MEGA_MOE_WEIGHT_EVICT_FIRST=1`, routed expert-weight demand TMA loads use the GPU L2 `EVICT_FIRST` replacement hint. The normal path remains the default (`0`), and shared-expert weights and scale-factor loads are unchanged. The hint is deliberately independent of speculative previous-active-expert prefetch. Prefetch remains excluded from this PR because its benefit is routing-sensitive. The switch is enabled only when `num_tokens <= 512`; larger batches always use the normal cache policy. ## Performance Measured on 8 SM100 GPU ranks with the DeepSeek-V4-Flash-FP8 configuration, FP8xFP4 MegaMoE, random routing with sequence length 8, one shared expert, hidden size 4096, intermediate size 2048, 256 total experts, and top-k 6. These measurements use the existing upstream/origin MegaMoE heuristic; no `v3`, `v6`, tiling, cluster, or prefetch experiment is included in this PR. Each value is the mean across 8 ranks from a correctness-plus-performance run. | Tokens per rank | `EVICT_NORMAL` | `EVICT_FIRST` | Improvement | | ---: | ---: | ---: | ---: | | 8 | 102.304 us | 96.851 us | +5.33% | | 16 | 110.812 us | 106.937 us | +3.50% | | 32 | 113.374 us | 109.755 us | +3.19% | | 512 | 140.852 us | 138.810 us | +1.45% | | 1024 | 211.897 us | 211.544 us | +0.17% | The 1024-token result is expected to be effectively unchanged because the runtime gate selects `EVICT_NORMAL` for that case. ## Repro Build the branch in a CUDA/SM100 environment, use a fresh JIT cache for each variant, and run the two variants in alternating order on an otherwise idle 8-GPU node. `num_shared_experts=1` and `num_max_tokens_per_rank=2048` intentionally match the measurements above. ```bash for tokens in 8 16 32 512 1024; do for hint in 0 1; do export DG_JIT_CACHE_DIR="/tmp/deepgemm-evict-first-${tokens}-${hint}" DG_MEGA_MOE_WEIGHT_EVICT_FIRST=$hint \ CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7 \ python tests/test_mega_moe.py \ --num-processes 8 \ --num-max-tokens-per-rank 2048 \ --num-tokens "$tokens" \ --hidden 4096 \ --intermediate-hidden 2048 \ --num-shared-experts 1 \ --num-experts 256 \ --num-topk 6 \ --mma-type fp8xfp4 \ --num-correctness-tests 1 done done ``` The commands above use only arguments available in the existing upstream test script and reproduce correctness plus a basic A/B measurement. The reported performance table was collected with a separate experiment harness that cycled a deterministic sequence of eight random routings during benchmarking. Collect fused-kernel latency and report rank means. The default path (`hint=0`) must retain `EVICT_NORMAL` behavior. ## Scope - Add a compile-time cache-hint parameter to the existing TMA copy helper. - Thread one boolean through MegaMoE runtime code generation. - Apply the hint only to routed expert-weight demand copies. - Document the opt-in environment variable. No routing, scheduler, tiling, cluster, synchronization, numerical, or public Python API behavior changes are included. ## Testing - `git diff --check` - `python -m py_compile tests/test_mega_moe.py` - Focused SM100 A/B benchmark and correctness runs described above - Default path remains `EVICT_NORMAL` when the variable is unset or `0` ## Known Limitations - This is an experimental, opt-in cache-policy hint and is not enabled by default. - The measured benefit is workload-dependent and decreases with larger token counts. - Speculative previous-active-expert prefetch is deliberately excluded from this PR and remains in the experiment archive for a separate proposal. --- <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:36 +00:00

bsakaguchi referenced this issue

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

[PR #422] [Feat] add support sm90 fp8 mega moe #434

bsakaguchi referenced this issue

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

[PR #440] Add a fused SM90 FP8 MegaMoE: one kernel, a lag-ordered schedule and a ring activation pool #449