Skip to content
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

Add Triton MLA Decode Rope Kernel #232

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Add Triton MLA Decode Rope Kernel #232

wants to merge 1 commit into from
+1,530 −177

Conversation

lucas-santos-amd
Copy link
Contributor

@lucas-santos-amd lucas-santos-amd commented Mar 20, 2025 •
edited
Loading

  • Moved the _fwd_kernel_stage2_asm Triton Kernel to aiter/mla.py.
  • Renamed decode_mla.py to mla_decode_ref.py and moved it to op_tests/triton/utils. For now, it will be used as reference to the unit tests of both the ASM and the Triton MLA decode rope implementations.
  • Added the Triton MLA Decode Rope and the stage2 Kernels to the mla_decode_rope.py file.

@lucas-santos-amd
Add Triton MLA Decode Rope Kernel
e65a7b4 
- Moved the _fwd_kernel_stage2_asm Triton Kernel to aiter/mla.py.
- Renamed decode_mla.py to mla_decode_ref.py and moved it to
 op_tests/triton/utils. For now, it will be used as reference to the
 unit tests of both the ASM and the Triton MLA decode rope
 implementations.
- Added the Triton MLA Decode Rope and the stage2 Kernels to the
  mla_decode_ref.py file.
@lucas-santos-amd lucas-santos-amd force-pushed the lusantos/mla_decode_rope_triton branch from ad03939 to e65a7b4 Compare March 21, 2025 13:14
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@vgokhale vgokhale Awaiting requested review from vgokhale

@rahulbatra85 rahulbatra85 Awaiting requested review from rahulbatra85

@jayfurmanek jayfurmanek Awaiting requested review from jayfurmanek

At least 1 approving review is required to merge this pull request.

Assignees
No one assigned
Labels
triton
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
1 participant
@lucas-santos-amd