Skip to content

Current scaling: two-stage amax kernel #369

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 13 commits into
base: dev
Choose a base branch
from
Draft

Current scaling: two-stage amax kernel #369

wants to merge 13 commits into from
+170 −11

Conversation

@matthiasdiener
Copy link
Contributor

@matthiasdiener matthiasdiener commented Nov 12, 2025 •
edited
Loading

Description

Implements a two-stage HIP kernel for the amax operation, as an alternative to the original implementation that uses atomic reductions. Make the two-stage kernel the default implementation. Users can use export NVTE_USE_ATOMIC_AMAX=1 to use the atomic amax kernel.

Fixes https://github.com/ROCm/frameworks-internal/issues/14303.

See https://github.com/ROCm/frameworks-internal/issues/14303#issuecomment-3554900809 for a performance analysis.

TODO:

  • Fix other call sites of nvte_compute_amax
  • Address FIXMEs in the code

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 matthiasdiener self-assigned this Nov 14, 2025
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@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

@ipanfilo ipanfilo Awaiting requested review from ipanfilo ipanfilo will be requested when the pull request is marked ready for review ipanfilo 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.

1 participant

@matthiasdiener