From 1428ea7addb8165a0f4c156501b9d02943a63bd8 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Fri, 23 Oct 2026 23:01:40 +0000 Subject: [PATCH 01/14] [CUDA] Tune HyperConnectionLowRankMix 4-row decode, rank-320 specialization, and deterministic prefill post-gate Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/hyper_connection_low_rank_mix.md | 19 +++ .../hyper_connection_low_rank_mix_impl.cu | 141 +++++++++++++----- .../hyper_connection_low_rank_mix_op_test.cc | 9 ++ 3 files changed, 128 insertions(+), 41 deletions(-) create mode 100644 docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md 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..e6053a5bf8a6e --- /dev/null +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -0,0 +1,19 @@ +# CUDA HyperConnectionLowRankMix decode tuning + +`ORT_HC_DECODE_ROWS` selects the existing fused-kernel row grouping for eligible +four-row decode inputs. The default is 1 (one row per work group), which +optimizes four-row speculative verification on Hopper (`SM90+`). Set it to 2 +before starting the process to use two-row work groups. + +Only values 1 and 2 are supported. Other input shapes and unsupported fused +configurations retain their existing dispatch. The operator schema, output +layout, and accumulation semantics are unchanged. + +`ORT_HYPER_CONNECTION_PDL` controls Hopper programmatic dependent launch +(`cudaLaunchAttributeProgrammaticStreamSerialization`) between `HcDownKernel` +and `HcUpPreKernel` on `SM90+` when weights are constant initializers. It +defaults to `1` (`true`). + +The rank-320 fused-decode reference test includes four-row inputs. Run that test +in separate processes with `ORT_HC_DECODE_ROWS=1` and `ORT_HC_DECODE_ROWS=2` to +exercise both dispatch choices. 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..24bb70476c14b 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 @@ -29,6 +29,7 @@ #include #include +#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/cub.cuh" #include "core/providers/cuda/cu_inc/cuda_type_helper.cuh" #include "core/providers/cuda/cuda_common.h" @@ -239,15 +240,15 @@ __global__ void __launch_bounds__(kThreads) #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[r][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]); + gate_c[r] = BitsToFloat(__ldcg(reinterpret_cast(prev_gate + row_idx[r] * branches + c))); const uint4* bp = reinterpret_cast(prev_out + row_idx[r] * 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[r][i] = qi < nvec ? __ldcg(bp + qi) : make_uint4(0u, 0u, 0u, 0u); } } } @@ -399,7 +400,7 @@ __global__ void __launch_bounds__(kThreads) } // CTA (hb, 0, z) owns features [16 hb, 16 hb + 16) of every branch for ROWS rows. -template +template __global__ void __launch_bounds__(kThreads) 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, @@ -407,7 +408,9 @@ __global__ void __launch_bounds__(kThreads) 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; @@ -429,8 +432,8 @@ __global__ void __launch_bounds__(kThreads) } __shared__ __align__(16) uint4 wsm[kWeightSlots * kThreads]; - __shared__ __align__(16) float sf[ROWS][kMaxFusedRank]; - __shared__ float sg[ROWS][kMaxFusedBranches * kUpTile]; + __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; @@ -438,7 +441,7 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int i = 0; i < kWeightSlots; ++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,25 +450,26 @@ __global__ void __launch_bounds__(kThreads) GridTrigger(); GridWait(); - float xv[ROWS][kMaxFusedBranches], inv[ROWS][kMaxFusedBranches], nw[kMaxFusedBranches]; + float xv[ROWS][kBranchStorage], inv[ROWS][kBranchStorage], nw[kBranchStorage]; 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) { + for (int b = 0; b < kBranchStorage; ++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]; + xv[r][b] = BitsToFloat(__ldcg(reinterpret_cast( + streams + row_idx[r] * K + static_cast(b) * hidden + h0 + hl))); + inv[r][b] = __ldcg(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 nq = rank / 4; const int cq = tid * 4; float4 va[ROWS]; #pragma unroll @@ -510,7 +514,9 @@ __global__ void __launch_bounds__(kThreads) const int r = tid / branches; const int t = tid % branches; 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 + (row_idx[r] * 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); @@ -523,12 +529,12 @@ __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) { 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(); } @@ -537,7 +543,7 @@ __global__ void __launch_bounds__(kThreads) #pragma unroll for (int i = 0; i < kWeightSlots; ++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 @@ -585,7 +591,7 @@ __global__ void __launch_bounds__(kThreads) for (int r = 0; r < ROWS; ++r) { 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]); sum = fmaf(n, sg[r][b * kUpTile + hl], sum); @@ -660,13 +666,40 @@ __global__ void PreMixKernel(const T* __restrict__ hn, const T* __restrict__ up, } 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 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; +#pragma unroll + for (int b = 0; b < 8; ++b) { + if (b < branches) acc[b] = fmaf(xv, to_float(wk[b]), acc[b]); + } + } +#pragma unroll + for (int b = 0; b < 8; ++b) { + if (b < branches) { +#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 < branches) { + 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 + tid] = from_float(u * gate_mul); + } } template @@ -782,7 +815,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 +836,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 +850,36 @@ 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); + 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,10 +935,27 @@ 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); } + static const int decode_rows = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_DECODE_ROWS", 1); + static const int up_decode_rows = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_UP_DECODE_ROWS", 1); + ORT_ENFORCE(decode_rows == 1 || decode_rows == 2, + "ORT_HC_DECODE_ROWS must be 1 or 2"); + if (config.rows == 4 && decode_rows == 1) { + if (up_decode_rows == 4) { + return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); + } + if (up_decode_rows == 2) { + 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); + } return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); } } @@ -926,7 +985,7 @@ 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)); + (void)inject; const int64_t down_count = rows * config.rank; ORT_RETURN_IF_NOT(BlocksFor(down_count) <= std::numeric_limits::max(), @@ -942,8 +1001,8 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ 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); + FusedInjectPostGateKernel<<(rows), kThreads, 0, stream>>>( + hn, tensors.inject_weight, tensors.post_gate, rows, config.branches, K, config.gate_div, config.gate_mul); return CUDA_CALL(cudaGetLastError()); } 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..f18b1becd5a5c 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,15 @@ TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { RunTest(1, 4, 2560, 320, false, true, kFp16Tol); } +TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { + for (int rows : {1, 2, 4}) { + 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); From d0d03f6d266c9501490ccb709dfe1710f2c8067b Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 00:37:59 +0000 Subject: [PATCH 02/14] [CUDA] Remove unused ORT_HC_UP_DECODE_ROWS instantiations in HyperConnectionLowRankMix Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hyper_connection_low_rank_mix_impl.cu | 39 +++++++------------ 1 file changed, 15 insertions(+), 24 deletions(-) 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 24bb70476c14b..29c3bbc032a9c 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 @@ -815,7 +815,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; @@ -836,10 +836,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) && - 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; @@ -850,36 +850,35 @@ 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 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); + 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); 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; if (config.branches == 4 && config.rank == 320) { - return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), + 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, + 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, @@ -943,18 +942,10 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c } static const int decode_rows = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_DECODE_ROWS", 1); - static const int up_decode_rows = - onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_UP_DECODE_ROWS", 1); ORT_ENFORCE(decode_rows == 1 || decode_rows == 2, "ORT_HC_DECODE_ROWS must be 1 or 2"); if (config.rows == 4 && decode_rows == 1) { - if (up_decode_rows == 4) { - return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); - } - if (up_decode_rows == 2) { - 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); + 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); } From aacae29ee017364dfe1ab5c827f625ed43e6ef98 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 04:10:07 +0000 Subject: [PATCH 03/14] fix(cuda): include env_var_utils.h after cuda_common.h in hyper_connection_low_rank_mix_impl.cu Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/bert/hyper_connection_low_rank_mix_impl.cu | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) 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 29c3bbc032a9c..e6b9ce4e841d4 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 @@ -29,11 +29,14 @@ #include #include -#include "core/platform/env_var_utils.h" #include "core/providers/cuda/cu_inc/cub.cuh" #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 { From 1904f3e064b40b80c857bd58cd8cb96e77bf67b5 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 15:17:42 +0000 Subject: [PATCH 04/14] perf(cuda): eliminate pq[8][ROWS] local register array in HcUpPreKernel Accumulate K-split partial sums directly into va[r] inside HcUpPreKernel instead of staging through a float4 pq[8][ROWS] local array, avoiding dynamic register indexing and saving ~0.35 ms/cycle across 96 launches at S=8. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hyper_connection_low_rank_mix_impl.cu | 25 +++++++------------ 1 file changed, 9 insertions(+), 16 deletions(-) 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 e6b9ce4e841d4..4fb41e4018f17 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 @@ -479,25 +479,18 @@ __global__ void __launch_bounds__(kThreads) for (int r = 0; r < ROWS; ++r) va[r] = make_float4(0.f, 0.f, 0.f, 0.f); if (tid < nq) { for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { - float4 pq[8][ROWS]; #pragma unroll for (int u = 0; u < 8; ++u) { + if (sp0 + u < p.splits) { #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; + for (int r = 0; r < ROWS; ++r) { + const float4 q = __ldcg(reinterpret_cast( + partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)); + va[r].x += q.x; + va[r].y += q.y; + va[r].z += q.z; + va[r].w += q.w; + } } } } From c40587556ba4ab69142c41ac3b1f1831f1251130 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 21:06:27 +0000 Subject: [PATCH 05/14] Optimize HyperConnectionLowRankMix HcUpPreKernel and HcDownKernel for 5 blocks/SM - Parallelize HcUpPreKernel partial reduction across ROWS*nq (160) threads and move xv/inv/nw global loads to the epilogue (60->48 regs, 4->5 blocks/SM, -22.5% kernel time). - Alias red over wsm, stage kAsyncWeightSlots=9 in HcDownKernel shared memory (42.05 KB/block), and set __launch_bounds__(256, 5) + MaxShared carveout so both HcDownKernel and HcUpPreKernel run at 5 blocks/SM on Hopper. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hyper_connection_low_rank_mix_impl.cu | 175 +++++++++++------- 1 file changed, 112 insertions(+), 63 deletions(-) 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 4fb41e4018f17..ce6652f4947e6 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 @@ -44,7 +44,8 @@ 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 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; @@ -171,7 +172,7 @@ struct UpParams { // 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 <= 2 ? 5 : 4) 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, @@ -198,9 +199,8 @@ __global__ void __launch_bounds__(kThreads) 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[kAsyncWeightSlots * 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; @@ -214,7 +214,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 < kAsyncWeightSlots; ++i) { const int row = rp + 128 * i; if (row < len) CpAsync16(&wsm[i * kThreads + tid], base + static_cast(row) * p.rank); } @@ -234,31 +234,27 @@ __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) { + uint4 sv[kMaxVecIters]; + uint4 bv[kMaxVecIters]; + float gate_c = 0.0f; const uint4* sp = reinterpret_cast(streams + row_idx[r] * 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 ? __ldcg(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] = BitsToFloat(__ldcg(reinterpret_cast(prev_gate + row_idx[r] * branches + c))); + gate_c = BitsToFloat(__ldcg(reinterpret_cast(prev_gate + row_idx[r] * branches + c))); const uint4* bp = reinterpret_cast(prev_out + row_idx[r] * hidden); #pragma unroll for (int i = 0; i < kMaxVecIters; ++i) { const int qi = tid + kThreads * i; - bv[r][i] = qi < nvec ? __ldcg(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) { @@ -266,12 +262,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]; @@ -320,12 +316,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 < kAsyncWeightSlots) + ? 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]); @@ -336,6 +336,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 @@ -381,6 +383,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 @@ -453,50 +457,27 @@ __global__ void __launch_bounds__(kThreads) GridTrigger(); GridWait(); - float xv[ROWS][kBranchStorage], inv[ROWS][kBranchStorage], nw[kBranchStorage]; - 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 < kBranchStorage; ++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] = BitsToFloat(__ldcg(reinterpret_cast( - streams + row_idx[r] * K + static_cast(b) * hidden + h0 + hl))); - inv[r][b] = __ldcg(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 = rank / 4; - const int cq = tid * 4; - float4 va[ROWS]; + if (ROWS * nq <= kThreads) { + if (tid < ROWS * nq) { + const int r = tid / nq; + const int cq = (tid % nq) * 4; + float4 va = make_float4(0.f, 0.f, 0.f, 0.f); + for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { #pragma unroll - for (int r = 0; r < ROWS; ++r) va[r] = make_float4(0.f, 0.f, 0.f, 0.f); - if (tid < nq) { - for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { -#pragma unroll - for (int u = 0; u < 8; ++u) { - if (sp0 + u < p.splits) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { + for (int u = 0; u < 8; ++u) { + if (sp0 + u < p.splits) { const float4 q = __ldcg(reinterpret_cast( partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)); - va[r].x += q.x; - va[r].y += q.y; - va[r].z += q.z; - va[r].w += q.w; + va.x += q.x; + va.y += q.y; + va.z += q.z; + va.w += q.w; } } } - } -#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}; + const float v[4] = {va.x, va.y, va.z, va.w}; float out[4]; #pragma unroll for (int m = 0; m < 4; ++m) { @@ -505,6 +486,40 @@ __global__ void __launch_bounds__(kThreads) } *reinterpret_cast(&sf[r][cq]) = make_float4(out[0], out[1], out[2], out[3]); } + } else { + 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) { + for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { +#pragma unroll + for (int u = 0; u < 8; ++u) { + if (sp0 + u < p.splits) { +#pragma unroll + for (int r = 0; r < ROWS; ++r) { + const float4 q = __ldcg(reinterpret_cast( + partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)); + va[r].x += q.x; + va[r].y += q.y; + va[r].z += q.z; + va[r].w += q.w; + } + } + } + } +#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]); + } + } } if (hb == 0 && tid < branches * ROWS) { const int r = tid / branches; @@ -570,30 +585,47 @@ __global__ void __launch_bounds__(kThreads) } } __syncthreads(); - if (tid < groups * 8) { + const int g8 = groups * 8; + if (ROWS * g8 <= kThreads) { + 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) * 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))); + } + } else if (tid < g8) { #pragma unroll for (int r = 0; r < ROWS; ++r) { float v = 0.0f; #pragma unroll - for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * (groups * 8) + tid]; + for (int q = 0; q < 8; ++q) v += red[(q * ROWS + r) * g8 + tid]; const int gg = tid >> 3; const int local = (gg & 1) * 8 + (tid & 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; + if ((h0 + hl) < hidden && row_ok[r]) { float sum = 0.0f; #pragma unroll 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 + row_idx[r] * K + static_cast(b) * hidden + h0 + hl))); + const float inv = __ldcg(inv_rms_in + row_idx[r] * 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[row_idx[r] * hidden + h0 + hl] = from_float(sum * p.reduction_scale); } } } @@ -850,6 +882,23 @@ 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) { + 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 row_groups = static_cast((config.rows + ROWS - 1) / ROWS); const dim3 down_grid(g.col_tiles + 1, g.splits, row_groups); From 2811979b3d38d4064880a9a6e1b51ad3ad596859 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Sat, 24 Oct 2026 23:40:48 +0000 Subject: [PATCH 06/14] lintrunner --- .../cuda/bert/hyper_connection_low_rank_mix_impl.cu | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) 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 ce6652f4947e6..e3c5b9369b4f1 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 @@ -41,12 +41,12 @@ 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 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 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; From 96485a2b833a8840aa069116ab2535f674e13dab Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 05:34:08 +0000 Subject: [PATCH 07/14] perf(cuda): eliminate HcDown/HcUpPre stack frame, default S=8 to 4x2 rows, and vectorize prefill - Replace runtime-indexed row_idx/row_ok local arrays with inline RowIdx/RowOk helpers, eliminating the 32-80B local-memory stack frame (STACK:0) across all HcDownKernel and HcUpPreKernel instantiations. - Default S=8 decode to LaunchFusedRows (DOWN_ROWS=4, UP_ROWS=2), halving HcDownKernel grid from 672 CTAs (1.02 waves) to 336 CTAs (0.64 waves) and reducing 97-pair wall-clock span by 307 us/cycle on H200. - Add 16-byte vectorized PreMixVec8Kernel (1.68x faster) and 4-row/CTA FusedInjectPostGateB4Kernel (2.39x faster) for unfused prefill (-5.34 ms at L=2048). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hyper_connection_low_rank_mix_impl.cu | 320 +++++++++++------- 1 file changed, 202 insertions(+), 118 deletions(-) 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 e3c5b9369b4f1..2956bf4921e32 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 @@ -168,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, ROWS <= 2 ? 5 : 4) +__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; @@ -191,15 +200,8 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) 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[kAsyncWeightSlots * kThreads]; + __shared__ __align__(16) uint4 wsm[kLocalAsyncWeightSlots * kThreads]; __shared__ __align__(16) uint16_t xs[ROWS][kMaxSliceRows]; __shared__ float sq_red[8][ROWS]; @@ -214,7 +216,7 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) if (col < p.rank) { const T* base = w_down + k0 * p.rank + col; #pragma unroll - for (int i = 0; i < kAsyncWeightSlots; ++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); } @@ -237,18 +239,20 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) float sum_sq[ROWS]; #pragma unroll for (int r = 0; r < ROWS; ++r) { + 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 + row_idx[r] * K + static_cast(c) * hidden); + 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[i] = qi < nvec ? __ldcg(sp + qi) : make_uint4(0u, 0u, 0u, 0u); } if constexpr (kPost) { - gate_c = BitsToFloat(__ldcg(reinterpret_cast(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; @@ -275,8 +279,8 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) 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, ROWS <= 2 ? 5 : 4) #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); @@ -322,7 +328,7 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) const int row = rp + 128 * i; if (row < len) { float wf[8]; - const uint4 w_vec = (i < kAsyncWeightSlots) + const uint4 w_vec = (i < kLocalAsyncWeightSlots) ? wsm[i * kThreads + tid] : __ldcg(reinterpret_cast(base + static_cast(row) * p.rank)); Unpack8(w_vec, wf); @@ -358,7 +364,9 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) #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); @@ -401,17 +409,20 @@ __global__ void __launch_bounds__(kThreads, ROWS <= 2 ? 5 : 4) 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) +__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) { + constexpr int kLocalWeightSlots = ROWS >= 4 ? 8 : kWeightSlots; const int hb = blockIdx.x; const int64_t row0 = static_cast(blockIdx.z) * ROWS; const int tid = threadIdx.x; @@ -429,16 +440,9 @@ __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) uint4 wsm[kLocalWeightSlots * kThreads]; __shared__ __align__(16) float sf[ROWS][RANK == 0 ? kMaxFusedRank : RANK]; __shared__ float sg[ROWS][kBranchStorage * kUpTile]; @@ -446,7 +450,7 @@ __global__ void __launch_bounds__(kThreads) 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 < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); } @@ -459,78 +463,44 @@ __global__ void __launch_bounds__(kThreads) // Fixed-order reduction of the K-split partial sums, then ScaledSiLU; four columns per thread. const int nq = rank / 4; - if (ROWS * nq <= kThreads) { - if (tid < ROWS * nq) { - const int r = tid / nq; - const int cq = (tid % nq) * 4; - float4 va = make_float4(0.f, 0.f, 0.f, 0.f); - for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { + 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) { #pragma unroll - for (int u = 0; u < 8; ++u) { - if (sp0 + u < p.splits) { - const float4 q = __ldcg(reinterpret_cast( - partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)); - va.x += q.x; - va.y += q.y; - va.z += q.z; - va.w += q.w; - } + for (int u = 0; u < 8; ++u) { + 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 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]); } - } else { - 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) { - for (int sp0 = 0; sp0 < p.splits; sp0 += 8) { + const float v[4] = {va.x, va.y, va.z, va.w}; + float out[4]; #pragma unroll - for (int u = 0; u < 8; ++u) { - if (sp0 + u < p.splits) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - const float4 q = __ldcg(reinterpret_cast( - partial + (row_idx[r] * p.splits + sp0 + u) * p.partial_stride + cq)); - va[r].x += q.x; - va[r].y += q.y; - va[r].z += q.z; - va[r].w += q.w; - } - } - } - } -#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 += __ldcg(partial + (row_idx[r] * p.splits + sp) * p.partial_stride + rank + t); + 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(); @@ -543,7 +513,7 @@ __global__ void __launch_bounds__(kThreads) 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 < rank) CpAsync16(&wsm[i * kThreads + tid], wbase + static_cast(row) * K); } @@ -552,7 +522,7 @@ __global__ void __launch_bounds__(kThreads) 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 < rank) { float wf[8]; @@ -612,20 +582,21 @@ __global__ void __launch_bounds__(kThreads) if (tid < ROWS * kUpTile) { const int r = tid / kUpTile; const int hl = tid % kUpTile; - if ((h0 + hl) < hidden && row_ok[r]) { + 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 < kBranchStorage; ++b) { if (b < branches) { const float nw = to_float(norm_w[static_cast(b) * hidden + h0 + hl]); const float xv = BitsToFloat(__ldcg(reinterpret_cast( - streams + row_idx[r] * K + static_cast(b) * hidden + h0 + hl))); - const float inv = __ldcg(inv_rms_in + row_idx[r] * branches + b); + 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); } } - 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); } } } @@ -676,6 +647,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) { @@ -693,6 +693,59 @@ __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); + } +} + template __global__ void FusedInjectPostGateKernel(const T* __restrict__ hn, const T* __restrict__ w_inject, T* __restrict__ gate_out, int64_t rows, int branches, int64_t K, @@ -843,7 +896,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; @@ -864,10 +917,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) && - 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; @@ -878,52 +931,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) { static const bool carveout_set = []() { - (void)cudaFuncSetAttribute(HcDownKernel, + (void)cudaFuncSetAttribute(HcDownKernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); - (void)cudaFuncSetAttribute(HcDownKernel, + (void)cudaFuncSetAttribute(HcDownKernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); - (void)cudaFuncSetAttribute(HcUpPreKernel, + (void)cudaFuncSetAttribute(HcUpPreKernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); - (void)cudaFuncSetAttribute(HcUpPreKernel, + (void)cudaFuncSetAttribute(HcUpPreKernel, cudaFuncAttributePreferredSharedMemoryCarveout, cudaSharedmemCarveoutMaxShared); (void)cudaGetLastError(); return true; }(); (void)carveout_set; - 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); + 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; if (config.branches == 4 && config.rank == 320) { - return LaunchWithProgrammaticSerialization(use_pdl, HcUpPreKernel, up_grid, dim3(kThreads), + 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, + 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, @@ -987,12 +1041,17 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c } static const int decode_rows = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_DECODE_ROWS", 1); - ORT_ENFORCE(decode_rows == 1 || decode_rows == 2, - "ORT_HC_DECODE_ROWS must be 1 or 2"); + ORT_ENFORCE(decode_rows == 1 || decode_rows == 2 || decode_rows == 4, + "ORT_HC_DECODE_ROWS must be 1, 2, or 4"); if (config.rows == 4 && decode_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); + static const int s8_decode_rows = + onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_S8_DECODE_ROWS", 4); + if (config.rows >= 8 && (s8_decode_rows == 4 || decode_rows == 4)) { + 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); } } @@ -1032,11 +1091,36 @@ 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())); + 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()); + } + } FusedInjectPostGateKernel<<(rows), kThreads, 0, stream>>>( hn, tensors.inject_weight, tensors.post_gate, rows, config.branches, K, config.gate_div, config.gate_mul); return CUDA_CALL(cudaGetLastError()); From 71978625139e3b84f2a8cceb32cc1ca26ca5938a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 06:47:10 +0000 Subject: [PATCH 08/14] [CUDA] Add vectorized NormVec8Kernel with register caching for unfused HyperConnectionLowRankMix Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hyper_connection_low_rank_mix_impl.cu | 108 +++++++++++++++++- 1 file changed, 102 insertions(+), 6 deletions(-) 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 2956bf4921e32..8a6561311297b 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 @@ -603,6 +603,70 @@ __global__ void __launch_bounds__(kThreads, 5) // ----------------------------------------------------------------------------- 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, @@ -1065,14 +1129,46 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ 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())); From 09c2ec1bcfd0b92d8d49c74e9ef5cd9efbb90f2a Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:52:02 +0000 Subject: [PATCH 09/14] [CUDA] HyperConnectionLowRankMix: remove the four-row decode special case Drop the rows == 4 -> 1x1 dispatch (and its ORT_HC_DECODE_ROWS default of 1); four-row inputs now use the generic 2x2 fused dispatch. ORT_HC_DECODE_ROWS=4 and ORT_HC_S8_DECODE_ROWS are kept for benchmarking the 8-row configuration. Remove code that is unreachable afterwards: the ROWS >= 4 weight-slot ternary and the second cross-warp reduction branch of HcUpPreKernel (it only runs with at most 2 rows), and the unused `inject` scratch buffer/parameter of the unfused launcher. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../bert/hyper_connection_low_rank_mix.cc | 4 +- .../hyper_connection_low_rank_mix_impl.cu | 43 +++++++------------ .../bert/hyper_connection_low_rank_mix_impl.h | 4 +- 3 files changed, 19 insertions(+), 32 deletions(-) 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 8a6561311297b..e2e56d579f0f2 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 @@ -422,7 +422,8 @@ __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) { - constexpr int kLocalWeightSlots = ROWS >= 4 ? 8 : kWeightSlots; + 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; @@ -555,28 +556,17 @@ __global__ void __launch_bounds__(kThreads, 5) } } __syncthreads(); + // 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 (ROWS * g8 <= kThreads) { - 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) * 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))); - } - } else if (tid < g8) { -#pragma unroll - for (int r = 0; r < ROWS; ++r) { - float v = 0.0f; + 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) * g8 + 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 (tid < ROWS * kUpTile) { @@ -1103,13 +1093,13 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c if (config.rows == 1) { 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). + // ORT_HC_S8_DECODE_ROWS=2 selects 2x2 for the 8-row decode shape; ORT_HC_DECODE_ROWS=4 forces 4x2 for it. + // Values 1 and 2 of ORT_HC_DECODE_ROWS no longer select anything (the dedicated four-row dispatch was removed). static const int decode_rows = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_DECODE_ROWS", 1); ORT_ENFORCE(decode_rows == 1 || decode_rows == 2 || decode_rows == 4, "ORT_HC_DECODE_ROWS must be 1, 2, or 4"); - if (config.rows == 4 && decode_rows == 1) { - return LaunchFusedRows(use_pdl, stream, config, tensors, g, partial, inv_rms, down, up); - } static const int s8_decode_rows = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_S8_DECODE_ROWS", 4); if (config.rows >= 8 && (s8_decode_rows == 4 || decode_rows == 4)) { @@ -1123,7 +1113,7 @@ 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; @@ -1176,7 +1166,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)); - (void)inject; const int64_t down_count = rows * config.rank; ORT_RETURN_IF_NOT(BlocksFor(down_count) <= std::numeric_limits::max(), @@ -1230,7 +1219,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 From 75a500243297b0a3304197244144ebafdc200e7e Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:52:15 +0000 Subject: [PATCH 10/14] [CUDA] HyperConnectionLowRankMix: support more than eight branches on the unfused path FusedInjectPostGateKernel only accumulated eight branch columns into an smem[8][8] buffer, so num_branches > 8 read out of bounds and left the ninth gate unaccumulated. Tile the kernel over branch columns (grid.y) so any num_branches is computed correctly. Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../bert/hyper_connection_low_rank_mix_impl.cu | 16 ++++++++++------ 1 file changed, 10 insertions(+), 6 deletions(-) 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 e2e56d579f0f2..2a77269ebafd1 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 @@ -800,6 +800,7 @@ __global__ void FusedInjectPostGateB4Kernel(const T* __restrict__ hn, const T* _ } } +// Grid (rows, ceil(branches / 8)): blockIdx.y selects a tile of up to 8 branch columns, so any num_branches works. template __global__ void FusedInjectPostGateKernel(const T* __restrict__ hn, const T* __restrict__ w_inject, T* __restrict__ gate_out, int64_t rows, int branches, int64_t K, @@ -807,33 +808,35 @@ __global__ void FusedInjectPostGateKernel(const T* __restrict__ hn, const T* __r 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; + const T* wk = w_inject + k * branches + b0; #pragma unroll for (int b = 0; b < 8; ++b) { - if (b < branches) acc[b] = fmaf(xv, to_float(wk[b]), acc[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 < branches) { + 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 < branches) { + 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 + tid] = from_float(u * gate_mul); + gate_out[row * branches + b0 + tid] = from_float(u * gate_mul); } } @@ -1206,7 +1209,8 @@ Status LaunchHyperConnectionLowRankMixUnfused(cublasHandle_t cublas, cudaStream_ return CUDA_CALL(cudaGetLastError()); } } - FusedInjectPostGateKernel<<(rows), kThreads, 0, stream>>>( + 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()); } From 8a2c2cd43572aff5f9a7bc789b0227f940ed754f Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:52:15 +0000 Subject: [PATCH 11/14] [CUDA] HyperConnectionLowRankMix tests: cover the 8-row decode dispatch and nine branches Add rows == 8 at (branches=4, hidden=2560, rank=320) to the rank-320 decode test (FP16/BF16, with and without prologue) so LaunchFusedRows is exercised, and add a more-than-eight-branch regression test (rows=1, C=9, H=8, R=8 for FP32/FP16/BF16, plus C=17). Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../hyper_connection_low_rank_mix_op_test.cc | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) 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 f18b1becd5a5c..b2a13eee5e126 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 @@ -220,7 +220,9 @@ TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { } TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { - for (int rows : {1, 2, 4}) { + // rows == 8 is the default speculative-decode shape (4-row down kernel + 2-row up kernel); 1, 2 and 4 cover the + // smaller row-group dispatches. + for (int rows : {1, 2, 4, 8}) { for (bool prologue : {false, true}) { RunTest(rows, 4, 2560, 320, prologue, true, kBf16Tol, 37); RunTest(rows, 4, 2560, 320, prologue, true, kFp16Tol, 53); @@ -278,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); From 533fff2a89e2d506d1ec0b0a836e5b7da985f1ef Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 22:52:40 +0000 Subject: [PATCH 12/14] docs: describe the HyperConnectionLowRankMix fused/unfused dispatch and env knobs Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/hyper_connection_low_rank_mix.md | 62 ++++++++++++++----- 1 file changed, 45 insertions(+), 17 deletions(-) diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md index e6053a5bf8a6e..3be1592a14f0e 100644 --- a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -1,19 +1,47 @@ # CUDA HyperConnectionLowRankMix decode tuning -`ORT_HC_DECODE_ROWS` selects the existing fused-kernel row grouping for eligible -four-row decode inputs. The default is 1 (one row per work group), which -optimizes four-row speculative verification on Hopper (`SM90+`). Set it to 2 -before starting the process to use two-row work groups. - -Only values 1 and 2 are supported. Other input shapes and unsupported fused -configurations retain their existing dispatch. The operator schema, output -layout, and accumulation semantics are unchanged. - -`ORT_HYPER_CONNECTION_PDL` controls Hopper programmatic dependent launch -(`cudaLaunchAttributeProgrammaticStreamSerialization`) between `HcDownKernel` -and `HcUpPreKernel` on `SM90+` when weights are constant initializers. It -defaults to `1` (`true`). - -The rank-320 fused-decode reference test includes four-row inputs. Run that test -in separate processes with `ORT_HC_DECODE_ROWS=1` and `ORT_HC_DECODE_ROWS=2` to -exercise both dispatch choices. +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 | selected by | + | ----------- | ---------- | ------------------------------------------------------------ | + | 1 | 1 x 1 | always | + | 2 .. 7 | 2 x 2 | always (this includes four-row inputs) | + | 8 (default) | 4 x 2 | default; `ORT_HC_S8_DECODE_ROWS=4` or `ORT_HC_DECODE_ROWS=4` | + | 8 | 2 x 2 | `ORT_HC_S8_DECODE_ROWS=2` (and `ORT_HC_DECODE_ROWS` != 4) | + + The 4 x 2 configuration is the speculative-decode default for 8 rows + (`batch * speculative tokens`). The fused path has no dedicated four-row + configuration. +* **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` use vectorized norm and pre-mix kernels, and the inject GEMV and + post-gate are fused in one deterministic kernel that is tiled over groups of + eight branches. + +## Environment variables + +Read once per process, so set them before starting it. + +* `ORT_HC_S8_DECODE_ROWS` (default `4`): for fused inputs with `rows >= 8`, `4` + selects the 4 x 2 kernels and any other value selects 2 x 2. +* `ORT_HC_DECODE_ROWS` (default `1`, allowed values `1`, `2`, `4`): value `4` also + selects the 4 x 2 kernels for `rows >= 8`. Values `1` and `2` select nothing; they + are accepted only for compatibility with the removed four-row dispatch. +* `ORT_HYPER_CONNECTION_PDL` (default `1`): controls Hopper programmatic dependent + launch (`cudaLaunchAttributeProgrammaticStreamSerialization`) between `HcDownKernel` + and `HcUpPreKernel`. 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, 4 and 8 rows. Run it in separate processes with +`ORT_HC_S8_DECODE_ROWS=4` (default) and `ORT_HC_S8_DECODE_ROWS=2` to exercise both +eight-row configurations. From 6078d627eff96dca35fea3fb5a11e1e091d7fbe9 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 23:01:15 +0000 Subject: [PATCH 13/14] [CUDA] HyperConnectionLowRankMix: drop the no-op ORT_HC_DECODE_ROWS knob, validate ORT_HC_S8_DECODE_ROWS, tighten docs and decode test Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/hyper_connection_low_rank_mix.md | 33 ++++++++++--------- .../hyper_connection_low_rank_mix_impl.cu | 10 ++---- .../hyper_connection_low_rank_mix_op_test.cc | 6 ++-- 3 files changed, 23 insertions(+), 26 deletions(-) diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md index 3be1592a14f0e..9b7d6bef5df5c 100644 --- a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -9,12 +9,12 @@ accumulation semantics are identical for both. `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 | selected by | - | ----------- | ---------- | ------------------------------------------------------------ | - | 1 | 1 x 1 | always | - | 2 .. 7 | 2 x 2 | always (this includes four-row inputs) | - | 8 (default) | 4 x 2 | default; `ORT_HC_S8_DECODE_ROWS=4` or `ORT_HC_DECODE_ROWS=4` | - | 8 | 2 x 2 | `ORT_HC_S8_DECODE_ROWS=2` (and `ORT_HC_DECODE_ROWS` != 4) | + | rows | row groups | selected by | + | ----------- | ---------- | --------------------------------------------- | + | 1 | 1 x 1 | always | + | 2 .. 7 | 2 x 2 | always (this includes four-row inputs) | + | 8 (default) | 4 x 2 | default (`ORT_HC_S8_DECODE_ROWS=4`) | + | 8 | 2 x 2 | `ORT_HC_S8_DECODE_ROWS=2` | The 4 x 2 configuration is the speculative-decode default for 8 rows (`batch * speculative tokens`). The fused path has no dedicated four-row @@ -22,26 +22,27 @@ accumulation semantics are identical for both. * **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` use vectorized norm and pre-mix kernels, and the inject GEMV and - post-gate are fused in one deterministic kernel that is tiled over groups of - eight branches. + `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 them before starting it. -* `ORT_HC_S8_DECODE_ROWS` (default `4`): for fused inputs with `rows >= 8`, `4` - selects the 4 x 2 kernels and any other value selects 2 x 2. -* `ORT_HC_DECODE_ROWS` (default `1`, allowed values `1`, `2`, `4`): value `4` also - selects the 4 x 2 kernels for `rows >= 8`. Values `1` and `2` select nothing; they - are accepted only for compatibility with the removed four-row dispatch. -* `ORT_HYPER_CONNECTION_PDL` (default `1`): controls Hopper programmatic dependent +* `ORT_HC_S8_DECODE_ROWS` (experimental, default `4`, allowed values `2` and `4`; + other values fail the fused launch): applies only to the fused path with + `rows == 8`. `4` selects the 4 x 2 kernels and `2` selects 2 x 2. It has no effect + for other row counts. +* `ORT_HYPER_CONNECTION_PDL` (experimental, default `1`): controls Hopper programmatic dependent launch (`cudaLaunchAttributeProgrammaticStreamSerialization`) between `HcDownKernel` and `HcUpPreKernel`. 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, 4 and 8 rows. Run it in separate processes with +covers 1, 2, 3, 4, 7 and 8 rows. Run it in separate processes with `ORT_HC_S8_DECODE_ROWS=4` (default) and `ORT_HC_S8_DECODE_ROWS=2` to exercise both eight-row configurations. 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 2a77269ebafd1..a807278d1237f 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 @@ -1097,15 +1097,11 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c 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). - // ORT_HC_S8_DECODE_ROWS=2 selects 2x2 for the 8-row decode shape; ORT_HC_DECODE_ROWS=4 forces 4x2 for it. - // Values 1 and 2 of ORT_HC_DECODE_ROWS no longer select anything (the dedicated four-row dispatch was removed). - static const int decode_rows = - onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_DECODE_ROWS", 1); - ORT_ENFORCE(decode_rows == 1 || decode_rows == 2 || decode_rows == 4, - "ORT_HC_DECODE_ROWS must be 1, 2, or 4"); + // Experimental: ORT_HC_S8_DECODE_ROWS=2 selects 2x2 instead of 4x2 for the 8-row shape (default 4). static const int s8_decode_rows = onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_S8_DECODE_ROWS", 4); - if (config.rows >= 8 && (s8_decode_rows == 4 || decode_rows == 4)) { + ORT_ENFORCE(s8_decode_rows == 2 || s8_decode_rows == 4, "ORT_HC_S8_DECODE_ROWS must be 2 or 4"); + if (config.rows >= 8 && s8_decode_rows == 4) { 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); 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 b2a13eee5e126..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 @@ -220,9 +220,9 @@ TEST(HyperConnectionLowRankMixTest, FusedMultipleKSlicesAndWeightChunks) { } TEST(HyperConnectionLowRankMixTest, FusedDecodeShapeRank320) { - // rows == 8 is the default speculative-decode shape (4-row down kernel + 2-row up kernel); 1, 2 and 4 cover the - // smaller row-group dispatches. - for (int rows : {1, 2, 4, 8}) { + // 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); From 4530b1ff780b82f0e6c032172904d48e9a86f764 Mon Sep 17 00:00:00 2001 From: Tianlei Wu Date: Mon, 26 Oct 2026 23:37:20 +0000 Subject: [PATCH 14/14] [CUDA] HyperConnectionLowRankMix: remove ORT_HC_S8_DECODE_ROWS, always use 4x2 for 8 rows; document PDL benchmark Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> --- .../cuda/hyper_connection_low_rank_mix.md | 51 +++++++++++-------- .../hyper_connection_low_rank_mix_impl.cu | 6 +-- 2 files changed, 31 insertions(+), 26 deletions(-) diff --git a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md index 9b7d6bef5df5c..348131a3db9b3 100644 --- a/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md +++ b/docs/contrib_ops/cuda/hyper_connection_low_rank_mix.md @@ -9,16 +9,15 @@ accumulation semantics are identical for both. `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 | selected by | - | ----------- | ---------- | --------------------------------------------- | - | 1 | 1 x 1 | always | - | 2 .. 7 | 2 x 2 | always (this includes four-row inputs) | - | 8 (default) | 4 x 2 | default (`ORT_HC_S8_DECODE_ROWS=4`) | - | 8 | 2 x 2 | `ORT_HC_S8_DECODE_ROWS=2` | - - The 4 x 2 configuration is the speculative-decode default for 8 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. + 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 @@ -30,19 +29,29 @@ accumulation semantics are identical for both. ## Environment variables -Read once per process, so set them before starting it. +Read once per process, so set it before starting it. -* `ORT_HC_S8_DECODE_ROWS` (experimental, default `4`, allowed values `2` and `4`; - other values fail the fused launch): applies only to the fused path with - `rows == 8`. `4` selects the 4 x 2 kernels and `2` selects 2 x 2. It has no effect - for other row counts. -* `ORT_HYPER_CONNECTION_PDL` (experimental, default `1`): controls Hopper programmatic dependent +* `ORT_HYPER_CONNECTION_PDL` (default `1`): controls Hopper programmatic dependent launch (`cudaLaunchAttributeProgrammaticStreamSerialization`) between `HcDownKernel` - and `HcUpPreKernel`. 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. + 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. Run it in separate processes with -`ORT_HC_S8_DECODE_ROWS=4` (default) and `ORT_HC_S8_DECODE_ROWS=2` to exercise both -eight-row configurations. +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_impl.cu b/onnxruntime/contrib_ops/cuda/bert/hyper_connection_low_rank_mix_impl.cu index a807278d1237f..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 @@ -1097,11 +1097,7 @@ Status LaunchHyperConnectionLowRankMixFused(const cudaDeviceProp& device_prop, c 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). - // Experimental: ORT_HC_S8_DECODE_ROWS=2 selects 2x2 instead of 4x2 for the 8-row shape (default 4). - static const int s8_decode_rows = - onnxruntime::ParseEnvironmentVariableWithDefault("ORT_HC_S8_DECODE_ROWS", 4); - ORT_ENFORCE(s8_decode_rows == 2 || s8_decode_rows == 4, "ORT_HC_S8_DECODE_ROWS must be 2 or 4"); - if (config.rows >= 8 && s8_decode_rows == 4) { + 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);