Add option to use triton tsyrk_ex kernels in muon - #49
Closed
skyw wants to merge 13 commits into
Closed
Conversation
skyw
force-pushed
the
skyw/support_syrk_in_ns_step
branch
from
October 9, 2025 02:47
5c91bdb to
32110a0
Compare
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
skyw
force-pushed
the
skyw/support_syrk_in_ns_step
branch
from
October 9, 2025 02:48
32110a0 to
ef94309
Compare
Contributor
Author
|
/ok to test ef94309 |
Signed-off-by: Hao Wu <skyw@nvidia.com>
Contributor
Author
|
/ok to test d6c89c9 |
Signed-off-by: Hao Wu <skyw@nvidia.com>
Contributor
Author
|
/ok to test c8f80ae |
gdengk
reviewed
Oct 9, 2025
| python tests/test_scalar_optimizers.py --device=cuda | ||
| python tests/test_spectral_clipping_utils.py | ||
| python tests/test_triton_kernels.py TritonKernelsIntegerInputTest | ||
| python tests/test_triton_kernels.py |
Contributor
There was a problem hiding this comment.
Tests of float dtype will fail. So such failure are allowed in L1 Tests?
| Output tensor of shape (N, N) | ||
| """ | ||
|
|
||
| assert a.dtype == torch.bfloat16, "Input tensor must be bfloat16" |
Contributor
There was a problem hiding this comment.
I think this assertion is still needed.
| logging.log_first_n(logging.INFO, "Using BF16 I/O kernels for Newton-Schulz iteration.", 1) | ||
| if use_syrk: | ||
| sm_version = torch.cuda.get_device_capability() | ||
| if sm_version in ((8, 0), (9, 0), (10, 0), (11, 0)): |
Contributor
There was a problem hiding this comment.
Feel it's better to be placed in the init function otherwise this is getting called for every param.
Signed-off-by: Hao Wu <skyw@nvidia.com>
Signed-off-by: Hao Wu <skyw@nvidia.com>
Contributor
Author
|
/ok to test dd5216b |
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
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
No description provided.