Update README and fix lint #23
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
meta-cla
bot
added
the
CLA Signed
|
@kwen2501 I've written a README and fixed some lint for releasing this at the PTC. Surya and I are working on the poster now (deadline is end of this week). Ping me if you want to talk about cleaning up this repo and what to put in the poster! |
kwen2501
reviewed
Sep 25, 2025
| return torch_symm_mem_ag_mm | ||
| if backend == "triton": | ||
| return kraken.all_gather.all_gather_matmul | ||
| return kraken.all_gather_fusion.all_gather_matmul |
There was a problem hiding this comment.
I think we can arrange ops in a couple categories:
kraken.comms: pure comm ops, liketwo_shot_all_reduce,all_to_all_on_device, etc.kraken.fused: here we put ops likeall_gather_matmulkraken.quantized: here we put ops with quantization, such as [mxfp8 moe training] mxfp8_on_device_all_to_all_v kernel pytorch/ao#3048
wdyt?
There was a problem hiding this comment.
@danielvegamyhre we created a bucket for you -- kraken.quantized :)
There was a problem hiding this comment.
sweet! happy to add the kernel here as well, as a standalone building block/example.
any plans to import and use kraken in torchtitan?
…g soon) quantized kernels
kwen2501
approved these changes
Sep 25, 2025
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
4 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.