[Triton] DS fused custom ops #607
Merged
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
reviewed
Jul 8, 2025
rahulbatra85
requested changes
Jul 8, 2025
Contributor
rahulbatra85
left a comment
rahulbatra85
left a comment
There was a problem hiding this comment.
Please see comments. Thanks!
…mains for batched gemm
rahulbatra85
approved these changes
Jul 16, 2025
cagrikymk
added a commit
that referenced
this pull request
Jul 30, 2025
* add fused concat * add fused elementwise and pytest * clean up * add fused quant * clean up * add fused mxfp4 quant and pytest * add gemm_a16w16_atomic and related tests * extend tests and add config files for bf16 GEMM with atomic add * fused quant. code cleanup * formatting changes * add rms norm quant tests and changes * update/add DS configs for GEMMs * add prequant config * fix bf16 gemm * fix fp4 gemm atomic add * black reformatting * tune fp4 prequant gemm atomic * fix batched fp4 prequant and add new DS configs * optimize shapes for deepseek * black reformatting * update bf16 atomic GEMM * add documentation for fused_concat * update rope qk cat * rename fused elementwise * add doc for fused quant * rename fused_quant into fused_mxfp4_quant * fix typo * fix a minor bug * black formatting * update comments and drop unused arg. * fix pre-quant GEMM tests based on func. sign. changes * fix pre-quant GEMM tests based on func. sign. changes * doc on fused_mul_add * corner case fix * corner case fix * test * Fix pytest errors with LRU cache inplace mod - big test case error remains for batched gemm * Black linting change * Fix linting error * Move inplace config operations to _get_config --------- Co-authored-by: Mehmet Cagri Kaymak <mehmet.kaymak@amd.com> Co-authored-by: Lukasz Burzawa <lukasz.burzawa@amd.com> Co-authored-by: William Zhou <William.Zhou@amd.com>
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
This PR contains the following:
aiter/ops/triton/fused_elementwise.pyaiter/ops/triton/fused_concat.pyaiter/ops/triton/fused_quant.pyaiter/ops/triton/gemm_a16w16.pyconfigsaiter/ops/triton/batched_gemm_afp4wfp4_pre_quant.pyaiter/ops/triton/gemm_a16w16_atomic.py