[Pytorch] Decoupling framework extensions from common module #1498
Conversation
This file contains 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
KshitijLakhani
force-pushed
the
klakhani/maint/decouple-te-in-fwexten
branch
from
99871f3 to
cd20471
Compare
|
/te-ci pytorch |
KshitijLakhani
force-pushed
the
klakhani/maint/decouple-te-in-fwexten
branch
from
7750ea4 to
c62be40
Compare
KshitijLakhani
commented
Feb 20, 2025
KshitijLakhani
commented
Feb 20, 2025
| #include "common/fused_attn/thd_utils.h" | ||
| #include "extensions.h" | ||
|
|
||
| using namespace transformer_engine::fused_attn; | ||
| #include "thd_util.cuh" |
There was a problem hiding this comment.
This is the new header file created in the framework extensions for the structs that were being used in attention.cu from transformer_engine/common/fused_attn/thd_utils.h.
KshitijLakhani
commented
Feb 20, 2025
| @@ -0,0 +1,59 @@ | |||
| /************************************************************************* | |||
There was a problem hiding this comment.
Is transformer_engine/common/include/transformer_engine a better place for this ?
|
/te-ci pytorch |
ptrendx
reviewed
Feb 20, 2025
ptrendx
reviewed
Feb 20, 2025
cyanguwa
reviewed
Feb 20, 2025
KshitijLakhani
force-pushed
the
klakhani/maint/decouple-te-in-fwexten
branch
from
adf5d8a to
56a8500
Compare
Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
…el kernels ONLY for invoking recompilation and not directly using the pre-compiled symbols in libtransformer.so Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
… common.h Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
…de header Code cleanup Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
for more information, see https://pre-commit.ci
Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
KshitijLakhani
force-pushed
the
klakhani/maint/decouple-te-in-fwexten
branch
from
40c3120 to
bd6182e
Compare
KshitijLakhani
changed the title
KshitijLakhani
changed the title
|
/te-ci pytorch |
cyanguwa
reviewed
Feb 21, 2025
| @@ -33,39 +86,78 @@ __forceinline__ __device__ int binary_search(int target, int *array, int len) { | |||
| /*************************************************************************************************** | |||
| * Support THD format for Context Parallel: Generate partitioned indices for input tokens | |||
| **************************************************************************************************/ | |||
|
|
|||
| // Templatizing this kernel ONLY so that when it is used in framework | |||
| // extensions, it is dynamically compiled | |||
There was a problem hiding this comment.
We could probably remove these comments now.
cyanguwa
reviewed
Feb 21, 2025
| #include <cuda.h> | ||
| #include <cuda_bf16.h> | ||
| /*************************************************************************************************** | ||
| * Support THD format for Context Parallel: softmax_lse related operations | ||
| **************************************************************************************************/ |
There was a problem hiding this comment.
Maybe have a cleanup of these comments. They seem unrelated to the functions below. Same for the Gradients ones.
…d_read_half_tensor_kernel Signed-off-by: Kshitij Janardan Lakhani <[email protected]> Code clean up Signed-off-by: Kshitij Janardan Lakhani <[email protected]>
KshitijLakhani
force-pushed
the
klakhani/maint/decouple-te-in-fwexten
branch
from
d32096e to
08b82a0
Compare
for more information, see https://pre-commit.ci
cyanguwa
approved these changes
Feb 21, 2025
ptrendx
approved these changes
Feb 21, 2025
13 tasks
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
Motivation: Currently, framework extensions has a dependency on
transformer_engine/common(libtransformer_engine.so) which is met via the symbols exposed bylibtransformer_engine.version. However, this could result in two problems:This PR is the first in a series of decoupling PRs to follow
Type of change
Changes
Files in the framework extensions identified to have this dependency in the first pass are -
In
pytorch/csrc/extensions:attention.cu,gemm.cpp,multi_tensor_adam.cuIn
jax/csrc/extensions:attention.cpp,utils.cuThis PR only decouples
pytorch/csrc/extensions/attention.cuby :nvte_*calls to replace the usage oftransformer_engine::Tensort_e/common/fused_attn/thd_util.hto a new filet_e/pytorch/csrc/thd_util.cuhthd_partition_indices_kernelandthd_read_half_tensor_kernelso that they are "re-compiled" on instantiation in the framework extensionsChecklist: