Skip to content

[AIXLA-150] Enable jax-triton kernels for ROCm #2

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Open
cschenjunlin wants to merge 72 commits into
base: main
Choose a base branch
from
Open

[AIXLA-150] Enable jax-triton kernels for ROCm #2

cschenjunlin wants to merge 72 commits into from

Conversation

@cschenjunlin
Copy link
Contributor

@cschenjunlin cschenjunlin commented Aug 12, 2025 •
edited by i-chaochen
Loading

Motivation

Enable existing examples for ROCm.
Enable Triton tutorials for ROCm.

Performance Comparison of Jax-Triton and Jax-Pallas Kernels

https://amd.atlassian.net/wiki/x/OBQDQg

Technical Details

Test Plan

Test Result

Submission Checklist

@cschenjunlin cschenjunlin changed the title rocm support Enable jax-triton kernels for ROCm Aug 13, 2025
@cschenjunlin
triton tutorial kernel 03 for rocm
31f8d00
@cschenjunlin
enable jax-triton kernel of seeded-dropout for rocm
53028af
@cschenjunlin
rocprofv3 performance analysis reports of add, softmax, matmul, dropo…
5e260a8 
…ut kernels
@cschenjunlin
jax-triton kernel of layer-norm for rocm
8d6f20d
@cschenjunlin
jax-triton kernel of layer-norm for rocm
57d8d58
@cschenjunlin
jax-triton kernel of layer-norm-fwd
0658371
@cschenjunlin
compare vec-add kernels of jax-pallas and jax-triton by rocprofv3
6e60074
@cschenjunlin
compare vec-add kernels of jax-pallas and jax-triton by rocprofv3
8e3193a
@cschenjunlin
compare vec-add kernels of jax-pallas and jax-triton by rocprofv3
a0eea84
@cschenjunlin
rocprof-compute performance analysis report of jax-triton kernel of v…
5411e39 
…ec-add
@cschenjunlin
rocprof-compute performance analysis report of jax-triton kernel of v…
b0a33c7 
…ec-add
@cschenjunlin
enable jax-pallas kernel of softmax
074cff6
@cschenjunlin
Copy link
Contributor Author

cschenjunlin commented Aug 26, 2025

Performance Comparison of Jax-Triton and Jax-Pallas Kernels
https://amd.atlassian.net/wiki/x/OBQDQg

@cschenjunlin cschenjunlin changed the title Enable jax-triton kernels for ROCm [AIXLA-150] Enable jax-triton kernels for ROCm Nov 15, 2025
@cschenjunlin
performance analysis for jax-triton, jax-ref, aiter-triton
dde799a
@cschenjunlin
finish mha_fwd and test
9bb42e0
@cschenjunlin
separate mha_bwd test from kernel
301cfcb
@cschenjunlin
remove some indent
c300a5a
@cschenjunlin
update .gitignore
44e4cb1
@cschenjunlin
finish mha_onekernel_bwd and test
3ed8c48
@cschenjunlin
finish mha_fused_bwd and test
632dec5
@cschenjunlin
unittest for mha
e978868
@cschenjunlin
rearrange some params
0b39042
@cschenjunlin
raw jax-pallas kernel
75dbfae
@cschenjunlin
add jax-pallas mha kernels
edf1df4
@cschenjunlin
fwd and bwd test passed
4bfb5b0
@cschenjunlin
add jax-pallas mha kernels
04be728
@cschenjunlin
add for
52b66d1
@cschenjunlin
pallas_mha test passed
1c86720
@cschenjunlin
add @jax.jit for mha host functions
d1436cf
@cschenjunlin
add synchronization for mha test
df58c99
@cschenjunlin
add mha kernels into jax_triton package
2138cf9
@cschenjunlin
support kernel import as package; update configs and static args
6a49db1
@cschenjunlin
add user interface for jax_triton.ops.mha
c2225f3
@cschenjunlin
jax_triton.op_tests: add unittest for mha
da0c24a
@cschenjunlin
use fixed dtype in einsum in jax_attn_reference
4101a04
@cschenjunlin
skip some unsupported cases
8c44355
@cschenjunlin
examples: use mha kernels from jax_triton; op_tests: TestCase for mha
126224b
@cschenjunlin
skip FUSED test, resulting in synchronization issues
6fe2196
i-chaochen
i-chaochen reviewed Dec 31, 2025
View reviewed changes
build_jax_rocm.sh
Copy link

@i-chaochen i-chaochen Dec 31, 2025

Choose a reason for hiding this comment

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

we don't need this build script in jax-triton project

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@i-chaochen i-chaochen i-chaochen left review comments

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

2 participants

@cschenjunlin @i-chaochen