Skip to content

NVRTC kernels for cast-transpose - #258

Merged
timmoon10 merged 22 commits into
NVIDIA:mainfrom
timmoon10:nvrtc-cast-transpose
Apr 19, 2024
Merged

NVRTC kernels for cast-transpose#258
timmoon10 merged 22 commits into
NVIDIA:mainfrom
timmoon10:nvrtc-cast-transpose

Conversation

@timmoon10

Copy link
Copy Markdown
Member

This is a follow-up to #138 that adds NVRTC kernels for cast-transpose.

I've also tweaked the kernel selection heuristics to be a bit more general. There are some performance differences at intermediate scale, but nothing too drastic:

Transpose kernel times

image

Evaluated FP8 transposes on an H100 PCIe with square matrices with nice dims (power of two or halfway between powers of two).

@timmoon10timmoon10 added the enhancement New feature or request label Jun 1, 2023
@timmoon10
timmoon10 requested a review from ptrendxJune 1, 2023 21:22
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10
timmoon10force-pushed the nvrtc-cast-transpose branch from dfc121b to aed892aCompareJune 1, 2023 21:24
@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

If there are no objections, I will merge this after the 1.6 release.

Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <tmoon@nvidia.com>
@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

Comment threadtransformer_engine/common/transpose/cast_transpose.cu
Comment threadtransformer_engine/common/transpose/cast_transpose.cu
const size_t num_tiles) {
__launch_bounds__(block_size)
cast_transpose_general_kernel(const IType * __restrict__ const input,
const CType * const noop,

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Was the restrict keyword intentionally omitted here?

Comment threadtransformer_engine/common/transpose/transpose.cu
Comment threadtransformer_engine/common/transpose/transpose.cu

@Oleg-GoncharovOleg-Goncharov left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM, besides the early return from the kernel configuration constructors

Comment threadtransformer_engine/common/transpose/cast_transpose.cu Outdated
Comment threadtransformer_engine/common/transpose/rtc/cast_transpose.cu Outdated
Comment threadtransformer_engine/common/transpose/cast_transpose.cu
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
@timmoon10

Copy link
Copy Markdown
MemberAuthor

/te-ci

@timmoon10
timmoon10 merged commit 14c1ecd into NVIDIA:mainApr 19, 2024
pggPL pushed a commit to pggPL/TransformerEngine that referenced this pull request May 23, 2024
* Add NVRTC kernels for cast-transpose
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Update copyright year
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Add noop flag to NVRTC cast-transpose kernel
Signed-off-by: Tim Moon <tmoon@nvidia.com>
* Apply suggestions from code review
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
---------
Signed-off-by: Tim Moon <tmoon@nvidia.com>
Signed-off-by: Tim Moon <4406448+timmoon10@users.noreply.github.com>
Signed-off-by: Pawel Gadzinski <pgadzinski@nvidia.com>
Sign up for freeto join this conversation on GitHub. Already have an account? Sign in to comment

Labels

enhancementNew feature or request

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants

@timmoon10@Oleg-Goncharov