Repository navigation
feat(rccl-tests): GIN SDMA ReduceScatter (-D 3) #11556
There are no files selected for viewing
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
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
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
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
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,166 @@ | ||
| /************************************************************************* | ||
| * Copyright (c) 2026 Advanced Micro Devices, Inc. All rights reserved. | ||
| * | ||
| * See LICENSE.txt for license information | ||
| ************************************************************************/ | ||
|
|
||
| // Device-side reduction ops for the GIN-SDMA ReduceScatter (deviceImpl == 3). | ||
| // | ||
| // These MUST bit-match rccl-tests' correctness oracle (verifiable.cu). The | ||
| // verifier's compare tolerance is 0 for sum/prod/min/max on EVERY type (2 for | ||
| // floating avg, 1 for fp8_e5m2 sum with >1 rank -- see ncclVerifiableVerify), so | ||
| // each op here mirrors the corresponding Reduce* struct in verifiable.cu exactly: | ||
| // * low-precision types (half / bf16 / fp8) convert to float, apply the op, | ||
| // then narrow back to T on EVERY pairwise step (the verifier folds the same | ||
| // way, so accumulating in float and narrowing once would NOT bit-match); | ||
| // * the caller folds contributions in ascending source-rank order -- the same | ||
| // order ncclVerifiablePrepareExpected uses; | ||
| // * avg premultiplies each element by 1/rankN for floating T (postOp identity) | ||
| // but sums-then-divides for integral T (preOp identity), exactly as | ||
| // ReduceAvg_Base<T, integral> does. | ||
| // | ||
| // mulsum / PreMulSum is intentionally NOT supported (deferred); the dispatch | ||
| // macro (SPECIALIZE_REDUCE_KERNEL in reduce_scatter.cu) returns nullptr for it, | ||
| // and fp8 prod is excluded there too. | ||
| // | ||
| // This header is device-only: include it AFTER common.h (which pulls in rccl.h | ||
| // for __half / hip_bfloat16 and rccl_float8.h for the fp8 types) so all element | ||
| // types and the HAVE_BF16 / HAVE_FP8 guards are in scope. | ||
|
|
||
| #ifndef GIN_SDMA_REDUCE_H_ | ||
| #define GIN_SDMA_REDUCE_H_ | ||
|
|
||
| #if defined(ENABLE_DEVICE_API) && NCCL_VERSION_CODE >= NCCL_VERSION(2,28,0) | ||
|
|
||
| #include <type_traits> | ||
| #if HAVE_BF16 | ||
| #include <hip/hip_bfloat16.h> // hip_bfloat16 (same type verifiable.cu folds with) | ||
| #endif | ||
|
|
||
| namespace gin_sdma_reduce { | ||
|
|
||
| // ------------------------------- sum (a + b) ------------------------------- | ||
| // Template covers all integer types (int8/uint8/int32/uint32/int64/uint64) plus | ||
| // float/double; the explicit overloads below are more specialized so they win | ||
| // for the low-precision types (matching verifiable.cu::ReduceSum). | ||
| template <class T> __device__ __forceinline__ T rsum(T a, T b) { return a + b; } | ||
| __device__ __forceinline__ __half rsum(__half a, __half b) { | ||
| return __float2half(__half2float(a) + __half2float(b)); | ||
| } | ||
| #if HAVE_BF16 | ||
| __device__ __forceinline__ hip_bfloat16 rsum(hip_bfloat16 a, hip_bfloat16 b) { | ||
| return hip_bfloat16(static_cast<float>(a) + static_cast<float>(b)); | ||
| } | ||
| #endif | ||
| #if HAVE_FP8 | ||
| __device__ __forceinline__ rccl_float8 rsum(rccl_float8 a, rccl_float8 b) { | ||
| return rccl_float8(static_cast<float>(a) + static_cast<float>(b)); | ||
| } | ||
| __device__ __forceinline__ rccl_bfloat8 rsum(rccl_bfloat8 a, rccl_bfloat8 b) { | ||
| return rccl_bfloat8(static_cast<float>(a) + static_cast<float>(b)); | ||
| } | ||
| #endif | ||
|
|
||
| // ------------------------------- prod (a * b) ------------------------------ | ||
| template <class T> __device__ __forceinline__ T rprod(T a, T b) { return a * b; } | ||
| __device__ __forceinline__ __half rprod(__half a, __half b) { | ||
| return __float2half(__half2float(a) * __half2float(b)); | ||
| } | ||
| #if HAVE_BF16 | ||
| __device__ __forceinline__ hip_bfloat16 rprod(hip_bfloat16 a, hip_bfloat16 b) { | ||
| return hip_bfloat16(static_cast<float>(a) * static_cast<float>(b)); | ||
| } | ||
| #endif | ||
| #if HAVE_FP8 | ||
| __device__ __forceinline__ rccl_float8 rprod(rccl_float8 a, rccl_float8 b) { | ||
| return static_cast<rccl_float8>(static_cast<float>(a) * static_cast<float>(b)); | ||
| } | ||
| __device__ __forceinline__ rccl_bfloat8 rprod(rccl_bfloat8 a, rccl_bfloat8 b) { | ||
| return static_cast<rccl_bfloat8>(static_cast<float>(a) * static_cast<float>(b)); | ||
| } | ||
| #endif | ||
|
|
||
| // ------------------------------- min / max --------------------------------- | ||
| template <class T> __device__ __forceinline__ T rmin(T a, T b) { return a < b ? a : b; } | ||
| template <class T> __device__ __forceinline__ T rmax(T a, T b) { return a > b ? a : b; } | ||
| __device__ __forceinline__ __half rmin(__half a, __half b) { | ||
| return __half2float(a) < __half2float(b) ? a : b; | ||
| } | ||
| __device__ __forceinline__ __half rmax(__half a, __half b) { | ||
| return __half2float(a) > __half2float(b) ? a : b; | ||
| } | ||
| #if HAVE_BF16 | ||
| __device__ __forceinline__ hip_bfloat16 rmin(hip_bfloat16 a, hip_bfloat16 b) { | ||
| return static_cast<float>(a) < static_cast<float>(b) ? a : b; | ||
| } | ||
| __device__ __forceinline__ hip_bfloat16 rmax(hip_bfloat16 a, hip_bfloat16 b) { | ||
| return static_cast<float>(a) > static_cast<float>(b) ? a : b; | ||
| } | ||
| #endif | ||
| #if HAVE_FP8 | ||
| __device__ __forceinline__ rccl_float8 rmin(rccl_float8 a, rccl_float8 b) { | ||
| return static_cast<float>(a) < static_cast<float>(b) ? a : b; | ||
| } | ||
| __device__ __forceinline__ rccl_float8 rmax(rccl_float8 a, rccl_float8 b) { | ||
| return static_cast<float>(a) > static_cast<float>(b) ? a : b; | ||
| } | ||
| __device__ __forceinline__ rccl_bfloat8 rmin(rccl_bfloat8 a, rccl_bfloat8 b) { | ||
| return static_cast<float>(a) < static_cast<float>(b) ? a : b; | ||
| } | ||
| __device__ __forceinline__ rccl_bfloat8 rmax(rccl_bfloat8 a, rccl_bfloat8 b) { | ||
| return static_cast<float>(a) > static_cast<float>(b) ? a : b; | ||
| } | ||
| #endif | ||
|
|
||
| // ---------------------------- avg premul factor ---------------------------- | ||
| // The 1/rankN factor, built in the same intermediate precision the verifier | ||
| // uses (ReduceAvg_Base<T,false>: T1 = float when sizeof(T) < sizeof(double), | ||
| // else double). Only used on the floating branch; integral avg divides in postOp. | ||
| template <class T> __device__ __forceinline__ T avgFactor(int rankN) { | ||
| return (T)(1.0f / (float)rankN); | ||
| } | ||
| template <> __device__ __forceinline__ double avgFactor<double>(int rankN) { | ||
| return 1.0 / (double)rankN; | ||
| } | ||
|
|
||
| // ------------------------ op-tagged pre / combine / post -------------------- | ||
| // `op` is an ncclRedOp_t value (ncclSum=0, ncclProd=1, ncclMax=2, ncclMin=3, | ||
| // ncclAvg=4). PreMulSum is never dispatched here. | ||
|
|
||
| template <class T> __device__ __forceinline__ T preOp(int op, T x, int rankN) { | ||
| // Floating avg premultiplies each element by 1/rankN before summing; integral | ||
| // avg (and every other op) leaves the element untouched (divides in postOp). | ||
| if (op == ncclAvg && !std::is_integral<T>::value) { | ||
| return rprod(avgFactor<T>(rankN), x); | ||
| } | ||
| return x; | ||
| } | ||
|
|
||
| template <class T> __device__ __forceinline__ T combine(int op, T a, T b) { | ||
|
Contributor
There was a problem hiding this comment. Nothing unit-tests
Contributor
Author
There was a problem hiding this comment. Fair, and this is the one I'd pick first out of the style items. Bit-matching Adding the
Contributor
Author
There was a problem hiding this comment. Agreed, and this is the coverage gap I would most want closed. |
||
| switch (op) { | ||
| case ncclProd: return rprod(a, b); | ||
| case ncclMax: return rmax(a, b); | ||
| case ncclMin: return rmin(a, b); | ||
| default: return rsum(a, b); // ncclSum and ncclAvg both accumulate by + | ||
| } | ||
| } | ||
|
|
||
| // Integral avg finishes with an integer divide of the (wrapped) sum; floating T | ||
| // already divided in preOp. Split via SFINAE so `x / rankN` is only instantiated | ||
| // for integral T (fp8/bf16/half have no unambiguous `operator/(T,int)`). | ||
| template <class T> __device__ __forceinline__ | ||
| typename std::enable_if<std::is_integral<T>::value, T>::type | ||
| postDivAvg(T x, int rankN) { return (T)(x / rankN); } | ||
| template <class T> __device__ __forceinline__ | ||
| typename std::enable_if<!std::is_integral<T>::value, T>::type | ||
| postDivAvg(T x, int /*rankN*/) { return x; } | ||
|
|
||
| template <class T> __device__ __forceinline__ T postOp(int op, T x, int rankN) { | ||
| if (op == ncclAvg) return postDivAvg<T>(x, rankN); | ||
| return x; | ||
| } | ||
|
|
||
| } // namespace gin_sdma_reduce | ||
|
|
||
| #endif // ENABLE_DEVICE_API && >= 2.28 | ||
| #endif // GIN_SDMA_REDUCE_H_ | ||
Oops, something went wrong.
Oops, something went wrong.
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.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
rsum,rprod,rminandrmaxrepeat the same widen-to-float-then-narrow body 16 times over the four low-precision types, andrprodhas already drifted: bf16 useship_bfloat16(...)construction at line 71 while fp8 usesstatic_cast<rccl_float8>(...)at 76, wherersumuses construction for both. Same meaning, but one per-type macro emitting all four would take this from about 60 lines to 14 and make the drift impossible.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Confirmed the drift you describe:
rprodconstructs for bf16 (line 71) butstatic_casts for fp8 (line 76), wherersumconstructs for both.Both forms are correct today — same widen-to-float, same result — so this is a readability and
future-drift concern rather than a live bug, which is why I'm not changing numeric code late in
review. Folding it into the single pass over this header together with the
GIN_SDMA_HOST_ONLYescape you asked for at line 139, so the macro rewrite lands with teststhat can actually verify it didn't change results.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agreed, including the drift:
rprodconstructship_bfloat16(...)for bf16 but takes a different route for fp8, which is exactly the kind of divergence a shared widen/narrow helper would prevent. Leaving open under AICOMRCCL-1803, and it pairs with the missing bit-match tests in thegin_sdma_reduce.hthread.