[JAX] Collective GEMM custom op + primitive + minimal supporting functions #1846
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
Signed-off-by: Alp Dener <[email protected]> CollectiveGemm XLA FFI op and FWD+BWD rules done Signed-off-by: Alp Dener <[email protected]> added missing custom VJP def Signed-off-by: Alp Dener <[email protected]> added comm overlap initializer/destroy Signed-off-by: Alp Dener <[email protected]> docstrings Signed-off-by: Alp Dener <[email protected]> compile errors resolved Signed-off-by: Alp Dener <[email protected]> added example Signed-off-by: Alp Dener <[email protected]> fixed compile and import issues Signed-off-by: Alp Dener <[email protected]> retooled how casting and transposes are done in FWD/BWD Signed-off-by: Alp Dener <[email protected]> some API tweaks/cleanup Signed-off-by: Alp Dener <[email protected]> added missing TE DType to pybind helper, and restored missing comm overlap config keys Signed-off-by: Alp Dener <[email protected]> fixed incorrect variable assignment order in bootstrapping Signed-off-by: Alp Dener <[email protected]> XLA custom op working, cublas unsupported parameter error underneath Signed-off-by: Alp Dener <[email protected]> fixed cublas invalid parameter error Signed-off-by: Alp Dener <[email protected]> both AG and RS overlaps tested working in BF16 Signed-off-by: Alp Dener <[email protected]> removed FWD/BWD impl to shrink PR down to the XLA custom op, JAX primitive, and minimum accompanying Python wrappers Signed-off-by: Alp Dener <[email protected]> fixed incorrect rowwise/columnwise placement of data and scales in XLA custom op Signed-off-by: Alp Dener <[email protected]> reverted some stuff Signed-off-by: Alp Dener <[email protected]>
This was referenced Jun 3, 2025
Signed-off-by: Alp Dener <[email protected]>
denera
force-pushed
the
jax/collective-gemm-api
branch
from
1a845e9 to
e92c81a
Compare
for more information, see https://pre-commit.ci
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.
Description
This PR introduces a new XLA custom op for calling
nvte_cublas_gemmor related comm+GEMM overlap algorithms, the accompanying JAX primitive, and bare minimum Python wrappers required to work with the custom call.FWD/BWD autograd implementation will be tackled in a separate upcoming PR.
Type of change
Checklist: