Skip to content

Chunk-GEMM compute-compensation tests (zero kernel change) - #13

Open
ForFishes wants to merge 3 commits into
PFCCLab:teramoefrom
ForFishes:teramoe-compensation
Open

ForFishes wants to merge 3 commits into
PFCCLab:teramoefrom
ForFishes:teramoe-compensation

Conversation

@ForFishes

Copy link
Copy Markdown

Summary

DeepGEMM-side validation that TeraMoE intra-node compute compensation needs zero kernel change.

A helper GPU can absorb peer experts' chunks purely by:

  1. extending the weight table b = [local experts ++ migrated peer experts], and
  2. pointing the migrated chunks' task_queue expert_idx at the appended slots.

The existing bf16_chunk_gemm_nn / chunk_weighted_swiglu / chunk_zip kernels already index the weight table by task.expert_idx and cover all groups, so no kernel modification is required.

Added

  • tests_overlap/test_compensation_chunk.py: a helper absorbs peer experts (extended-table slots); chunk-GEMM output vs the m_grouped reference is bitwise diff=0 (both all-valid and migrated-only), plus a build_helper_schedule constructor.
  • tests_overlap/test_fp8_compensation_chunk.py: FP8 variant.

Testing

  • Chunk-GEMM (o1/o3) vs m_grouped reference: bitwise diff=0 for all-valid and migrated-only schedules.

🤖 Generated with Claude Code

root and others added 3 commits September 21, 2026 14:28
DeepGEMM-side validation that intra-node compute compensation needs zero
kernel change: a helper GPU absorbs peer experts' chunks purely by
extending the weight table (b=[local experts ++ migrated peer experts])
and pointing the migrated chunks' task_queue expert_idx at the appended
slots — the existing bf16_chunk_gemm_nn / chunk_weighted_swiglu / chunk_zip
kernels index the weight table by task.expert_idx and cover all groups.

- test_compensation_chunk.py: helper absorbs peer experts; chunk-GEMM
  output vs m_grouped reference is bitwise diff=0 (all-valid and
  migrated-only), plus a build_helper_schedule constructor.
- test_fp8_compensation_chunk.py: FP8 variant.

Co-Authored-By: Claude Code <noreply@anthropic.com>
Co-Authored-By: Claude Code <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant