feat(moe): support HIP and Triton backends for moe_permute, default to Triton by zhenhuang12 · Pull Request #455 · AMD-AGI/Primus-Turbo

zhenhuang12 · 2026-08-12T01:41:58Z

Description

moe_permute / moe_unpermute previously only had the HIP extension path. This PR makes the
backend selectable between HIP (BackendType.TURBO) and a new Triton one
(BackendType.TRITON), defaults to Triton for the plain (non-padded, non-FP8) case, and makes
the permute roundtrip traceable under torch.compile.

Fixes AIMA-219, AIMA-220

Type of change

Documentation change (change only to the documentation, either a fix or a new content)
Bug fix (non-breaking change which fixes an issue)
New feature (non-breaking change which adds functionality)
Breaking change (fix or feature that would cause existing functionality to not work as expected)
Infra/Build change
Code refactoring

Changes

Please list the changes introduced in this PR:

Add Triton permute / unpermute kernels and a backend argument selecting between HIP and
Triton. Padding and FP8 still go to HIP; everything else defaults to Triton. Permute and
unpermute must use the same backend.
Wrap the permute stages in torch.library.custom_op with register_fake meta impls, so the
roundtrip captures as a single graph.
Breaking: moe_permute takes keyword-only routing_map / topk_indices instead of
expert_map, and no longer returns num_dispatched_tokens (7 → 6 values). The row bound now
comes from DeepEP through the dispatcher.
Dispatcher: reuse DeepEP's already-synced host tokens_per_expert instead of forcing a second
device sync, and tighten a few argument checks.
Fix GroupedGemmFunc.backward folding uninitialized worst-case tail rows into grad_b for the
single-group case, plus several permute edge cases (zero tokens, padding, num_topk).
Fix deep_ep.cpp reserving an idle hardware queue when force_current_stream is set.
Move the permute / indices-converter impls under pytorch/kernels/moe/ and drop dead code.

Checklist:

The functionality is complete
I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes

GroupedGemmFunc's backward took a dense gemm_impl path when len(group_lens) == 1. That reduces over all M rows of `a`, but under the sync-free MoE path `a` is worst-case allocated and its tail rows are never written -- the dense GEMM folds that uninitialized memory into grad_b. Route the single-group case through grouped_gemm_variable_k_impl like the multi-group branch already does, so the reduction stops at group_offs. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

moe_permute/moe_unpermute gain a `backend` argument selecting between the HIP C++ extension and new Triton kernels. Both backends consume the same HIP-built row_id_map, so the Triton path is a drop-in swap of the scatter/ gather kernels only -- preprocessing is unchanged. `backend` defaults to "triton"; pass "hip" to fall back to the extension. DeepEPTokenDispatcher no longer needs to request the Triton path explicitly. Adds permute_with_hip_row_map / unpermute_with_hip_row_map to the Triton moe permutation module, a standalone benchmark, and tests that run both backends and assert the Triton results match HIP for topk probs, both expert_map dtypes, and padded and unpadded layouts. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

moe_permute no longer computes and returns a dispatched-token count; the row bound is whatever DeepEP already knows, so the return tuple drops from 7 to 6 values and the strided row_id_map[:, -1] read is gone. - Route num_dispatched_tokens through the dispatcher instead of the kernel. - Cache the device-side row bound in DeepEPTokenDispatcher: DeepEP's recv height is stable, so the steady state stops allocating a torch.full per layer per step, and moe_unpermute's backward reuses the same tensor. - Drop the dead permute_with_hip_row_map / unpermute_with_hip_row_map paths and make backend auto-selection symmetric between permute and unpermute. - Move the triton permute and indices-converter impls under kernels/moe and fix the stale paths in LICENSE. Verified in-container on 8x MI355: 97 passed / 13 skipped for test_moe_permute.py + test_token_dispatcher.py, 36 passed for the router and mask tests, 3 passed for fused_mega_moe. A profiler sweep over a steady-state dispatch/combine/backward reports D2H=0 H2D=0 for TRITON pad=0 (host and cuda counts) and TURBO pad=128, with a positive control confirming the counter sees real copies. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Wrap the three moe_permute stages (preprocess / permute / unpermute) in torch.library.custom_op with explicit schema strings -- infer_schema cannot express the optional tensor returns -- plus register_fake meta impls so the whole roundtrip captures as a single graph. The preprocess op no longer passes ``probs`` straight through: a custom op may not return one of its own inputs. Backends now hand back ``multihot_probs``, non-None only when TRITON actually converted the top-k probs, and the python wrapper resolves it, keeping the public 8-tuple contract unchanged. Guard ``ctx.set_materialize_grads`` behind ``is_compiling()`` in both autograd Functions; dynamo cannot trace that ctx method and would break the graph. Eager and compiled forward/backward stay bit-identical. Tests: 84 passed / 13 skipped in tests/pytorch/ops/test_moe_permute.py (4 new torch.compile cases), 15 passed in test_token_dispatcher.py. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

An idle high-priority pool stream still consumes one of GPU_MAX_HW_QUEUES. When force_current_stream is set, comm_stream is never used, so alias it to the caller's current stream instead of reserving a hardware queue for nothing. The env is read directly here because force_current_stream is declared after comm_stream and is not initialised yet at this point in the member init list. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

The Triton-vs-HIP permute comparison served the backend bring-up and is not part of the shipped surface; nothing references it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

…dding - treat num_permuted_tokens == 0 as a valid caller-provided bound - return 0 permuted rows in the meta kernel when there are no tokens - keep indices_position_map alive via save_for_backward - add the pad_multiple slack to the dispatcher permute bound - validate num_topk against topk_indices width and token_indices shape - reject deepep_num_worst_tokens without cuda tokens-per-expert - forward use_fp8 into the moe_unpermute default backend choice Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Copilot

Pull request overview

This PR refactors and extends the PyTorch MoE permute/unpermute path to support selecting between HIP (“TURBO”) and a new Triton backend, defaulting to Triton for the unpadded/non-FP8 case, and adds torch.compile traceability via torch.library custom ops.

Changes:

Adds Triton-backed permute/unpermute implementation and backend selection logic (with TURBO retained for padding/FP8/capacity-related behavior).
Updates DeepEP token dispatcher to reuse DeepEP-provided token counts and tightens input/shape validation.
Updates and expands unit tests to cover both backends, padding/capacity edge cases, and torch.compile tracing behavior; removes/relocates older MoE permutation utilities under pytorch/kernels/moe/.

Reviewed changes

Copilot reviewed 14 out of 14 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/pytorch/ops/test_moe_permute.py	Expands coverage across TURBO/TRITON backends, adds compile and edge-case tests, updates API expectations.
tests/pytorch/modules/test_token_dispatcher.py	Improves multi-rank failure handling and strengthens dispatcher correctness checks (counts, padding alignment).
primus_turbo/triton/moe/permutation.py	Formats/adjusts Triton permutation kernels used by the TRITON backend.
primus_turbo/pytorch/ops/moe/permutation.py	Deletes legacy MoE permutation ops (moved/replaced by kernel backend layer).
primus_turbo/pytorch/ops/moe/moe_permute.py	Introduces backend selection + new API surface, routes to backend implementations, adds compile-friendly handling.
primus_turbo/pytorch/ops/moe/indices_converter.py	Deletes legacy indices converter (replaced by fused kernel impl).
primus_turbo/pytorch/ops/moe/init.py	Stops exporting removed legacy modules.
primus_turbo/pytorch/ops/grouped_gemm.py	Fixes single-group backward grad_b computation to avoid folding uninitialized tail rows.
primus_turbo/pytorch/modules/moe/token_dispatcher.py	Reuses DeepEP host counts, adjusts capacity sizing logic, tightens token_indices shape checks.
primus_turbo/pytorch/kernels/moe/moe_permute_impl.py	New backend abstraction + torch.library custom ops for preprocess/permute/unpermute (TURBO + TRITON).
primus_turbo/pytorch/kernels/moe/fused_moe_indices_converter_impl.py	New fused topk<->multihot converter kernel implementation for TRITON paths.
primus_turbo/pytorch/kernels/moe/init.py	Exposes new MoE kernel backend APIs.
LICENSE	Updates third-party attribution paths after file moves.
csrc/pytorch/deep_ep/deep_ep.cpp	Avoids reserving an idle HW queue by selecting current stream when force_current_stream is enabled.

Suppressed comments (1)

primus_turbo/pytorch/ops/moe/moe_permute.py:247

_MoEUnpermute.forward calls ctx.save_for_backward(row_id_map, num_dispatched_tokens_tensor), but num_dispatched_tokens_tensor is optional and is typically None at call sites. save_for_backward cannot accept None, so moe_unpermute will raise at runtime.

Persist the optional bound as a ctx attribute and only save_for_backward the row_id_map tensor.

        ctx.save_for_backward(row_id_map, num_dispatched_tokens_tensor)
        return outputs

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

+        ctx.save_for_backward(row_id_map, num_dispatched_tokens, indices_position_map)
        return (


Formatting only, no behavior change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Copilot

Pull request overview

Copilot reviewed 14 out of 14 changed files in this pull request and generated 1 comment.

Suppressed comments (1)

primus_turbo/pytorch/ops/moe/moe_permute.py:330

moe_permute currently allows both routing_map and topk_indices to be passed. The backends resolve this differently (TURBO prefers topk_indices; TRITON prefers routing_map), which can change results depending on backend selection. Reject ambiguous calls by requiring exactly one of the two inputs.

    if routing_map is None and topk_indices is None:
        raise ValueError("moe_permute: one of routing_map / topk_indices must be provided")

+        # Must be an upper bound: the kernel stores at dst row without clamping.
+        if num_permuted_tokens >= 0:
+            num_permuted = int(num_permuted_tokens)
+        else:
+            num_permuted = int(tokens_per_expert.sum().item())


zhenhuang12 and others added 7 commits August 8, 2026 11:15

chore(benchmark): drop bench_moe_permute

dd44086

The Triton-vs-HIP permute comparison served the backend bring-up and is not part of the shipped surface; nothing references it. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

zhenhuang12 requested a review from xiaobochen-amd as a code owner August 12, 2026 01:41

Copilot AI lite review requested due to automatic review settings August 12, 2026 01:41

zhenhuang12 requested review from RuibinCheung and wenxie-amd as code owners August 12, 2026 01:41

Copilot started reviewing on behalf of zhenhuang12 August 12, 2026 01:43 View session

Copilot AI reviewed Aug 12, 2026

View reviewed changes

Comment thread primus_turbo/pytorch/ops/moe/moe_permute.py

Comment on lines +125 to 126

ctx.save_for_backward(row_id_map, num_dispatched_tokens, indices_position_map)

return (

style(moe): apply ruff-format and clang-format

5560491

Formatting only, no behavior change. Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>

Copilot AI review requested due to automatic review settings August 12, 2026 02:01

Copilot started reviewing on behalf of zhenhuang12 August 12, 2026 02:03 View session

Copilot AI reviewed Aug 12, 2026

View reviewed changes

Navigation Menu

Uh oh!

feat(moe): support HIP and Triton backends for moe_permute, default to Triton - #455

feat(moe): support HIP and Triton backends for moe_permute, default to Triton#455
zhenhuang12 wants to merge 8 commits into
mainfrom
feat/zhuang12/moe-permute-triton-default

zhenhuang12 commented Aug 12, 2026

Uh oh!

Copilot AI left a comment

Uh oh!

Copilot AI left a comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

		ctx.save_for_backward(row_id_map, num_dispatched_tokens, indices_position_map)
		return (

Uh oh!

Conversation

zhenhuang12 commented Aug 12, 2026

Description

Type of change

Changes

Checklist:

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants