diff --git a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu index 12a68e0c18..6933be7989 100644 --- a/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu +++ b/src/ops/sparse_moe/decode/sparse_moe_decode_kernels.cu @@ -153,7 +153,7 @@ struct W8Codec { } }; -template +template __device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const std::uint8_t* high, const std::uint8_t* scales, int row0, int row1, const __nv_bfloat16* x, int k_begin, int k_end, @@ -168,32 +168,64 @@ __device__ __forceinline__ void dot_two_rows(const std::uint8_t* codes, const st // Four adjacent Q4 groups form one 128-byte warp transaction. Each lane owns eight // consecutive K values, so one mantissa decode feeds eight FP32 FMAs instead of issuing // four scalar code-pair/decode iterations. + // + // QuadsInFlight quads are in flight at a time. All of their code words and scales are + // issued before any is decoded, so the decode and the FMA chain of one quad cover the + // memory latency of the next. One quad in flight leaves the load unit idle for the whole + // decode, which pays off where the path is latency-bound rather than bandwidth-bound: at + // the operator bench's trace-like point it demands 55% of the measured DRAM read ceiling. + // + // The two callers want different depths and get them. The T >= 2 path-tiled kernel takes + // the default 2; sparse_moe_d3_nine_warp_kernel, the T = 1 decode path, keeps 1, because + // at one token per launch it has a fifth of the occupancy and pairing costs it 6% in the + // round and 10% on the operator bench. + // + // The groups are still visited in ascending order at either depth, so the order of + // additions into acc0 and acc1 is unchanged and the output is bit-identical. + constexpr int kQuadsInFlight = QuadsInFlight; + // What the pairing needs is that the span this loop walks, (k_end - k_begin) / + // Codec::kGroupK, is a multiple of 4 * kQuadsInFlight. That is a property of the + // arguments and cannot be asserted here; every caller passes k_begin = 0 and + // k_end = K, which reduces it to the compile-time condition below. + static_assert(kGroups % (4 * kQuadsInFlight) == 0); const int lane_group = lane >> 3; const int lane_in_group = lane & 7; - for (int group_base = first_group; group_base < last_group; group_base += 4) { - const int group = group_base + lane_group; - const std::int64_t index0 = static_cast(row0) * kGroups + group; - const std::int64_t index1 = static_cast(row1) * kGroups + group; - const std::uint32_t packed0 = - *reinterpret_cast(codes + index0 * 32 + lane_in_group * 4); - const std::uint32_t packed1 = - *reinterpret_cast(codes + index1 * 32 + lane_in_group * 4); - const auto scale0 = *reinterpret_cast(scales + index0 * 2); - const auto scale1 = *reinterpret_cast(scales + index1 * 2); - float weights0[8]; - float weights1[8]; - Q4SimtDecodeAtom::decode_eight(packed0, scale0, weights0); - Q4SimtDecodeAtom::decode_eight(packed1, scale1, weights1); - const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); - const float2 x0 = bf16x2_bits_to_float2(input.x); - const float2 x1 = bf16x2_bits_to_float2(input.y); - const float2 x2 = bf16x2_bits_to_float2(input.z); - const float2 x3 = bf16x2_bits_to_float2(input.w); - const float values[8] = {x0.x, x0.y, x1.x, x1.y, x2.x, x2.y, x3.x, x3.y}; + for (int group_base = first_group; group_base < last_group; + group_base += 4 * kQuadsInFlight) { + std::uint32_t packed0[kQuadsInFlight]; + std::uint32_t packed1[kQuadsInFlight]; + std::uint16_t scale0[kQuadsInFlight]; + std::uint16_t scale1[kQuadsInFlight]; +#pragma unroll + for (int quad = 0; quad < kQuadsInFlight; ++quad) { + const int group = group_base + quad * 4 + lane_group; + const std::int64_t index0 = static_cast(row0) * kGroups + group; + const std::int64_t index1 = static_cast(row1) * kGroups + group; + packed0[quad] = *reinterpret_cast(codes + index0 * 32 + + lane_in_group * 4); + packed1[quad] = *reinterpret_cast(codes + index1 * 32 + + lane_in_group * 4); + scale0[quad] = *reinterpret_cast(scales + index0 * 2); + scale1[quad] = *reinterpret_cast(scales + index1 * 2); + } +#pragma unroll + for (int quad = 0; quad < kQuadsInFlight; ++quad) { + const int group = group_base + quad * 4 + lane_group; + float weights0[8]; + float weights1[8]; + Q4SimtDecodeAtom::decode_eight(packed0[quad], scale0[quad], weights0); + Q4SimtDecodeAtom::decode_eight(packed1[quad], scale1[quad], weights1); + const uint4 input = load_vec(x + group * Codec::kGroupK + lane_in_group * 8); + const float2 x0 = bf16x2_bits_to_float2(input.x); + const float2 x1 = bf16x2_bits_to_float2(input.y); + const float2 x2 = bf16x2_bits_to_float2(input.z); + const float2 x3 = bf16x2_bits_to_float2(input.w); + const float values[8] = {x0.x, x0.y, x1.x, x1.y, x2.x, x2.y, x3.x, x3.y}; #pragma unroll - for (int item = 0; item < 8; ++item) { - acc0 = fmaf(weights0[item], values[item], acc0); - acc1 = fmaf(weights1[item], values[item], acc1); + for (int item = 0; item < 8; ++item) { + acc0 = fmaf(weights0[item], values[item], acc0); + acc1 = fmaf(weights1[item], values[item], acc1); + } } } } else if constexpr (Codec::kD3SingleValuePerLane) { @@ -246,9 +278,9 @@ __global__ void sparse_moe_d3_nine_warp_kernel( pdl::wait_for_dependencies(); const int expert = ids[warp]; const int row_base = expert * 1024; - dot_two_rows(routed_codes, routed_high, routed_scales, row_base + j, - row_base + kIntermediate + j, x_shared, 0, kHidden, gate, - up); + dot_two_rows(routed_codes, routed_high, routed_scales, + row_base + j, row_base + kIntermediate + j, x_shared, + 0, kHidden, gate, up); } else { dot_two_rows(shared_codes, nullptr, shared_scales, j, kIntermediate + j, x_shared, 0, kHidden, gate, up);