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

Comienza por la interfaz megablocks.ops existente que utilizan las capas MoE y sigue su enrutamiento, permutación y operaciones agrupadas de expertos. Revisa el comportamiento de referencia propuesto en pure-PyTorch y el alcance de las pruebas de conformidad antes de la implementación. Se considera terminado cuando los contratos de las operaciones y las pruebas de referencia estén definidos, dejando las implementaciones de dispatch y específicas de aceleradores como trabajo de seguimiento independiente.

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.