Skip to content

Musa support native kernels - #376

Draft
Arlo-mt wants to merge 6 commits into
RL-Align:mainfrom
Arlo-mt:MUSA-support-native-kernels
Draft

Musa support native kernels#376
Arlo-mt wants to merge 6 commits into
RL-Align:mainfrom
Arlo-mt:MUSA-support-native-kernels

Conversation

@Arlo-mt

@Arlo-mt Arlo-mt commented Sep 1, 2026

Copy link
Copy Markdown

No description provided.

@coderabbitai

coderabbitai Bot commented Sep 1, 2026

Copy link
Copy Markdown

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true

Comment @coderabbitai help to get the list of available commands.

mt added 2 commits September 2, 2026 17:02
Signed-off-by: mt <mt@mt.localdomain>
Signed-off-by: mt <mt@mt.localdomain>

@zhangj1an zhangj1an left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks! It is already almost there, after addressing the following we can merge it.

Please refer to PR 204 (#204) as an example when adding new kernels.

  1. Test breakage: tests/test_rms_norm.py calls torch.musa.is_available() unguarded and fails with AttributeError on non-MUSA platforms (verified). Use device_ctx.is_musa or a hasattr guard.
  2. Please register musa kernels in rl_engine/kernels/gtest/operator_specs.py ("musa" candidates), docs/operators/*.md Backends rows, rl_engine/_C.pyi stubs for the new symbols, and a benchmark script.
  3. Bitwise evidence: for copy/elementwise ops (embedding/rope/silu/swiglu) add torch.equal bitwise tests vs the same-dtype PyTorch reference (fp32 internally, single cast at the end); for reductions (logp/lm_head/rmsnorm/attention/det_gemm) add bitwise batch-invariance tests + contract-tolerance accuracy tests, mirroring tests/test_batch_invariant_logp.py.
  4. Please rebase your branch on the test branch and change the PR base to test. Also update TestMusaPlatform expectations (the native-first musa map conflicts with the current assertions on both main and test), and coordinate the musa priority map with #375 (decide merge order / final native-first + triton-fallback map).
  5. Fix DCO sign-off and the unreachable device_ctx.is_musa branch in tests/test_logp.py.

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.

2 participants