Skip to content

[WIP] Speedup amax triton #385

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
Draft
wants to merge 35 commits into
base: dev
Choose a base branch
from
Draft

[WIP] Speedup amax triton #385

wants to merge 35 commits into from

Conversation

@matthiasdiener
Copy link
Contributor

@matthiasdiener matthiasdiener commented Nov 26, 2025

Description

Please include a brief summary of the changes, relevant motivation and context.

Fixes # (issue)

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:

  • Change A
  • Change B

Checklist:

  • I have read and followed the contributing guidelines
  • 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

matthiasdiener and others added 30 commits November 12, 2025 14:10
@matthiasdiener
Current scaling: two-stage amax kernel
c15d93b
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
51fab36
@matthiasdiener
bugfix graph capture
ae35e4c
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
77a68a7
@matthiasdiener
outline workspace allocation
c0d8e73
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
6c3507d
@matthiasdiener
Proper allocation of workspace
3c9de07
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
91249cc
@matthiasdiener
add a test to compare the accuracy of both amax implementations
be0e0c8
@matthiasdiener
add possibility to force using previous (atomic) kernel
bce34da
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
8c388cc
@matthiasdiener
2-stage Triton amax
73c8d4e
@matthiasdiener
add copyrights
6388604
@matthiasdiener
don't add extra template to kernel
9e6586f
@matthiasdiener
make amax_kernel_threads usable in pytorch
18292bf
@matthiasdiener
update remaining calls to nvte_compute_amax
a389455
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
d87ab8a
@matthiasdiener
Merge branch 'speedup-amax-kernel' into speedup-amax-triton
7d9ee16
@matthiasdiener
additional copyrights
fd5dead
@matthiasdiener
avoid workspace allocations if NVTE_USE_ATOMIC_AMAX is set
16d3bf9
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
50b34aa
@matthiasdiener
remove use_block_amax parameter, more cleanups
ef532b1
@matthiasdiener
Factor workspace allocation into function
f933ef3
@matthiasdiener
expand test slightly
7d4054e
Revert "expand test slightly"
63cff98 
This reverts commit 7d4054e.
@matthiasdiener
guard by HIP macro, address review comments
c7d44a7
@matthiasdiener
bugfix workspace.data.dptr
f92b926
@matthiasdiener
various cleanups
eba552e
@matthiasdiener
Merge branch 'dev' into speedup-amax-kernel
0d6a177
@matthiasdiener
Merge branch 'speedup-amax-kernel' into speedup-amax-triton
19901a0
@matthiasdiener matthiasdiener self-assigned this Nov 26, 2025
@matthiasdiener matthiasdiener changed the base branch from dev to speedup-amax-kernel November 26, 2025 22:16
@matthiasdiener matthiasdiener changed the base branch from speedup-amax-kernel to dev November 26, 2025 22:53
ipanfilo
ipanfilo reviewed Nov 27, 2025
View reviewed changes
transformer_engine/common/recipe/current_scaling.cu
atomicMaxFloat(amax, max);
}
#else
atomicMaxFloat(amax, max);
Copy link
Collaborator

@ipanfilo ipanfilo Nov 27, 2025

Choose a reason for hiding this comment

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

This line looks incorrectly indented

@matthiasdiener matthiasdiener changed the title Speedup amax triton [WIP] Speedup amax triton Nov 27, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@ipanfilo ipanfilo ipanfilo left review comments

@wenchenvincent wenchenvincent Awaiting requested review from wenchenvincent wenchenvincent will be requested when the pull request is marked ready for review wenchenvincent is a code owner

@wangye805 wangye805 Awaiting requested review from wangye805 wangye805 will be requested when the pull request is marked ready for review wangye805 is a code owner

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

Assignees

@matthiasdiener matthiasdiener

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@matthiasdiener @ipanfilo