[RFC] Define accelerator-agnostic backend boundaries for core MoE operations

MoE層で使用されている既存のmegablocks.opsインターフェースから始め、そのrouting、permutation、grouped expert operationsを追跡します。実装前に、提案されているpure-PyTorchの参照動作とconformance-testの範囲を確認します。operation contractとreference testが定義され、dispatchおよびaccelerator固有の実装が独立したフォローアップ作業として残されていれば完了です。

Summary

MegaBlocks already exposes megablocks.ops as the interface used by its MoE layers, but it does not currently have a consistent internal boundary between operation semantics and accelerator-specific implementations. Routing metadata operations depend on the megablocks_ops CUDA extension, token permutation uses Triton kernels, and grouped expert computation depends on the CUDA-specific grouped_gemm package.

This RFC proposes explicit internal contracts for core operations and grouped GEMM. A pure-PyTorch reference implementation would define device-independent correctness behavior, while optimized accelerator implementations could be developed and reviewed independently.

Motivation

• #97 requests AMD and ROCm support. A maintainer noted that the remaining CUDA kernels could be replaced with Torch operations or Triton kernels and welcomed contributions.
• #182 explores pure-PyTorch routing and permutation implementations for NPU tensors while preserving the CUDA paths.
• #187 explores ROCm support through HIPification and build-system changes.

These efforts require common decisions about operation semantics, dispatch, optional dependencies, and conformance testing. Defining those boundaries once would reduce duplication and make accelerator support easier to review and maintain.

Proposed direction

The existing megablocks.ops functions would remain the interface used by MoE layers. Their input, output, dtype, autograd, and edge-case behavior would be defined explicitly. An internal dispatch mechanism would select implementations from tensor device and runtime capabilities.

A pure-PyTorch backend would serve as the conformance reference. It would not replace optimized implementations by default. Existing CUDA kernels would remain the optimized CUDA path. Grouped GEMM would use a separate provider contract because its implementation and dependency requirements differ from routing and permutation operations. The existing grouped and sparse choices would continue to describe algorithms rather than accelerator platforms.

Contribution approach

I will drive the implementation. If this direction is acceptable, I would begin with a change defining operation contracts and pure-PyTorch reference tests. Dispatch, dependency separation, and accelerator-specific implementations would follow as independent pull requests.

Feedback requested

Before implementation, I would appreciate the maintainers' feedback on whether these backend boundaries fit the intended direction of MegaBlocks and whether the proposed incremental contribution approach is suitable

AI tool use

• Grammarly was used to improve grammar and phrasing.
• Cursor was used to inspect this repository, trace accelerator-specific dependencies, and review the claims of this issue.

I reviewed the final content and take responsibility for the proposal and any subsequent contributions.