warp specializied tma persistent kernel, step-4c, add iter grouped warp reduction #4315
+284
−56
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
|
!test |
|
Review updated until commit ba0b6d6 Description
Changes walkthrough 📝
PR Reviewer Guide 🔍Here are some key observations to aid the review process:
|
|
!test |
|
!test |
|
!test |
|
!test |
|
!test |
|
!test |
|
!test --dev |
liqiangxl
force-pushed
the
llu/ws_tma_4c_grouped
branch
from
5a9216c to
bcc6d91
Compare
|
!test |
|
!test --dev |
|
Fusion IR CUDA Kernel: |
|
!test |
liqiangxl
commented
Apr 30, 2025
| constexpr unsigned int align_size = sizeof(T) * N; | ||
| static_assert(align_size <= 16, "max allowed vect r/w is 16 bytes"); | ||
| // [warp_idx, N] | ||
| // [w0r0, w0r1, w0r2, w0r3, w1r0, w1r1, w1r2, w1r3] |
There was a problem hiding this comment.
The smem access pattern is illustrated in the figures:
jjsjann123
reviewed
May 2, 2025
| for (auto cached_tv : cached_inputs) { | ||
| if (cached_tv->hasBroadcast() && | ||
| is_redu_mapped_to_bcast(inner_reference_tv, cached_tv)) { | ||
| cached_tv->axis(2)->parallelize(ParallelType::Vectorize); |
There was a problem hiding this comment.
We should check for contiguity flag as well as allocation domain for vectorization.
There was a problem hiding this comment.
Good point. Heuristic ensures the iteration dim is divisible by the unroll factor. Here, we only need to further confirm all the iteration domains are contiguous. Also extended tests to test case with non-contig input.
auto can_vectorize = [](TensorView* redu_tv, TensorView* bcast_tv) {
const auto& alloc_dom_1 = redu_tv->getMaybeAllocationDomain();
const auto& alloc_dom_2 = bcast_tv->getMaybeAllocationDomain();
if (alloc_dom_1.size() != alloc_dom_2.size()) {
return false;
}
const auto& contiguity = bcast_tv->domain()->contiguity();
for (int i = 0; i < (int)alloc_dom_1.size(); i++) {
if (alloc_dom_1[i]->isReduction()) {
break;
}
if (!contiguity[i].has_value() || !contiguity[i].value()) {
return false;
}
}
return true;
};
There was a problem hiding this comment.
Thanks a lot for patching this.
A follow up ask is, can we add a test covering this case, when the inner most dimension is not contiguous.
csrc/codegen.cpp
Outdated
| std::pair<IterDomain*, IterDomain*> reduction_dims, | ||
| bool is_all_reduce) { | ||
| NVF_ERROR( | ||
| is_all_reduce, |
There was a problem hiding this comment.
This feels strange to pass an arg and assert on it.
There was a problem hiding this comment.
moved to before function call.
NVF_ERROR(
grouped_rop->isAllreduce(),
"iterGroupedStaticWarpAllReduce should be used for allreduce.");
return genGroupedWarpReduction(
csrc/codegen.cpp
Outdated
| func_args.arg(genStaticCast(genPtrType(output->dtype()), "shared_mem")); | ||
|
|
||
| ArgumentBuilder template_args; | ||
| if (reduction_dims.first->getParallelType() == ParallelType::TIDx && |
There was a problem hiding this comment.
nitpick, assert or check reduction_dims.first != nullptr
csrc/codegen.cpp
Outdated
| << ";\n"; | ||
| } else { | ||
| NVF_THROW( | ||
| "Grouped warp reduction is only supported for TIDx reduction with no second dimension"); |
There was a problem hiding this comment.
nitpick, if we are just throwing here, might as well convert this to an assert to avoid branching.
There was a problem hiding this comment.
good point. Moved to the out of the function.
NVF_ERROR(
lparams_.bdimx() % 128 == 0,
"iterGroupedStaticWarpAllReduce() requires bdimx % 128 == 0.");
NVF_ERROR(
grouped_rop->isAllreduce(),
"iterGroupedStaticWarpAllReduce should be used for allreduce.");
NVF_ERROR(
reduction_dims.first->getParallelType() == ParallelType::TIDx &&
reduction_dims.second == nullptr,
"Grouped warp reduction is only supported for TIDx reduction with no second dimension.");
return genGroupedWarpReduction(
(int)num_grouped_iterations,
output,
input,
grouped_rop->initVal(0),
op_type,
grouped_rop->predicate());
runtime/warp.cu
Outdated
|
|
||
| // sizeof(T) * K = sizeof(uint64_t) | ||
| // require alginment of sizeof(T) * K to safely cast between T and uint64_t | ||
| // shfl uses uint64_t to reduce number of shuffles |
There was a problem hiding this comment.
nitpick question: since this is runtime file, I don't think we can do concepts there, can we?
There was a problem hiding this comment.
didn't check that, but // Array structure ensures data is aligned for safe casting to uint64_t. So we actually don't need to add extra checks.
template <typename scalar_t, int size, int align_size = 1>
struct alignas(sizeof(scalar_t) * align_size) Array {
|
!test |
jjsjann123
reviewed
May 2, 2025
| const T inp_val[N], | ||
| Func reduction_op, | ||
| T* shared_mem, | ||
| uint32_t threadIdx_x, |
There was a problem hiding this comment.
dumb question, why are we passing threadIdx_x as arg, instead of just using threadIdx.x as-is here?
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.
This PR follows #4288. It is step-4cd of implementing warp specializied tma persistent kernel described in the design doc.
Changes:
(1) Add iteration grouped warp reduction, where broadcast is also fused with reduction, step-4c.
(2) Vectorized load of cached input with inner bcast, step-4d.
Details of iteration grouped warp reduction:
Input: Each thread has
Uelements with typteT, whereUis the iteration dim unroll factor.Algorithm:
packedWarpReduce(), where elements of typeTare packed intouint64_tto minimize shuffle instructions during warp reduction.Output: All threads have a copy of the reduction results.
Fusion IR changes
Using
NVFUSER_DUMP=fusion_ir,cuda_to_file ./test_nvfuser --gtest_filter=TmaWarpSpecializedTest.SimpleFusion/ws_1_dtype___bfloat_batch_1056_hidden_4096 2>&1 |tee 1.logas an example.The major change is the grouped reduction, e.g.
T5_l_float[iblockIdx.y72{132}, iS71{( ceilDiv(( ceilDiv(i0, 2) ), 132) )}, iG70{2}, rthreadIdx.x68{( ceilDiv(( ceilDiv(i2, 4) ), 8) )}_p] ca_pos( 2 ) produce_pos( 2 )CUDA code changes
(1) grouped reduction, it uses static CTA shape and fused with broadcast.
(2) vectorized load of the inner broadcast Tv (only when it exists, e.g. RMS Norm Bwd)
It is transformed as: