diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md new file mode 100644 index 0000000000000..348131a3db9b3 --- /dev/null +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -0,0 +1,57 @@ +# CUDA HyperConnectionLowRankMix decode tuning + +The CUDA kernel has two implementations. The operator schema, output layout, and +accumulation semantics are identical for both. + +* **Fused path** (two kernels, `HcDownKernel` + `HcUpPreKernel`): 16-bit types + (FP16/BF16), compute capability 8.0+, `rows <= 8`, `num_branches` in + {1, 2, 4, 8}, `hidden % 8 == 0`, `hidden <= 4096`, `rank % 8 == 0`, + `8 <= rank <= 512`, and 16-byte aligned tensors. The row group handled by one + CTA depends on the row count (down rows x up rows): + + | rows | row groups | + | ---- | ---------- | + | 1 | 1 x 1 | + | 2..7 | 2 x 2 (this includes four-row inputs) | + | 8 | 4 x 2 | + + The 4 x 2 configuration is the speculative-decode shape for 8 rows + (`batch * speculative tokens`). The fused path has no dedicated four-row + configuration, and the row grouping is not configurable. +* **Unfused path** (any other shape, any type): branch-wise RMS norm, cuBLAS GEMMs + with FP32 accumulation, and elementwise kernels. It supports an arbitrary + `num_branches`, any row count, and unaligned tensors. 16-bit inputs with + `hidden % 8 == 0`, `hidden <= 4096` and 16-byte aligned tensors use vectorized + norm and pre-mix kernels. The inject GEMV and post-gate are fused in one + deterministic kernel: a four-row kernel for 16-bit `num_branches == 4` with an + 8-byte aligned `inject_weight`, otherwise a one-row kernel tiled over groups of + eight branches (so `num_branches > 8` is supported). + +## Environment variables + +Read once per process, so set it before starting it. + +* `ORT_HYPER_CONNECTION_PDL` (default `1`): controls Hopper programmatic dependent + launch (`cudaLaunchAttributeProgrammaticStreamSerialization`) between `HcDownKernel` + and `HcUpPreKernel` of the fused path. It applies on `SM90+` with CUDA 12.3+, when the + weight inputs (`norm_scale`, `down_weight`, `up_weight`, `inject_weight`) are + constant initializers; `0` forces plain launches. + +The rank-320 decode test (`HyperConnectionLowRankMixTest.FusedDecodeShapeRank320`) +covers 1, 2, 3, 4, 7 and 8 rows. + +## Benchmark + +rc3 model, INT4, H200, DFlash2 S=8 decode, ms per cycle at 2K / 8K / 32K context +(5 repetitions, noise about +-0.03 ms): + +| configuration | 2K | 8K | 32K | +| ---------------------------------- | ----- | ----- | ----- | +| default (4 x 2 rows, PDL on) | 14.93 | 15.07 | 15.30 | +| 2 x 2 rows for 8 rows | 15.00 | 15.11 | 15.38 | +| `ORT_HYPER_CONNECTION_PDL=0` | 15.09 | 15.25 | 15.48 | + +The PDL comparison pass measured 14.92 / 15.10 / 15.33 ms with PDL against +15.09 / 15.25 / 15.48 ms without, so PDL gives about 1%. The 2 x 2 variant for 8 rows +was 0.04-0.08 ms slower, so the row grouping is fixed at 4 x 2 and is no longer +configurable. diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc index ab2d89e01f391..e8c5871a09e7d 100644 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix.cc @@ -163,11 +163,9 @@ Status HyperConnectionLowRankMix::ComputeInternal(OpKernelContext* context) c auto hn = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); auto down = GetScratchBuffer(rows * static_cast(rank), GetComputeStream(context)); auto up = GetScratchBuffer(rows * static_cast(features), GetComputeStream(context)); - auto inject = GetScratchBuffer(rows * static_cast(branches), GetComputeStream(context)); return LaunchHyperConnectionLowRankMixUnfused( GetCublasHandle(context), Stream(context), config, tensors, reinterpret_cast(hn.get()), - reinterpret_cast(down.get()), reinterpret_cast(up.get()), - reinterpret_cast(inject.get())); + reinterpret_cast(down.get()), reinterpret_cast(up.get())); } template class HyperConnectionLowRankMix; diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu index a801360eafbf2..59c9626b0035b 100644 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu @@ -33,15 +33,20 @@ #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cuda_common.h" +// Include env_var_utils.h after cuda_common.h: the latter transitively pulls in provider_api.h, +// which defines SHARED_PROVIDER and suppresses env_var_utils.h's own logging.h include. +#include "core/platform/env_var_utils.h" + namespace onnxruntime::contrib::cuda { namespace { constexpr int kThreads = 256; -constexpr int kColTile = 16; // down columns per CTA of the down kernel (two 16-byte weight vectors) -constexpr int kUpTile = 16; // features per CTA of the up kernel -constexpr int kMaxSliceRows = 1280; // K rows handled by one CTA of the down kernel -constexpr int kWeightSlots = 10; // 16-byte weight vectors per thread held in shared memory -constexpr int kMaxVecIters = 2; // 16-byte vectors per thread in the norm pass +constexpr int kColTile = 16; // down columns per CTA of the down kernel (two 16-byte weight vectors) +constexpr int kUpTile = 16; // features per CTA of the up kernel +constexpr int kMaxSliceRows = 1280; // K rows handled by one CTA of the down kernel +constexpr int kWeightSlots = 10; // 16-byte weight vectors per thread in the down pass +constexpr int kAsyncWeightSlots = 9; // slots staged in shared memory so 5 CTAs/SM fit in 228 KB on sm_90+ +constexpr int kMaxVecIters = 2; // 16-byte vectors per thread in the norm pass constexpr int kMaxFusedHidden = kMaxVecIters * kThreads * 8; constexpr int kMaxFusedRank = 512; constexpr int kMaxFusedBranches = 8; @@ -163,15 +168,24 @@ struct UpParams { float gate_mul; }; +__device__ __forceinline__ bool RowOk(int64_t row0, int r, int64_t total_rows) { + return row0 + r < total_rows; +} + +__device__ __forceinline__ int64_t RowIdx(int64_t row0, int r, int64_t total_rows) { + return (row0 + r < total_rows) ? (row0 + r) : (total_rows - 1); +} + // CTA (j, s, z): K slice s (rows [k0, k0 + len) of the flattened stream, inside branch c) of the ROWS rows // z * ROWS ... z * ROWS + ROWS - 1, which share every weight read. // j < ceil(R / 16) owns down columns [16 j, 16 j + 16); the last j owns the C inject columns. template -__global__ void __launch_bounds__(kThreads) +__global__ void __launch_bounds__(kThreads, ROWS >= 4 ? 4 : 5) HcDownKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, const T* __restrict__ prev_gate, const T* __restrict__ norm_w, const T* __restrict__ w_down, const T* __restrict__ w_inject, T* __restrict__ streams_out, float* __restrict__ partial, float* __restrict__ inv_rms_out, DownParams p) { + constexpr int kLocalAsyncWeightSlots = ROWS >= 4 ? 8 : kAsyncWeightSlots; const int j = blockIdx.x; const int s = blockIdx.y; const int64_t row0 = static_cast(blockIdx.z) * ROWS; @@ -186,17 +200,9 @@ __global__ void __launch_bounds__(kThreads) const int col_tiles = (p.rank + kColTile - 1) / kColTile; const bool is_inject = (j == col_tiles); const int64_t k0 = static_cast(c) * hidden + h0; - int64_t row_idx[ROWS]; - bool row_ok[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - row_ok[r] = row0 + r < p.rows; - row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; // keep loads of unused rows in bounds - } - __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; + __shared__ __align__(16) uint4 wsm[kLocalAsyncWeightSlots * kThreads]; __shared__ __align__(16) uint16_t xs[ROWS][kMaxSliceRows]; - __shared__ float red[8][ROWS][kColTile]; __shared__ float sq_red[8][ROWS]; const int cg = tid & 1; @@ -210,7 +216,7 @@ __global__ void __launch_bounds__(kThreads) if (col < p.rank) { const T* base = w_down + k0 * p.rank + col; #pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { + for (int i = 0; i < kLocalAsyncWeightSlots; ++i) { const int row = rp + 128 * i; if (row < len) CpAsync16(&wsm[i * kThreads + tid], base + static_cast(row) * p.rank); } @@ -230,31 +236,29 @@ __global__ void __launch_bounds__(kThreads) GridWait(); const int nvec = hidden / 8; - uint4 sv[ROWS][kMaxVecIters]; - uint4 bv[ROWS][kMaxVecIters]; - float gate_c[ROWS]; + float sum_sq[ROWS]; #pragma unroll for (int r = 0; r < ROWS; ++r) { - const uint4* sp = reinterpret_cast(streams + row_idx[r] * K + static_cast(c) * hidden); + const int64_t r_idx = RowIdx(row0, r, p.rows); + const bool r_ok = RowOk(row0, r, p.rows); + uint4 sv[kMaxVecIters]; + uint4 bv[kMaxVecIters]; + float gate_c = 0.0f; + const uint4* sp = reinterpret_cast(streams + r_idx * K + static_cast(c) * hidden); #pragma unroll for (int i = 0; i < kMaxVecIters; ++i) { const int qi = tid + kThreads * i; - sv[r][i] = qi < nvec ? sp[qi] : make_uint4(0u, 0u, 0u, 0u); + sv[i] = qi < nvec ? __ldcg(sp + qi) : make_uint4(0u, 0u, 0u, 0u); } if constexpr (kPost) { - gate_c[r] = to_float(prev_gate[row_idx[r] * branches + c]); - const uint4* bp = reinterpret_cast(prev_out + row_idx[r] * hidden); + gate_c = BitsToFloat(__ldcg(reinterpret_cast(prev_gate + r_idx * branches + c))); + const uint4* bp = reinterpret_cast(prev_out + r_idx * hidden); #pragma unroll for (int i = 0; i < kMaxVecIters; ++i) { const int qi = tid + kThreads * i; - bv[r][i] = qi < nvec ? bp[qi] : make_uint4(0u, 0u, 0u, 0u); + bv[i] = qi < nvec ? __ldcg(bp + qi) : make_uint4(0u, 0u, 0u, 0u); } } - } - - float sum_sq[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { sum_sq[r] = 0.0f; #pragma unroll for (int i = 0; i < kMaxVecIters; ++i) { @@ -262,12 +266,12 @@ __global__ void __launch_bounds__(kThreads) if (qi < nvec) { const int h = qi * 8; float v[8]; - Unpack8(sv[r][i], v); + Unpack8(sv[i], v); if constexpr (kPost) { float b[8]; - Unpack8(bv[r][i], b); + Unpack8(bv[i], b); #pragma unroll - for (int m = 0; m < 8; ++m) v[m] = Rnd(fmaf(gate_c[r], b[m], v[m])); + for (int m = 0; m < 8; ++m) v[m] = Rnd(fmaf(gate_c, b[m], v[m])); } #pragma unroll for (int m = 0; m < 8; ++m) sum_sq[r] += v[m] * v[m]; @@ -275,8 +279,8 @@ __global__ void __launch_bounds__(kThreads) const uint4 packed = Pack8(v); *reinterpret_cast(&xs[r][h - h0]) = packed; if constexpr (kPost) { - if (j == 0 && row_ok[r]) { - *reinterpret_cast(streams_out + row_idx[r] * K + static_cast(c) * hidden + h) = + if (j == 0 && r_ok) { + *reinterpret_cast(streams_out + r_idx * K + static_cast(c) * hidden + h) = packed; } } @@ -295,7 +299,9 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int w = 0; w < 8; ++w) total += sq_red[w][r]; inv_rms[r] = rsqrtf(total / static_cast(hidden) + p.epsilon); - if (j == 0 && piece == 0 && tid == 0 && row_ok[r]) inv_rms_out[row_idx[r] * branches + c] = inv_rms[r]; + if (j == 0 && piece == 0 && tid == 0 && RowOk(row0, r, p.rows)) { + inv_rms_out[RowIdx(row0, r, p.rows) * branches + c] = inv_rms[r]; + } if (norm_active) { float x[8], w[8]; Unpack8(*reinterpret_cast(&xs[r][tid * 8]), x); @@ -316,12 +322,16 @@ __global__ void __launch_bounds__(kThreads) for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; } if (col < p.rank) { + const T* base = w_down + k0 * p.rank + col; #pragma unroll for (int i = 0; i < kWeightSlots; ++i) { const int row = rp + 128 * i; if (row < len) { float wf[8]; - Unpack8(wsm[i * kThreads + tid], wf); + const uint4 w_vec = (i < kLocalAsyncWeightSlots) + ? wsm[i * kThreads + tid] + : __ldcg(reinterpret_cast(base + static_cast(row) * p.rank)); + Unpack8(w_vec, wf); #pragma unroll for (int r = 0; r < ROWS; ++r) { const float xv = BitsToFloat(xs[r][row]); @@ -332,6 +342,8 @@ __global__ void __launch_bounds__(kThreads) } } // Lanes of a warp share cg = lane & 1 and differ in rp: reduce over lane bits 1..4. + __syncthreads(); + float (*red)[ROWS][kColTile] = reinterpret_cast(wsm); #pragma unroll for (int r = 0; r < ROWS; ++r) { #pragma unroll @@ -352,7 +364,9 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int q = 0; q < 8; ++q) v += red[q][r][t]; const int out_col = j * kColTile + t; - if (out_col < p.rank && row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + out_col] = v; + if (out_col < p.rank && RowOk(row0, r, p.rows)) { + partial[(RowIdx(row0, r, p.rows) * p.splits + s) * p.partial_stride + out_col] = v; + } } } else { const T* wi = reinterpret_cast(wsm); @@ -377,6 +391,8 @@ __global__ void __launch_bounds__(kThreads) } } } + __syncthreads(); + float (*red)[ROWS][kColTile] = reinterpret_cast(wsm); #pragma unroll for (int r = 0; r < ROWS; ++r) { #pragma unroll @@ -393,21 +409,27 @@ __global__ void __launch_bounds__(kThreads) float v = 0.0f; #pragma unroll for (int q = 0; q < 8; ++q) v += red[q][r][t]; - if (row_ok[r]) partial[(row_idx[r] * p.splits + s) * p.partial_stride + p.rank + t] = v; + if (RowOk(row0, r, p.rows)) { + partial[(RowIdx(row0, r, p.rows) * p.splits + s) * p.partial_stride + p.rank + t] = v; + } } } } // CTA (hb, 0, z) owns features [16 hb, 16 hb + 16) of every branch for ROWS rows. -template -__global__ void __launch_bounds__(kThreads) +template +__global__ void __launch_bounds__(kThreads, 5) HcUpPreKernel(const T* __restrict__ streams, const T* __restrict__ norm_w, const T* __restrict__ w_up, const float* __restrict__ partial, const float* __restrict__ inv_rms_in, T* __restrict__ pre_out, T* __restrict__ gate_out, UpParams p) { + static_assert(ROWS <= 2, "the cross-warp gate reduction below assumes ROWS * 16 * branches <= kThreads"); + constexpr int kLocalWeightSlots = kWeightSlots; const int hb = blockIdx.x; const int64_t row0 = static_cast(blockIdx.z) * ROWS; const int tid = threadIdx.x; - const int branches = p.branches; + const int branches = BRANCHES == 0 ? p.branches : BRANCHES; + const int rank = RANK == 0 ? p.rank : RANK; + constexpr int kBranchStorage = BRANCHES == 0 ? kMaxFusedBranches : BRANCHES; const int hidden = p.hidden; const int64_t K = static_cast(branches) * hidden; const int h0 = hb * kUpTile; @@ -419,26 +441,19 @@ __global__ void __launch_bounds__(kThreads) const int rp = tid / groups; const int col = h0 + half_id * 8; const bool valid = col < hidden; - const int chunk_rows = phases * kWeightSlots; - int64_t row_idx[ROWS]; - bool row_ok[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - row_ok[r] = row0 + r < p.rows; - row_idx[r] = row_ok[r] ? row0 + r : p.rows - 1; - } + const int chunk_rows = phases * kLocalWeightSlots; - __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; - __shared__ __align__(16) float sf[ROWS][kMaxFusedRank]; - __shared__ float sg[ROWS][kMaxFusedBranches * kUpTile]; + __shared__ __align__(16) uint4 wsm[kLocalWeightSlots * kThreads]; + __shared__ __align__(16) float sf[ROWS][RANK == 0 ? kMaxFusedRank : RANK]; + __shared__ float sg[ROWS][kBranchStorage * kUpTile]; // Prefetch the first chunk of up weights before waiting for the down kernel. const T* wbase = w_up + static_cast(c) * hidden + col; if (valid) { #pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { + for (int i = 0; i < kLocalWeightSlots; ++i) { const int row = rp + phases * i; - if (row < p.rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); + if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); } } CpCommit(); @@ -447,73 +462,46 @@ __global__ void __launch_bounds__(kThreads) GridTrigger(); GridWait(); - float xv[ROWS][kMaxFusedBranches], inv[ROWS][kMaxFusedBranches], nw[kMaxFusedBranches]; - const int hl = tid; // feature owned by threads < 16 in the final pre-mix - const bool owner = tid < kUpTile && (h0 + hl) < hidden; - if (owner) { -#pragma unroll - for (int b = 0; b < kMaxFusedBranches; ++b) { - if (b < branches) { - nw[b] = to_float(norm_w[static_cast(b) * hidden + h0 + hl]); -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - xv[r][b] = to_float(streams[row_idx[r] * K + static_cast(b) * hidden + h0 + hl]); - inv[r][b] = inv_rms_in[row_idx[r] * branches + b]; - } - } - } - } - // Fixed-order reduction of the K-split partial sums, then ScaledSiLU; four columns per thread. - const int nq = p.rank / 4; - const int cq = tid * 4; - float4 va[ROWS]; -#pragma unroll - for (int r = 0; r < ROWS; ++r) va[r] = make_float4(0.f, 0.f, 0.f, 0.f); - if (tid < nq) { + const int nq = rank / 4; + for (int idx = tid; idx < ROWS * nq; idx += kThreads) { + const int r = idx / nq; + const int cq = (idx % nq) * 4; + const int64_t r_idx = RowIdx(row0, r, p.rows); + float4 va = make_float4(0.f, 0.f, 0.f, 0.f); for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { - float4 pq[8][ROWS]; #pragma unroll for (int u = 0; u < 8; ++u) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - pq[u][r] = sp0 + u < p.splits - ? __ldcg(reinterpret_cast( - partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)) - : make_float4(0.f, 0.f, 0.f, 0.f); - } - } -#pragma unroll - for (int u = 0; u < 8; ++u) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - va[r].x += pq[u][r].x; - va[r].y += pq[u][r].y; - va[r].z += pq[u][r].z; - va[r].w += pq[u][r].w; + if (sp0 + u < p.splits) { + const float4 q = __ldcg(reinterpret_cast( + partial + (r_idx * p.splits + sp0 + u) * p.partial_stride + cq)); + va.x += q.x; + va.y += q.y; + va.z += q.z; + va.w += q.w; } } } + const float v[4] = {va.x, va.y, va.z, va.w}; + float out[4]; #pragma unroll - for (int r = 0; r < ROWS; ++r) { - const float v[4] = {va[r].x, va[r].y, va[r].z, va[r].w}; - float out[4]; -#pragma unroll - for (int m = 0; m < 4; ++m) { - const float z = Rnd(Rnd(v[m]) * p.silu_alpha); - out[m] = Rnd(z * Rnd(SigmoidF32(z))); - } - *reinterpret_cast(&sf[r][cq]) = make_float4(out[0], out[1], out[2], out[3]); + for (int m = 0; m < 4; ++m) { + const float z = Rnd(Rnd(v[m]) * p.silu_alpha); + out[m] = Rnd(z * Rnd(SigmoidF32(z))); } + *reinterpret_cast(&sf[r][cq]) = make_float4(out[0], out[1], out[2], out[3]); } if (hb == 0 && tid < branches * ROWS) { const int r = tid / branches; const int t = tid % branches; + const int64_t r_idx = RowIdx(row0, r, p.rows); float v = 0.0f; - for (int sp = 0; sp < p.splits; ++sp) v += partial[(row_idx[r] * p.splits + sp) * p.partial_stride + p.rank + t]; + for (int sp = 0; sp < p.splits; ++sp) { + v += __ldcg(partial + (r_idx * p.splits + sp) * p.partial_stride + rank + t); + } const float inject = Rnd(v); const float u = Rnd(SigmoidF32(Rnd(inject / p.gate_div))); - if (row_ok[r]) gate_out[row_idx[r] * branches + t] = from_float(u * p.gate_mul); + if (RowOk(row0, r, p.rows)) gate_out[r_idx * branches + t] = from_float(u * p.gate_mul); } __syncthreads(); @@ -523,21 +511,21 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int m = 0; m < 8; ++m) acc[r][m] = 0.0f; } - for (int chunk0 = 0; chunk0 < p.rank; chunk0 += chunk_rows) { + for (int chunk0 = 0; chunk0 < rank; chunk0 += chunk_rows) { if (chunk0 > 0 && valid) { #pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { + for (int i = 0; i < kLocalWeightSlots; ++i) { const int row = chunk0 + rp + phases * i; - if (row < p.rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); + if (row < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); } CpCommit(); } CpWaitAll(); if (valid) { #pragma unroll - for (int i = 0; i < kWeightSlots; ++i) { + for (int i = 0; i < kLocalWeightSlots; ++i) { const int row = chunk0 + rp + phases * i; - if (row < p.rank) { + if (row < rank) { float wf[8]; Unpack8(wsm[i * kThreads + tid], wf); #pragma unroll @@ -568,36 +556,107 @@ __global__ void __launch_bounds__(kThreads) } } __syncthreads(); - if (tid < groups * 8) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - float v = 0.0f; + // The up kernel runs with at most 2 rows, so ROWS * groups * 8 = 16 * ROWS * branches <= kThreads (branches <= 8). + const int g8 = groups * 8; + if (tid < ROWS * g8) { + const int r = tid / g8; + const int idx = tid % g8; + float v = 0.0f; #pragma unroll - for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * (groups * 8) + tid]; - const int gg = tid >> 3; - const int local = (gg & 1) * 8 + (tid & 7); - sg[r][(gg >> 1) * kUpTile + local] = Rnd(SigmoidF32(Rnd(v))); - } + for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * g8 + idx]; + const int gg = idx >> 3; + const int local = (gg & 1) * 8 + (idx & 7); + sg[r][(gg >> 1) * kUpTile + local] = Rnd(SigmoidF32(Rnd(v))); } __syncthreads(); - if (owner) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { + if (tid < ROWS * kUpTile) { + const int r = tid / kUpTile; + const int hl = tid % kUpTile; + const int64_t r_idx = RowIdx(row0, r, p.rows); + if ((h0 + hl) < hidden && RowOk(row0, r, p.rows)) { float sum = 0.0f; #pragma unroll - for (int b = 0; b < kMaxFusedBranches; ++b) { + for (int b = 0; b < kBranchStorage; ++b) { if (b < branches) { - const float n = Rnd(xv[r][b] * inv[r][b] * nw[b]); + const float nw = to_float(norm_w[static_cast(b) * hidden + h0 + hl]); + const float xv = BitsToFloat(__ldcg(reinterpret_cast( + streams + r_idx * K + static_cast(b) * hidden + h0 + hl))); + const float inv = __ldcg(inv_rms_in + r_idx * branches + b); + const float n = Rnd(xv * inv * nw); sum = fmaf(n, sg[r][b * kUpTile + hl], sum); } } - if (row_ok[r]) pre_out[row_idx[r] * hidden + h0 + hl] = from_float(sum * p.reduction_scale); + pre_out[r_idx * hidden + h0 + hl] = from_float(sum * p.reduction_scale); } } } // ----------------------------------------------------------------------------- unfused path kernels +template +__global__ void NormVec8Kernel(const T* __restrict__ streams, const T* __restrict__ prev_out, + const T* __restrict__ prev_gate, const T* __restrict__ norm_w, T* __restrict__ hn, + T* __restrict__ streams_out, int branches, int hidden, float epsilon) { + __shared__ float sq_red[8]; + const int tid = threadIdx.x; + const int64_t group = blockIdx.x; + const int64_t row = group / branches; + const int c = static_cast(group - row * branches); + const int64_t base = group * hidden; + const int nvec = hidden >> 3; + float gate_c = 0.0f; + if constexpr (kPost) gate_c = to_float(__ldg(prev_gate + group)); + + const uint4* sp = reinterpret_cast(streams + base); + const uint4* bp = kPost ? reinterpret_cast(prev_out + row * hidden) : nullptr; + uint4* sout_p = (kPost && streams_out != nullptr) ? reinterpret_cast(streams_out + base) : nullptr; + + float v_reg[kMaxVecIters][8]; + float sum_sq = 0.0f; +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + if (qi < nvec) { + Unpack8(__ldg(sp + qi), v_reg[i]); + if constexpr (kPost) { + float b[8]; + Unpack8(__ldg(bp + qi), b); +#pragma unroll + for (int m = 0; m < 8; ++m) v_reg[i][m] = Rnd(fmaf(gate_c, b[m], v_reg[i][m])); + if (sout_p != nullptr) { + sout_p[qi] = Pack8(v_reg[i]); + } + } +#pragma unroll + for (int m = 0; m < 8; ++m) sum_sq += v_reg[i][m] * v_reg[i][m]; + } + } +#pragma unroll + for (int off = 16; off > 0; off >>= 1) sum_sq += __shfl_xor_sync(0xffffffffu, sum_sq, off); + if ((tid & 31) == 0) sq_red[tid >> 5] = sum_sq; + __syncthreads(); + + float total = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) total += sq_red[w]; + const float inv_rms = rsqrtf(total / static_cast(hidden) + epsilon); + + const uint4* nwp = reinterpret_cast(norm_w + static_cast(c) * hidden); + uint4* hnp = reinterpret_cast(hn + base); +#pragma unroll + for (int i = 0; i < kMaxVecIters; ++i) { + const int qi = tid + kThreads * i; + if (qi < nvec) { + float nw[8]; + Unpack8(__ldg(nwp + qi), nw); + float out[8]; +#pragma unroll + for (int m = 0; m < 8; ++m) out[m] = v_reg[i][m] * inv_rms * nw[m]; + hnp[qi] = Pack8(out); + } + } +} + template __global__ void NormKernel(const T* __restrict__ streams, const T* __restrict__ prev_out, const T* __restrict__ prev_gate, const T* __restrict__ norm_w, T* __restrict__ hn, @@ -642,6 +701,35 @@ __global__ void ScaledSiluKernel(T* __restrict__ x, int64_t count, float alpha) x[i] = from_float(z * sig); } +template +__global__ void PreMixVec8Kernel(const T* __restrict__ hn, const T* __restrict__ up, T* __restrict__ pre, + uint32_t vec_count, uint32_t vec_cols, int branches, int hidden, + float reduction_scale) { + const uint32_t vec_i = blockIdx.x * blockDim.x + threadIdx.x; + if (vec_i >= vec_count) return; + const uint32_t row = vec_i / vec_cols; + const int h = static_cast((vec_i - row * vec_cols) * 8u); + const int64_t K = static_cast(branches) * hidden; + float sum[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + for (int b = 0; b < branches; ++b) { + const int64_t idx = static_cast(row) * K + static_cast(b) * hidden + h; + float hn_v[8], up_v[8]; + Unpack8(__ldg(reinterpret_cast(hn + idx)), hn_v); + Unpack8(__ldg(reinterpret_cast(up + idx)), up_v); +#pragma unroll + for (int m = 0; m < 8; ++m) { + const float gate = Rnd(SigmoidF32(up_v[m])); + sum[m] = fmaf(hn_v[m], gate, sum[m]); + } + } + float out[8]; +#pragma unroll + for (int m = 0; m < 8; ++m) { + out[m] = sum[m] * reduction_scale; + } + *reinterpret_cast(pre + static_cast(vec_i) * 8) = Pack8(out); +} + template __global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, T* __restrict__ pre, int64_t count, int branches, int hidden, float reduction_scale) { @@ -659,14 +747,97 @@ __global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, pre[i] = from_float(sum * reduction_scale); } +template +__global__ void FusedInjectPostGateB4Kernel(const T* __restrict__ hn, const T* __restrict__ w_inject, + T* __restrict__ gate_out, int64_t rows, int64_t K, + float gate_div, float gate_mul) { + const int64_t row0 = static_cast(blockIdx.x) * ROWS_PER_CTA; + if (row0 >= rows) return; + const int tid = threadIdx.x; + __shared__ float smem[8][ROWS_PER_CTA][4]; + float acc[ROWS_PER_CTA][4]; +#pragma unroll + for (int r = 0; r < ROWS_PER_CTA; ++r) { +#pragma unroll + for (int b = 0; b < 4; ++b) acc[r][b] = 0.0f; + } + const int valid_rows = static_cast(min(static_cast(ROWS_PER_CTA), rows - row0)); + for (int64_t k = tid; k < K; k += kThreads) { + const uint2 wk_u2 = __ldg(reinterpret_cast(w_inject + k * 4)); + float wf[4]; + Unpack2(wk_u2.x, wf[0], wf[1]); + Unpack2(wk_u2.y, wf[2], wf[3]); +#pragma unroll + for (int r = 0; r < ROWS_PER_CTA; ++r) { + if (r < valid_rows) { + const float xv = to_float(hn[(row0 + r) * K + k]); +#pragma unroll + for (int b = 0; b < 4; ++b) { + acc[r][b] = fmaf(xv, wf[b], acc[r][b]); + } + } + } + } +#pragma unroll + for (int r = 0; r < ROWS_PER_CTA; ++r) { +#pragma unroll + for (int b = 0; b < 4; ++b) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) acc[r][b] += __shfl_down_sync(0xffffffffu, acc[r][b], off); + if ((tid & 31) == 0) smem[tid >> 5][r][b] = acc[r][b]; + } + } + __syncthreads(); + if (tid < valid_rows * 4) { + const int r = tid >> 2; + const int b = tid & 3; + float v = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) v += smem[w][r][b]; + const float t = Rnd(Rnd(v) / gate_div); + const float u = Rnd(SigmoidF32(t)); + gate_out[(row0 + r) * 4 + b] = from_float(u * gate_mul); + } +} + +// Grid (rows, ceil(branches / 8)): blockIdx.y selects a tile of up to 8 branch columns, so any num_branches works. template -__global__ void PostGateKernel(const T* __restrict__ inject, T* __restrict__ gate, int64_t count, float gate_div, - float gate_mul) { - const int64_t i = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - if (i >= count) return; - const float t = Rnd(to_float(inject[i]) / gate_div); - const float u = Rnd(SigmoidF32(t)); - gate[i] = from_float(u * gate_mul); +__global__ void FusedInjectPostGateKernel(const T* __restrict__ hn, const T* __restrict__ w_inject, + T* __restrict__ gate_out, int64_t rows, int branches, int64_t K, + float gate_div, float gate_mul) { + const int64_t row = blockIdx.x; + if (row >= rows) return; + const int tid = threadIdx.x; + const int b0 = blockIdx.y * 8; + const int tile = min(8, branches - b0); + const T* x_row = hn + row * K; + __shared__ float smem[8][8]; + float acc[8] = {0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f, 0.f}; + for (int64_t k = tid; k < K; k += kThreads) { + const float xv = to_float(x_row[k]); + const T* wk = w_inject + k * branches + b0; +#pragma unroll + for (int b = 0; b < 8; ++b) { + if (b < tile) acc[b] = fmaf(xv, to_float(wk[b]), acc[b]); + } + } +#pragma unroll + for (int b = 0; b < 8; ++b) { + if (b < tile) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) acc[b] += __shfl_down_sync(0xffffffffu, acc[b], off); + if ((tid & 31) == 0) smem[tid >> 5][b] = acc[b]; + } + } + __syncthreads(); + if (tid < tile) { + float v = 0.0f; +#pragma unroll + for (int w = 0; w < 8; ++w) v += smem[w][tid]; + const float t = Rnd(Rnd(v) / gate_div); + const float u = Rnd(SigmoidF32(t)); + gate_out[row * branches + b0 + tid] = from_float(u * gate_mul); + } } template @@ -782,7 +953,7 @@ Status LaunchWithProgrammaticSerialization(bool use_pdl, Kernel kernel, dim3 gri // griddepcontrol is compiled only for sm_90+; a kernel image JIT-compiled from older PTX would not contain it, // so programmatic dependent launch must not be requested for such an image (or for a driver older than 12.3). -template +template bool IsProgrammaticLaunchAvailable() { #if defined(CUDART_VERSION) && CUDART_VERSION >= 12030 static thread_local int cached_device = -1; @@ -803,9 +974,10 @@ bool IsProgrammaticLaunchAvailable() { return cudaFuncGetAttributes(&attributes, kernel) == cudaSuccess && attributes.ptxVersion >= 90 && attributes.binaryVersion >= 90; }; - available = available && image_has_griddepcontrol(HcDownKernel) && - image_has_griddepcontrol(HcDownKernel) && - image_has_griddepcontrol(HcUpPreKernel); + available = available && image_has_griddepcontrol(HcDownKernel) && + image_has_griddepcontrol(HcDownKernel) && + image_has_griddepcontrol(HcUpPreKernel) && + image_has_griddepcontrol(HcUpPreKernel); (void)cudaGetLastError(); cached_device = device; cached_available = available; @@ -816,28 +988,53 @@ bool IsProgrammaticLaunchAvailable() { #endif } -template +template Status LaunchFusedRows(bool use_pdl, cudaStream_t stream, const HyperConnectionLowRankMixConfig& config, const HyperConnectionLowRankMixTensors& tensors, const FusedGeometry& g, float* partial, float* inv_rms, const DownParams& down, const UpParams& up) { - use_pdl = use_pdl && IsProgrammaticLaunchAvailable(); - const unsigned row_groups = static_cast((config.rows + ROWS - 1) / ROWS); - const dim3 down_grid(g.col_tiles + 1, g.splits, row_groups); - const dim3 up_grid((config.hidden + kUpTile - 1) / kUpTile, 1, row_groups); + static const bool carveout_set = []() { + (void)cudaFuncSetAttribute(HcDownKernel, + cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxShared); + (void)cudaFuncSetAttribute(HcDownKernel, + cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxShared); + (void)cudaFuncSetAttribute(HcUpPreKernel, + cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxShared); + (void)cudaFuncSetAttribute(HcUpPreKernel, + cudaFuncAttributePreferredSharedMemoryCarveout, + cudaSharedmemCarveoutMaxShared); + (void)cudaGetLastError(); + return true; + }(); + (void)carveout_set; + use_pdl = use_pdl && IsProgrammaticLaunchAvailable(); + const unsigned down_row_groups = static_cast((config.rows + DOWN_ROWS - 1) / DOWN_ROWS); + const unsigned up_row_groups = static_cast((config.rows + UP_ROWS - 1) / UP_ROWS); + const dim3 down_grid(g.col_tiles + 1, g.splits, down_row_groups); + const dim3 up_grid((config.hidden + kUpTile - 1) / kUpTile, 1, up_row_groups); if (tensors.prev_block_output != nullptr) { ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( - use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, + use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, tensors.down_weight, tensors.inject_weight, tensors.streams_out, partial, inv_rms, down)); } else { ORT_RETURN_IF_ERROR(LaunchWithProgrammaticSerialization( - use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, + use_pdl, HcDownKernel, down_grid, dim3(kThreads), stream, tensors.streams, static_cast(nullptr), static_cast(nullptr), tensors.norm_scale, tensors.down_weight, tensors.inject_weight, static_cast(nullptr), partial, inv_rms, down)); } const T* up_streams = tensors.prev_block_output != nullptr ? tensors.streams_out : tensors.streams; - return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, + if (config.branches == 4 && config.rank == 320) { + return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), + stream, up_streams, tensors.norm_scale, tensors.up_weight, + static_cast(partial), + static_cast(inv_rms), tensors.pre_output, + tensors.post_gate, up); + } + return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), stream, up_streams, tensors.norm_scale, tensors.up_weight, static_cast(partial), static_cast(inv_rms), tensors.pre_output, @@ -893,11 +1090,17 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c g.splits, g.partial_stride, config.epsilon}; UpParams up{config.rows, config.branches, config.hidden, config.rank, g.splits, g.partial_stride, config.silu_alpha, config.reduction_scale, config.gate_div, config.gate_mul}; - const bool use_pdl = device_prop.major >= 9 && weights_are_constant; + static const bool enable_pdl = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HYPER_CONNECTION_PDL", true); + const bool use_pdl = enable_pdl && device_prop.major >= 9 && weights_are_constant; if (config.rows == 1) { return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); } - return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + // Fused row-group dispatch: rows == 1 -> 1x1, rows 2..7 -> 2x2, rows == 8 -> 4x2 (down x up rows per CTA). + if (config.rows >= 8) { + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + } + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); } } @@ -905,20 +1108,52 @@ template Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, const HyperConnectionLowRankMixConfig& config, const HyperConnectionLowRankMixTensors& tensors, - T* hn, T* down, T* up, T* inject) { + T* hn, T* down, T* up) { const int64_t rows = config.rows; const int64_t K = static_cast(config.branches) * config.hidden; const int64_t groups = rows * config.branches; ORT_RETURN_IF_NOT(groups <= std::numeric_limits::max(), "CUDA launch requires too many blocks"); + const bool use_vec8_norm = + sizeof(T) == 2 && config.hidden % 8 == 0 && config.hidden <= kMaxFusedHidden && + (reinterpret_cast(tensors.streams) & 15u) == 0 && + (reinterpret_cast(tensors.norm_scale) & 15u) == 0 && + (reinterpret_cast(hn) & 15u) == 0 && + (tensors.prev_block_output == nullptr || + ((reinterpret_cast(tensors.prev_block_output) & 15u) == 0 && + (tensors.streams_out == nullptr || (reinterpret_cast(tensors.streams_out) & 15u) == 0))); if (tensors.prev_block_output != nullptr) { - NormKernel<<(groups), kThreads, 0, stream>>>( - tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, - tensors.streams_out, config.branches, config.hidden, config.epsilon); + if constexpr (sizeof(T) == 2) { + if (use_vec8_norm) { + NormVec8Kernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, + tensors.streams_out, config.branches, config.hidden, config.epsilon); + } else { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, + tensors.streams_out, config.branches, config.hidden, config.epsilon); + } + } else { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, tensors.prev_block_output, tensors.prev_post_gate, tensors.norm_scale, hn, + tensors.streams_out, config.branches, config.hidden, config.epsilon); + } } else { - NormKernel<<(groups), kThreads, 0, stream>>>( - tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, - config.epsilon); + if constexpr (sizeof(T) == 2) { + if (use_vec8_norm) { + NormVec8Kernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, + config.epsilon); + } else { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, + config.epsilon); + } + } else { + NormKernel<<(groups), kThreads, 0, stream>>>( + tensors.streams, nullptr, nullptr, tensors.norm_scale, hn, nullptr, config.branches, config.hidden, + config.epsilon); + } } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); @@ -926,7 +1161,6 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ ORT_RETURN_IF_ERROR(math_mode.Set(cublas, CUBLAS_MATH_DISALLOW_REDUCED_PRECISION_REDUCTION)); ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.rank, K, hn, tensors.down_weight, down)); - ORT_RETURN_IF_ERROR(Gemm(cublas, rows, config.branches, K, hn, tensors.inject_weight, inject)); const int64_t down_count = rows * config.rank; ORT_RETURN_IF_NOT(BlocksFor(down_count) <= std::numeric_limits::max(), @@ -937,13 +1171,39 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ ORT_RETURN_IF_ERROR(Gemm(cublas, rows, K, config.rank, down, tensors.up_weight, up)); const int64_t pre_count = rows * config.hidden; - ORT_RETURN_IF_NOT(BlocksFor(pre_count) <= std::numeric_limits::max(), - "CUDA launch requires too many blocks"); - PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( - hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); + if constexpr (sizeof(T) == 2) { + const int64_t vec_count = pre_count / 8; + if (config.hidden % 8 == 0 && vec_count <= std::numeric_limits::max() && + (reinterpret_cast(hn) & 15u) == 0 && + (reinterpret_cast(up) & 15u) == 0 && + (reinterpret_cast(tensors.pre_output) & 15u) == 0) { + PreMixVec8Kernel<<(BlocksFor(vec_count)), kThreads, 0, stream>>>( + hn, up, tensors.pre_output, static_cast(vec_count), + static_cast(config.hidden / 8), config.branches, config.hidden, config.reduction_scale); + } else { + ORT_RETURN_IF_NOT(BlocksFor(pre_count) <= std::numeric_limits::max(), + "CUDA launch requires too many blocks"); + PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( + hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); + } + } else { + ORT_RETURN_IF_NOT(BlocksFor(pre_count) <= std::numeric_limits::max(), + "CUDA launch requires too many blocks"); + PreMixKernel<<(BlocksFor(pre_count)), kThreads, 0, stream>>>( + hn, up, tensors.pre_output, pre_count, config.branches, config.hidden, config.reduction_scale); + } ORT_RETURN_IF_ERROR(CUDA_CALL(cudaGetLastError())); - PostGateKernel<<(BlocksFor(groups)), kThreads, 0, stream>>>( - inject, tensors.post_gate, groups, config.gate_div, config.gate_mul); + if constexpr (sizeof(T) == 2) { + if (config.branches == 4 && (reinterpret_cast(tensors.inject_weight) & 7u) == 0) { + const int64_t row_blocks = (rows + 3) / 4; + FusedInjectPostGateB4Kernel<<(row_blocks), kThreads, 0, stream>>>( + hn, tensors.inject_weight, tensors.post_gate, rows, K, config.gate_div, config.gate_mul); + return CUDA_CALL(cudaGetLastError()); + } + } + const dim3 inject_grid(static_cast(rows), static_cast((config.branches + 7) / 8)); + FusedInjectPostGateKernel<<>>( + hn, tensors.inject_weight, tensors.post_gate, rows, config.branches, K, config.gate_div, config.gate_mul); return CUDA_CALL(cudaGetLastError()); } @@ -955,7 +1215,7 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ const HyperConnectionLowRankMixTensors&, float*, bool); \ template Status LaunchHyperConnectionLowRankMixUnfused( \ cublasHandle_t, cudaStream_t, const HyperConnectionLowRankMixConfig&, \ - const HyperConnectionLowRankMixTensors&, T*, T*, T*, T*); + const HyperConnectionLowRankMixTensors&, T*, T*, T*); INSTANTIATE(float) INSTANTIATE(half) diff --git a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h index bb8d867723672..861ea51299df6 100644 --- a/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h +++ b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.h @@ -58,11 +58,11 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c float* scratch, bool weights_are_constant = false); // General path built from cuBLAS GEMMs and elementwise kernels. Scratch buffers (in elements of T): -// hn [rows, C * H], down [rows, R], up [rows, C * H], inject [rows, C]. +// hn [rows, C * H], down [rows, R], up [rows, C * H]. template Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_t stream, const HyperConnectionLowRankMixConfig& config, const HyperConnectionLowRankMixTensors& tensors, - T* hn, T* down, T* up, T* inject); + T* hn, T* down, T* up); } // namespace onnxruntime::contrib::cuda diff --git a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc index 6ed8209b9bc38..644918d930b4f 100644 --- a/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc +++ b/onnxruntime/test/contrib_ops/hyper_connection_low_rank_mix_op_test.cc @@ -219,6 +219,17 @@ TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { RunTest(1, 4, 2560, 320, false, true, kFp16Tol); } +TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { + // rows == 8 is the default speculative-decode shape (4-row down kernel + 2-row up kernel); 1 uses the single-row + // kernels and 2..7 the 2-row kernels (3 and 7 exercise the partial last row group). + for (int rows : {1, 2, 3, 4, 7, 8}) { + for (bool prologue : {false, true}) { + RunTest(rows, 4, 2560, 320, prologue, true, kBf16Tol, 37); + RunTest(rows, 4, 2560, 320, prologue, true, kFp16Tol, 53); + } + } +} + TEST(HyperConnectionLowRankMixTest, UnfusedManyRows) { RunTest(9, 4, 16, 8, true, true, kFp16Tol); RunTest(9, 4, 16, 8, true, false, kFp16Tol); @@ -269,6 +280,17 @@ TEST(HyperConnectionLowRankMixTest, UnfusedUnsupportedShape) { RunTest(1, 4, 12, 16, false, false, kBf16Tol); } +TEST(HyperConnectionLowRankMixTest, UnfusedMoreThanEightBranches) { + // The inject/post-gate kernel tiles the branch columns in groups of eight: C = 9 needs a second (partial) tile and + // C = 17 a third. + for (bool prologue : {false, true}) { + RunTest(1, 9, 8, 8, prologue, true, kFp32Tol); + RunTest(1, 9, 8, 8, prologue, true, kFp16Tol); + RunTest(1, 9, 8, 8, prologue, true, kBf16Tol); + RunTest(3, 17, 16, 8, prologue, false, kFp16Tol); + } +} + TEST(HyperConnectionLowRankMixTest, Float32) { RunTest(1, 4, 64, 16, true, true, kFp32Tol); RunTest(5, 2, 24, 8, false, true, kFp32Tol);