Repository navigation
megablocks: add ROCm backend (AITER Triton grouped GEMM) - #916
Conversation
|
/kernel-bot build megablocks |
|
Build request processed. Command: Dispatched (1):
|
| from ._ops import ops as _compiled_ops | ||
| from . import ops |
There was a problem hiding this comment.
can you explain why we need both _compiled_ops and ops exposed?
There was a problem hiding this comment.
Inside the megablocks package there are two different "ops" things that happen to share the same name:
- The C++ kernels (compiled .so) : sort, histogram, etc. as raw Torch ops. Fast, no autograd, takes all the args (e.g. sort(x, end_bit, x_out, iota_out) : 4 args).
- A Python folder also called ops/ : wrappers around those C++ kernels that add autograd support. Same names, different signatures (e.g. sort(x, end_bit=None) : 2 args, returns a tuple).
There was a problem hiding this comment.
Before my change, the single name ops was bound to the compiled namespace, then silently overwritten to the submodule when from . import layers ran ; so half the file got one and half got the other depending on call order. That's the bug that broke megablocks.sort. Giving them distinct names makes the intent explicit and removes the order-dependence.
There was a problem hiding this comment.
thanks for the explanation and the fix!
|
/kernel-bot merge-and-upload megablocks |
|
Build request failed. Command: Failure: Failed to merge PR before build/upload. Check mergeability and required checks. |
|
/kernel-bot merge-and-upload megablocks |
|
Build request processed. Command: Merge result: PR is already merged. Continuing with build/upload. Dispatched (1):
|
* feat: allow kernel-bot to run security analysis. * include security in the build. * megablocks for rocm (#916) * feat: update merge-and-upload permissions (#930) * [Liger Kernels] Adding TiledMLP kernels (#918) * chore: adding liger tiled mlp kernels * liger-kernels: test with Triton ROCm binary wheels * chore: align layer to pure --------- Co-authored-by: Daniël de Kok <me@danieldk.eu> * finegrained-fp8: MXFP4 recipes (#923) * add fp4 matmul kernels for deepseek v4 flash Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * add missing source code.. Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * update to w4a8 Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * update Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * perf optimization add add benchmark_perf.py script Signed-off-by: Wang, Yi <yi.a.wang@intel.com> * structure * standarize all kernels and testing around the same api * better testing framework * fix batched with sentinels * tests with sentinels as well * test static and support ue8m0 * specify v2 as the version * cleanup * naming convention * update readme * remove bench * format * update readme * trigger * skip tests on < sm89 * drop sm89 testing, too error prone * Use triton-rocm-fixes branch to pass build * Update build.toml * Fix Torch ops * Remove custom path insertion This breaks tests in kernel-builder devshells --------- Signed-off-by: Wang, Yi <yi.a.wang@intel.com> Co-authored-by: Wang, Yi <yi.a.wang@intel.com> Co-authored-by: Daniël de Kok <me@danieldk.eu> * address review feedback * Revert "megablocks for rocm (#916)" This reverts commit e06068f. * Revert "feat: update merge-and-upload permissions (#930)" This reverts commit a913728. * Revert "[Liger Kernels] Adding TiledMLP kernels (#918)" This reverts commit 3d9bb2f. * Revert "finegrained-fp8: MXFP4 recipes (#923)" This reverts commit 9f9d5e1. * security-and-build --------- Signed-off-by: Wang, Yi <yi.a.wang@intel.com> Co-authored-by: BADAOUI Abdennacer <106801897+Abdennacer-Badaoui@users.noreply.github.com> Co-authored-by: drbh <david.richard.holtz@gmail.com> Co-authored-by: Aritra Roy Gosthipaty <aritra.born2fly@gmail.com> Co-authored-by: Daniël de Kok <me@danieldk.eu> Co-authored-by: Ilyas Moutawwakil <57442720+IlyasMoutawwakil@users.noreply.github.com> Co-authored-by: Wang, Yi <yi.a.wang@intel.com>
Adds AMD ROCm support to the megablocks kernel alongside the existing cuda/xpu/cpu backends.
[kernel.megablocks_rocm]block (gfx942 + gfx950) reusing the same CUDA op sources via hipify; two small CUDA-safe guards (_LDGfor__ldg,MEGABLOCKS_CUBforcub::DeviceHistogram) plug the gaps in torch's hipify symbol table.gmm/ptgmm, vendored undertorch-ext/megablocks/_grouped_gemm_triton/with an adapter mapping megablocks'gmm(a, b, c, batch_sizes, trans_a, trans_b)semantics.grouped_gemm/backend.pydispatches bytorch.version.hip.torch_binding.cppregisters the 5 ops underROCM_KERNEL(HIP →torch::kCUDA); CUTLASSgmmstays CUDA-only.megablocks/__init__.pywherefrom . import layersrebindsopsto the python submodule, breaking the top-levelmegablocks.sortwrapper on both backends.Validated on gfx942 (MI300X): ops correctness, grouped-GEMM autograd, full
MegaBlocksMoeMLPforward; both via local cpp-extension build (torch 2.10) and via the realnixkernel-builder artifact (torch212-cxx11-rocm72) loaded in the matching dev shell.