From 93c3353dfe80e99ab663bc536fd3ddfbd3af8ccb Mon Sep 17 00:00:00 2001 From: giveen Date: Fri, 18 Sep 2026 20:08:39 -0600 Subject: [PATCH 1/4] perf(ops): double-buffer the Q4 k-split MMA staging pipeline Stage K group g+1 codes, scales and activation tiles with cp.async while group g runs on the tensor cores. The stage count and launch bounds derive from the static shared-memory footprint: two stages for 8/16-column tiles, one for wider tiles. Numerics are unchanged. RTX 5090, CUDA 13.3, cold-cache public Op benches (median): Q4 LinearSwiGLU 34816x5120 T=4/8/16 77.8/79.9/106.5 -> 75.2/79.7/102.4 us Q4 Linear 4096x5120 T=2/6/10 20.5/20.5/22.5 -> 15.8/16.4/18.4 us Q4 Linear 7168x5120 T=4/8 24.6/24.6 -> 20.5/22.5 us In-model MTP3 T=4 SwiGLU kernel median 75.4 -> 73.1 us (nsys, qwen3_8_27b). Q4 linear/linear_swiglu/linear_add/linear_topk oracle suites pass. --- src/ops/linear/q4/q4_ksplit_mma.cuh | 115 ++++++++++++++++++++++------ 1 file changed, 91 insertions(+), 24 deletions(-) diff --git a/src/ops/linear/q4/q4_ksplit_mma.cuh b/src/ops/linear/q4/q4_ksplit_mma.cuh index cf6437409a..0037cfa1b9 100644 --- a/src/ops/linear/q4/q4_ksplit_mma.cuh +++ b/src/ops/linear/q4/q4_ksplit_mma.cuh @@ -1,5 +1,18 @@ #pragma once +// Q4G64 RowSplit K-split MMA contraction for small column extents. +// +// Eight K-split warps cooperatively own one 16-row output tile; each warp evaluates a disjoint +// 64-wide K slice of every 512-wide K group, then the CTA reduces FP32 partials in shared +// memory. Codes are decoded to exact BF16 integers, contracted per 64-wide group on the tensor +// cores, and each group's binary16 scale is applied once to the FP32 group accumulator. +// +// Staging is a cp.async ring of complete K groups (codes, scales and each warp's activation +// tile). Q4 stores half a byte per weight, so a single-buffered CTA leaves too few bytes in +// flight to reach DRAM bandwidth; with two stages the loads for group g+1 overlap the MMA work +// on group g. Wide column tiles fall back to one stage because two would not fit the static +// shared-memory limit. + #include "ops/common/mma.cuh" #include "ops/common/memory.cuh" @@ -37,8 +50,37 @@ struct Q4KSplitMmaSchedule { static constexpr int kGroupK = kKWarps * kTileKPerWarp; static constexpr int kRowsPerCta = 16; static constexpr int kRowsPerLoaderWarp = kRowsPerCta / kKWarps; + // Static shared memory available to one CTA for cp.async staging buffers. + static constexpr int kStaticSharedBudget = 48 * 1024; +}; + +// One cp.async staging buffer: the CTA's 16-row Q4 code slab for one 512-wide K group, each +// K-split warp's activation tile, and the per-row group scales. +template +struct alignas(16) Q4KSplitStage { + std::uint8_t codes[Q4KSplitMmaSchedule::kRowsPerCta][Q4KSplitMmaSchedule::kGroupK / 2]; + __nv_bfloat16 activations[Q4KSplitMmaSchedule::kKWarps] + [TileCols * Q4KSplitMmaSchedule::kTileKPerWarp]; + std::uint16_t scales[Q4KSplitMmaSchedule::kRowsPerCta][Q4KSplitMmaSchedule::kKWarps]; }; +// Pipeline depth: two stages when the ring fits the static shared-memory budget, else one. +template +__host__ __device__ constexpr int q4_ksplit_stages() { + return 2 * static_cast(sizeof(Q4KSplitStage)) <= + Q4KSplitMmaSchedule::kStaticSharedBudget + ? 2 + : 1; +} + +// Two-stage residency is shared-memory bound (four 25 KiB CTAs per 100 KiB SM for 8-column +// tiles, two 41 KiB CTAs for 16-column tiles), so the register cap follows that residency +// instead of the single-stage six-CTA target. +template +__host__ __device__ constexpr int q4_ksplit_min_blocks_per_sm() { + return q4_ksplit_stages() == 2 ? 4 : Q4KSplitMmaSchedule::kMinBlocksPerSm; +} + __device__ __forceinline__ int q4_ksplit_swizzle_64(int row, int col) { return (((col >> 3) ^ (row & 7)) << 3) | (col & 7); } @@ -58,7 +100,7 @@ __device__ __forceinline__ unsigned q4_ksplit_bf16_pair(std::uint8_t packed) { template -__launch_bounds__(256, 6) __global__ +__launch_bounds__(Q4KSplitMmaSchedule::kThreads, q4_ksplit_min_blocks_per_sm()) __global__ void q4_ksplit_mma_kernel(const __nv_bfloat16* __restrict__ x, const std::uint8_t* __restrict__ codes, const std::uint8_t* __restrict__ scales, @@ -74,25 +116,23 @@ __launch_bounds__(256, 6) __global__ constexpr int kCodeRowBytes = kHidden / 2; constexpr int kTileCols = TileCols; constexpr int kNt = kTileCols / 8; + constexpr int kStages = q4_ksplit_stages(); static_assert(kTileCols >= 8 && kTileCols <= 32 && (kTileCols % 8) == 0); static_assert(ActiveCols >= 1 && ActiveCols <= kTileCols && ActiveCols > kTileCols - 8); static_assert((kHidden % kGroupK) == 0); static_assert(RowPolicy::kOutputRowsPerCta <= kRowsPerCta); + static_assert(kStages == 1 || kStages == 2); - union SharedStorage { - struct { - std::uint8_t codes[kRowsPerCta][kGroupK / 2]; - __nv_bfloat16 activations[kWarps][kTileCols * kTileK]; - std::uint16_t scales[kRowsPerCta][kWarps]; - } staging; + using Stage = Q4KSplitStage; + static_assert(sizeof(Stage) % 16 == 0, "stage buffers must keep 16-byte cp.async alignment"); + union SharedStorage { + Stage staging[kStages]; float partial[kWarps * kNt * 32 * 4]; }; + static_assert(sizeof(SharedStorage) <= Q4KSplitMmaSchedule::kStaticSharedBudget); __shared__ __align__(16) SharedStorage shared; - auto& code_shared = shared.staging.codes; - auto& x_shared = shared.staging.activations; - auto& scale_shared = shared.staging.scales; const int tid = static_cast(threadIdx.x); const int warp = tid >> 5; @@ -103,7 +143,8 @@ __launch_bounds__(256, 6) __global__ const int row0 = static_cast(blockIdx.x) * RowPolicy::kOutputRowsPerCta; const int live_columns = MaskedColumns ? columns : ActiveCols; - const auto stage_x = [&](int group_k0) { + const auto stage_x = [&](int group_k0, Stage& stage) { + auto& x_shared = stage.activations; constexpr int kItemsPerSplit = ActiveCols * (kTileK / 8); for (int item = lane; item < kItemsPerSplit; item += 32) { const int col = item / (kTileK / 8); @@ -122,7 +163,9 @@ __launch_bounds__(256, 6) __global__ } }; - const auto stage_weight = [&](int group_k0) { + const auto stage_weight = [&](int group_k0, Stage& stage) { + auto& code_shared = stage.codes; + auto& scale_shared = stage.scales; #pragma unroll for (int row_item = 0; row_item < Schedule::kRowsPerLoaderWarp; ++row_item) { const int row = warp * Schedule::kRowsPerLoaderWarp + row_item; @@ -148,16 +191,38 @@ __launch_bounds__(256, 6) __global__ const int warp_koff = k_split * kTileK; float acc[kNt][4] = {}; - stage_weight(0); - stage_x(0); + stage_weight(0, shared.staging[0]); + stage_x(0, shared.staging[0]); cp_commit(); - cp_wait<0>(); - __syncthreads(); #pragma unroll for (int group_index = 0; group_index < kGroups; ++group_index) { - const int group_k0 = group_index * kGroupK; - float group_acc[kNt][4] = {}; + const int group_k0 = group_index * kGroupK; + const bool has_next = group_index + 1 < kGroups; + + // Two stages: issue group g+1 into the other buffer before consuming group g, then + // wait for everything except that newest group. One stage: group g+1 is issued after + // this group's trailing barrier, so wait for all outstanding copies here. + if constexpr (kStages == 2) { + if (has_next) { + Stage& next = shared.staging[(group_index + 1) % kStages]; + stage_weight(group_k0 + kGroupK, next); + stage_x(group_k0 + kGroupK, next); + cp_commit(); + cp_wait<1>(); + } else { + cp_wait<0>(); + } + } else { + cp_wait<0>(); + } + __syncthreads(); + + const Stage& current = shared.staging[group_index % kStages]; + const auto& code_shared = current.codes; + const auto& x_shared = current.activations; + const auto& scale_shared = current.scales; + float group_acc[kNt][4] = {}; #pragma unroll for (int ks = 0; ks < 4; ++ks) { @@ -188,13 +253,15 @@ __launch_bounds__(256, 6) __global__ acc[nt][3] = fmaf(group_acc[nt][3], bot_scale, acc[nt][3]); } - if (group_index + 1 < kGroups) { - __syncthreads(); - stage_weight(group_k0 + kGroupK); - stage_x(group_k0 + kGroupK); - cp_commit(); - cp_wait<0>(); + if (has_next) { + // Every warp has finished reading this buffer before it is refilled: by the + // single-stage issue below, or by the two-stage prefetch one iteration later. __syncthreads(); + if constexpr (kStages == 1) { + stage_weight(group_k0 + kGroupK, shared.staging[0]); + stage_x(group_k0 + kGroupK, shared.staging[0]); + cp_commit(); + } } } From 9fedc856d5b148a2ab039641753b1bca23fa1847 Mon Sep 17 00:00:00 2001 From: giveen Date: Fri, 18 Sep 2026 20:49:41 -0600 Subject: [PATCH 2/4] perf(ops): add Q5 k-split MMA routes for small batches MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replace the superseded small-T Q5 LinearAdd kernels and retune affected Q5 dispatch boundaries around the shared k-split MMA implementation. 🤖 Generated with Codebuff Co-Authored-By: Codebuff --- src/ops/linear/q5/q5_ksplit_mma.cuh | 399 ++++++++++++++++++ src/ops/linear/q5/shapes/n5120_k17408.cu | 9 +- src/ops/linear/q5/shapes/n5120_k6144.cu | 9 +- src/ops/linear/q5/shapes/n6144_k5120.cu | 9 +- src/ops/linear/q5/shapes/n7168_k5120.cu | 9 +- .../linear_add/q5/q5_linear_add_gemm_mma.cu | 16 - .../linear_add/q5/q5_linear_add_gemm_simt.cu | 75 ---- src/ops/linear_add/q5/q5_linear_add_kernels.h | 8 +- .../linear_add/q5/q5_linear_add_ksplit_mma.cu | 51 +++ src/ops/linear_add/q5/q5_linear_add_plan.cpp | 37 +- src/ops/linear_add/q5/q5_linear_add_plan.h | 4 +- src/ops/linear_add/sources.cmake | 2 +- tests/ops/linear/test_q5_a16.cpp | 36 ++ tests/ops/linear_add/test_q5_a16.cpp | 25 +- 14 files changed, 519 insertions(+), 170 deletions(-) create mode 100644 src/ops/linear/q5/q5_ksplit_mma.cuh delete mode 100644 src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu create mode 100644 src/ops/linear_add/q5/q5_linear_add_ksplit_mma.cu diff --git a/src/ops/linear/q5/q5_ksplit_mma.cuh b/src/ops/linear/q5/q5_ksplit_mma.cuh new file mode 100644 index 0000000000..ebb9b8ccc3 --- /dev/null +++ b/src/ops/linear/q5/q5_ksplit_mma.cuh @@ -0,0 +1,399 @@ +#pragma once + +// Q5G64 RowSplit K-split MMA contraction for small column extents: out[N,T] = W[N,K] . x[K,T]. +// +// Eight K-split warps cooperatively own one 16-row output tile; each warp evaluates a disjoint +// 64-wide K slice (exactly one quantization group) of every 512-wide K group, then the CTA +// reduces FP32 partials in shared memory. Q5 codes (nibble plane plus the fifth-bit plane) are +// decoded to exact BF16 integers, contracted per group on the tensor cores, and each group's +// binary16 scale is applied once to the FP32 group accumulator. +// +// Staging is a two-deep cp.async ring of complete K groups (codes, high bits, scales and each +// warp's activation tile) so the loads of group g+1 overlap the MMA work on group g. Compared +// with the warp-per-row SIMT route, the activation tile is read once per 16 rows instead of +// once per row and the per-weight FMA work moves to the tensor cores, which is what the T>1 +// verification extents need to stay on the DRAM roofline. + +#include "core/device.h" +#include "core/tensor.h" +#include "core/weight.h" +#include "ops/common/mma.cuh" +#include "ops/common/memory.cuh" +#include "ops/linear/q5/q5_rowsplit_storage.cuh" + +#include +#include + +#include +#include + +namespace ninfer::ops::detail { + +struct Q5KSplitMmaSchedule { + static constexpr int kKWarps = 8; + static constexpr int kThreads = kKWarps * 32; + static constexpr int kTileKPerWarp = Q5RowSplitStorage::kGroupK; + static constexpr int kGroupK = kKWarps * kTileKPerWarp; + static constexpr int kRowsPerCta = 16; + static constexpr int kRowsPerLoaderWarp = kRowsPerCta / kKWarps; + static constexpr int kCodeBytesPerGroup = Q5RowSplitStorage::kCodeBytesPerGroup; + static constexpr int kHighBytesPerGroup = Q5RowSplitStorage::kHighBytesPerGroup; + static constexpr int kStaticSharedBudget = 48 * 1024; +}; + +// One cp.async stage: the CTA's 16-row code and high-bit slabs for one 512-wide K group, the +// binary16 scale of each row's eight groups, and each K-split warp's activation tile. +template +struct alignas(16) Q5KSplitStage { + std::uint8_t codes[Q5KSplitMmaSchedule::kRowsPerCta] + [Q5KSplitMmaSchedule::kKWarps * Q5KSplitMmaSchedule::kCodeBytesPerGroup]; + std::uint8_t high[Q5KSplitMmaSchedule::kRowsPerCta] + [Q5KSplitMmaSchedule::kKWarps * Q5KSplitMmaSchedule::kHighBytesPerGroup]; + std::uint16_t scales[Q5KSplitMmaSchedule::kRowsPerCta][Q5KSplitMmaSchedule::kKWarps]; + __nv_bfloat16 activations[Q5KSplitMmaSchedule::kKWarps] + [TileCols * Q5KSplitMmaSchedule::kTileKPerWarp]; +}; + +// Ring depth: the deepest ring that fits the static shared-memory budget. The N=5120 shapes +// launch only 320 CTAs (about two per SM), so each group costs one memory latency unless +// several groups are in flight per CTA; extra stages cost nothing there. +template +__host__ __device__ constexpr int q5_ksplit_stages() { + constexpr int kStage = static_cast(sizeof(Q5KSplitStage)); + constexpr int kBudget = Q5KSplitMmaSchedule::kStaticSharedBudget; + return 3 * kStage <= kBudget ? 3 : (2 * kStage <= kBudget ? 2 : 1); +} + +// Residency is shared-memory bound; the register cap follows the CTAs that actually fit. +template +__host__ __device__ constexpr int q5_ksplit_min_blocks_per_sm() { + constexpr int kFit = (100 * 1024) / (q5_ksplit_stages() * + static_cast(sizeof(Q5KSplitStage))); + return kFit < 1 ? 1 : (kFit > 4 ? 4 : kFit); +} + +__device__ __forceinline__ int q5_ksplit_swizzle_64(int row, int col) { + return (((col >> 3) ^ (row & 7)) << 3) | (col & 7); +} + +// Fragment reads touch the same column of rows gid..gid+7, whose 256-byte code rows and +// 64-byte high-bit rows would otherwise map to one bank. XOR-swizzling the 16-byte chunk index +// by a row-derived key spreads the eight rows over eight distinct banks. +__device__ __forceinline__ int q5_ksplit_code_offset(int row, int byte) { + return (((byte >> 4) ^ (row & 7)) << 4) | (byte & 15); +} + +__device__ __forceinline__ int q5_ksplit_high_offset(int row, int byte) { + return (((byte >> 4) ^ ((row >> 1) & 3)) << 4) | (byte & 15); +} + +union Q5KSplitBf16PairBits { + __nv_bfloat162 pair; + unsigned bits; +}; + +// One packed code byte holds K pair (2b, 2b+1) in its low/high nibbles; the matching high-bit +// byte holds their fifth bits at 2*(b&3) and 2*(b&3)+1. Signed code: (u ^ 16) - 16. +__device__ __forceinline__ unsigned q5_ksplit_bf16_pair(std::uint8_t packed, std::uint8_t high, + int shift) { + const int q0 = ((static_cast(packed & 0x0fu) | (((high >> shift) & 1) << 4)) ^ 0x10) - 0x10; + const int q1 = + ((static_cast(packed >> 4) | (((high >> (shift + 1)) & 1) << 4)) ^ 0x10) - 0x10; + Q5KSplitBf16PairBits result; + result.pair = __floats2bfloat162_rn(static_cast(q0), static_cast(q1)); + return result.bits; +} + +// AddResidual selects the LinearAdd form: out already holds the residual and receives +// residual + projection; otherwise out receives the projection. +template +__launch_bounds__(Q5KSplitMmaSchedule::kThreads, q5_ksplit_min_blocks_per_sm()) __global__ + void q5_ksplit_mma_kernel(const __nv_bfloat16* __restrict__ x, + const std::uint8_t* __restrict__ codes, + const std::uint8_t* __restrict__ high, + const std::uint8_t* __restrict__ scales, + __nv_bfloat16* __restrict__ out, int columns) { + using Schedule = Q5KSplitMmaSchedule; + constexpr int kHidden = InputRows; + constexpr int kTileK = Schedule::kTileKPerWarp; + constexpr int kWarps = Schedule::kKWarps; + constexpr int kRowsPerCta = Schedule::kRowsPerCta; + constexpr int kGroupK = Schedule::kGroupK; + constexpr int kGroups = kHidden / kGroupK; + constexpr int kGroupsPerRow = kHidden / Q5RowSplitStorage::kGroupK; + constexpr int kCodeRowBytes = kGroupsPerRow * Schedule::kCodeBytesPerGroup; + constexpr int kHighRowBytes = kGroupsPerRow * Schedule::kHighBytesPerGroup; + constexpr int kScaleRowBytes = kGroupsPerRow * 2; + constexpr int kTileCols = TileCols; + constexpr int kNt = kTileCols / 8; + constexpr int kStages = q5_ksplit_stages(); + constexpr int kPrefetch = kStages - 1; + static_assert(kTileCols >= 8 && kTileCols <= 32 && (kTileCols % 8) == 0); + static_assert(ActiveCols >= 1 && ActiveCols <= kTileCols && ActiveCols > kTileCols - 8); + static_assert((kHidden % kGroupK) == 0 && kGroups >= 1); + static_assert(OutputRows % kRowsPerCta == 0); + // Every row's code, high-bit and scale planes must start 16-byte aligned for cp.async. + static_assert(kCodeRowBytes % 16 == 0 && kHighRowBytes % 16 == 0 && kScaleRowBytes % 16 == 0); + static_assert(kStages >= 1 && kStages <= 4); + + using Stage = Q5KSplitStage; + static_assert(sizeof(Stage) % 16 == 0, "stage buffers must keep 16-byte cp.async alignment"); + + union SharedStorage { + Stage staging[kStages]; + float partial[kWarps * kNt * 32 * 4]; + }; + static_assert(sizeof(SharedStorage) <= Schedule::kStaticSharedBudget); + + __shared__ __align__(16) SharedStorage shared; + + const int tid = static_cast(threadIdx.x); + const int warp = tid >> 5; + const int lane = tid & 31; + const int gid = lane >> 2; + const int lid = lane & 3; + const int k_split = warp; + const int row0 = static_cast(blockIdx.x) * kRowsPerCta; + // blockIdx.y selects a kTileCols-wide column tile; a single-tile launch has offset 0. + const int column_offset = static_cast(blockIdx.y) * kTileCols; + const int live_columns = min(kTileCols, columns - column_offset); + x += static_cast(column_offset) * kHidden; + out += static_cast(column_offset) * OutputRows; + + const auto stage_x = [&](int group_k0, Stage& stage) { + constexpr int kItemsPerSplit = ActiveCols * (kTileK / 8); + for (int item = lane; item < kItemsPerSplit; item += 32) { + const int col = item / (kTileK / 8); + const int k8 = item - col * (kTileK / 8); + auto* dst = &stage.activations[warp][col * kTileK + q5_ksplit_swizzle_64(col, k8 * 8)]; + const int source = col < live_columns ? col : 0; + cp_async_zfill<16>(dst, + &x[static_cast(source) * kHidden + group_k0 + + warp * kTileK + k8 * 8], + col < live_columns ? 16 : 0); + } + }; + + // Per row and 512-wide group: 256 code bytes (16 chunks), 64 high-bit bytes (4 chunks) and + // 16 scale bytes (1 chunk); 21 lanes of the loader warp carry one row. + const auto stage_weight = [&](int group_k0, Stage& stage) { + constexpr int kCodeChunks = kWarps * Schedule::kCodeBytesPerGroup / 16; + constexpr int kHighChunks = kWarps * Schedule::kHighBytesPerGroup / 16; + static_assert(kCodeChunks + kHighChunks + 1 <= 32); +#pragma unroll + for (int row_item = 0; row_item < Schedule::kRowsPerLoaderWarp; ++row_item) { + const int row = warp * Schedule::kRowsPerLoaderWarp + row_item; + const int weight_row = row0 + row; + if (lane < kCodeChunks) { + cp_async<16, Cache::cg>(&stage.codes[row][q5_ksplit_code_offset(row, lane * 16)], + codes + static_cast(weight_row) * kCodeRowBytes + + group_k0 / 2 + lane * 16); + } else if (lane < kCodeChunks + kHighChunks) { + const int chunk = lane - kCodeChunks; + cp_async<16, Cache::cg>(&stage.high[row][q5_ksplit_high_offset(row, chunk * 16)], + high + static_cast(weight_row) * kHighRowBytes + + group_k0 / 8 + chunk * 16); + } else if (lane == kCodeChunks + kHighChunks) { + cp_async<16>(&stage.scales[row][0], + scales + static_cast(weight_row) * kScaleRowBytes + + (group_k0 / Q5RowSplitStorage::kGroupK) * 2); + } + } + }; + + const int b_rin = lane & 7; + const int b_koff = ((lane >> 3) & 1) << 3; + const int code_off = k_split * Schedule::kCodeBytesPerGroup; + const int high_off = k_split * Schedule::kHighBytesPerGroup; + const int hshift = 2 * lid; + float acc[kNt][4] = {}; + + if constexpr (kStages == 1) { + stage_weight(0, shared.staging[0]); + stage_x(0, shared.staging[0]); + cp_commit(); + } else { +#pragma unroll + for (int s = 0; s < kPrefetch; ++s) { + if (s < kGroups) { + stage_weight(s * kGroupK, shared.staging[s]); + stage_x(s * kGroupK, shared.staging[s]); + } + cp_commit(); + } + } + + // Full unrolling of a long K loop (34 groups at K=17408) thrashes the instruction cache; + // a partial unroll that is a multiple of every ring depth keeps buffer indexing cheap. + constexpr int kGroupUnroll = kGroups <= 12 ? kGroups : 6; +#pragma unroll kGroupUnroll + for (int group_index = 0; group_index < kGroups; ++group_index) { + const int group_k0 = group_index * kGroupK; + const bool has_next = group_index + 1 < kGroups; + + // Ring of kStages: keep kPrefetch groups in flight by refilling the buffer consumed in + // the previous iteration (its trailing barrier retired every reader); one group is + // committed every iteration so wait_group means "group g has landed". With a + // single stage, group g+1 is issued after this group's trailing barrier instead. + if constexpr (kStages >= 2) { + const int fetch = group_index + kPrefetch; + if (fetch < kGroups) { + Stage& next = shared.staging[fetch % kStages]; + stage_weight(fetch * kGroupK, next); + stage_x(fetch * kGroupK, next); + } + cp_commit(); + cp_wait(); + } else { + cp_wait<0>(); + } + __syncthreads(); + + const Stage& current = shared.staging[group_index % kStages]; + float group_acc[kNt][4] = {}; + +#pragma unroll + for (int ks = 0; ks < 4; ++ks) { + // Byte b = ks*8 + lid (K pair 2b within the warp's group) for af0/af1 and b + 4 for + // af2/af3; their fifth bits live in high bytes b>>2 = 2*ks and 2*ks + 1. + const int byte_col = code_off + ks * 8 + lid; + const int high_col = high_off + ks * 2; + const auto code_at = [&](int row, int byte) { + return current.codes[row][q5_ksplit_code_offset(row, byte)]; + }; + const auto high_at = [&](int row, int byte) { + return current.high[row][q5_ksplit_high_offset(row, byte)]; + }; + const unsigned af0 = q5_ksplit_bf16_pair(code_at(gid, byte_col), + high_at(gid, high_col), hshift); + const unsigned af1 = q5_ksplit_bf16_pair(code_at(gid + 8, byte_col), + high_at(gid + 8, high_col), hshift); + const unsigned af2 = q5_ksplit_bf16_pair(code_at(gid, byte_col + 4), + high_at(gid, high_col + 1), hshift); + const unsigned af3 = q5_ksplit_bf16_pair(code_at(gid + 8, byte_col + 4), + high_at(gid + 8, high_col + 1), hshift); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) { + unsigned bf0, bf1; + const int br = nt * 8 + b_rin; + ldmatrix_x2(bf0, bf1, + smem_addr(¤t.activations[k_split] + [br * kTileK + q5_ksplit_swizzle_64( + br, ks * 16 + b_koff)])); + mma_bf16(group_acc[nt][0], group_acc[nt][1], group_acc[nt][2], group_acc[nt][3], + af0, af1, af2, af3, bf0, bf1); + } + } + + const float top_scale = __half2float(__ushort_as_half(current.scales[gid][k_split])); + const float bot_scale = __half2float(__ushort_as_half(current.scales[gid + 8][k_split])); +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) { + acc[nt][0] = fmaf(group_acc[nt][0], top_scale, acc[nt][0]); + acc[nt][1] = fmaf(group_acc[nt][1], top_scale, acc[nt][1]); + acc[nt][2] = fmaf(group_acc[nt][2], bot_scale, acc[nt][2]); + acc[nt][3] = fmaf(group_acc[nt][3], bot_scale, acc[nt][3]); + } + + if (has_next) { + __syncthreads(); + if constexpr (kStages == 1) { + stage_weight(group_k0 + kGroupK, shared.staging[0]); + stage_x(group_k0 + kGroupK, shared.staging[0]); + cp_commit(); + } + } + } + + __syncthreads(); + auto* partial = shared.partial; + if ((k_split & 1) != 0) { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) { + store_vec(partial + ((k_split * kNt + nt) * 32 + lane) * 4, + make_float4(acc[nt][0], acc[nt][1], acc[nt][2], acc[nt][3])); + } + } + __syncthreads(); + + if ((k_split & 1) == 0) { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) { + const float4 partner = + load_vec(partial + (((k_split + 1) * kNt + nt) * 32 + lane) * 4); + acc[nt][0] += partner.x; + acc[nt][1] += partner.y; + acc[nt][2] += partner.z; + acc[nt][3] += partner.w; + if (k_split != 0) { + store_vec(partial + ((k_split * kNt + nt) * 32 + lane) * 4, + make_float4(acc[nt][0], acc[nt][1], acc[nt][2], acc[nt][3])); + } + } + } + __syncthreads(); + + if (k_split == 0) { +#pragma unroll + for (int nt = 0; nt < kNt; ++nt) { + float4 sum = make_float4(acc[nt][0], acc[nt][1], acc[nt][2], acc[nt][3]); +#pragma unroll + for (int split = 2; split < kWarps; split += 2) { + const float4 value = + load_vec(partial + ((split * kNt + nt) * 32 + lane) * 4); + sum.x += value.x; + sum.y += value.y; + sum.z += value.z; + sum.w += value.w; + } + const int col0 = nt * 8 + 2 * lid; + const auto store = [&](int col, int row, float value) { + __nv_bfloat16* destination = out + static_cast(col) * OutputRows + row; + if constexpr (AddResidual) { value += __bfloat162float(*destination); } + *destination = __float2bfloat16_rn(value); + }; + if (col0 < live_columns) { + store(col0, row0 + gid, sum.x); + store(col0, row0 + gid + 8, sum.z); + } + if (col0 + 1 < live_columns) { + store(col0 + 1, row0 + gid, sum.y); + store(col0 + 1, row0 + gid + 8, sum.w); + } + } + } +} + +// Exact-shape host launch. Capacity is the compile-time column capacity of one column tile; +// the live column count is a runtime argument, so one instance serves every T in +// (Capacity - 8, Capacity]. MaxColumns > Capacity launches ceil(T / Capacity) column tiles +// along blockIdx.y, streaming the weights once per tile; that stays far ahead of the wide GEMM +// tiles, which compute padded columns, until T approaches a few tiles. +template +void launch_q5_ksplit_mma(const Tensor& x, const Weight& weight, Tensor& out, + cudaStream_t stream) { + constexpr int kTileCols = (Capacity + 7) / 8 * 8; + static_assert(MaxColumns >= Capacity); + static_assert(MaxColumns == Capacity || Capacity == kTileCols, + "column tiling requires full-width tiles"); + if (weight.n != OutputRows || weight.padded_shape[1] != InputRows) { + throw std::invalid_argument("q5 K-split MMA: weight geometry differs from instance"); + } + const std::int32_t columns = x.ne[1]; + if (columns < 1 || columns > MaxColumns) { + throw std::invalid_argument("q5 K-split MMA: column count exceeds instance capacity"); + } + const dim3 grid(OutputRows / Q5KSplitMmaSchedule::kRowsPerCta, + static_cast((columns + kTileCols - 1) / kTileCols)); + q5_ksplit_mma_kernel + <<>>( + static_cast(x.data), + static_cast(weight.qdata), + static_cast(weight.qhigh), + static_cast(weight.scales), + static_cast<__nv_bfloat16*>(out.data), columns); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/linear/q5/shapes/n5120_k17408.cu b/src/ops/linear/q5/shapes/n5120_k17408.cu index a931d52468..0c10ca7bef 100644 --- a/src/ops/linear/q5/shapes/n5120_k17408.cu +++ b/src/ops/linear/q5/shapes/n5120_k17408.cu @@ -1,16 +1,11 @@ #include "ops/linear/q5/q5_shapes.h" -#include "ops/linear/q5/q5_ksplit_launch.cuh" +#include "ops/linear/q5/q5_ksplit_mma.cuh" namespace ninfer::ops::detail { Q5Launch select_q5_n5120_k17408(std::int32_t tokens) { if (tokens == 1) return launch_q5_split4_c1_k17408; - if (tokens <= 2) return launch_q5_ksplit<17408, 2, 2>; - if (tokens <= 3) return launch_q5_ksplit<17408, 3, 2>; - if (tokens <= 4) return launch_q5_ksplit<17408, 4, 2>; - if (tokens <= 5) return launch_q5_ksplit<17408, 5, 2>; - if (tokens <= 6) return launch_q5_ksplit<17408, 6, 2>; - if (tokens <= 15) return launch_q5_simt_r8_c4; + if (tokens <= 96) return launch_q5_ksplit_mma<5120, 17408, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; if (tokens <= 256) return launch_q5_mma_r32_c128; return launch_q5_mma_r64_c128; diff --git a/src/ops/linear/q5/shapes/n5120_k6144.cu b/src/ops/linear/q5/shapes/n5120_k6144.cu index add335e62e..d10289e2a1 100644 --- a/src/ops/linear/q5/shapes/n5120_k6144.cu +++ b/src/ops/linear/q5/shapes/n5120_k6144.cu @@ -1,16 +1,11 @@ #include "ops/linear/q5/q5_shapes.h" -#include "ops/linear/q5/q5_ksplit_launch.cuh" +#include "ops/linear/q5/q5_ksplit_mma.cuh" namespace ninfer::ops::detail { Q5Launch select_q5_n5120_k6144(std::int32_t tokens) { if (tokens == 1) return launch_q5_split4_c1_k6144; - if (tokens <= 2) return launch_q5_ksplit<6144, 2, 2>; - if (tokens <= 3) return launch_q5_ksplit<6144, 3, 2>; - if (tokens <= 4) return launch_q5_ksplit<6144, 4, 2>; - if (tokens <= 5) return launch_q5_ksplit<6144, 5, 2>; - if (tokens <= 6) return launch_q5_ksplit<6144, 6, 2>; - if (tokens <= 15) return launch_q5_simt_r8_c4; + if (tokens <= 96) return launch_q5_ksplit_mma<5120, 6144, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; if (tokens <= 256) return launch_q5_mma_r32_c128; return launch_q5_mma_r64_c128; diff --git a/src/ops/linear/q5/shapes/n6144_k5120.cu b/src/ops/linear/q5/shapes/n6144_k5120.cu index e1c2452621..b600db54d2 100644 --- a/src/ops/linear/q5/shapes/n6144_k5120.cu +++ b/src/ops/linear/q5/shapes/n6144_k5120.cu @@ -1,16 +1,13 @@ #include "ops/linear/q5/q5_shapes.h" +#include "ops/linear/q5/q5_ksplit_mma.cuh" #include "ops/linear/q5/q5_ksplit_launch.cuh" namespace ninfer::ops::detail { Q5Launch select_q5_n6144_k5120(std::int32_t tokens) { if (tokens == 1) return launch_q5_split4_c1_k5120; - if (tokens <= 2) return launch_q5_ksplit<5120, 2, 4>; - if (tokens <= 3) return launch_q5_ksplit<5120, 3, 4>; - if (tokens <= 4) return launch_q5_ksplit<5120, 4, 4>; - if (tokens <= 5) return launch_q5_ksplit<5120, 5, 4>; - if (tokens <= 6) return launch_q5_ksplit<5120, 6, 4>; - if (tokens <= 11) return launch_q5_simt_r8_c4; + if (tokens == 2) return launch_q5_ksplit<5120, 2, 4>; + if (tokens <= 96) return launch_q5_ksplit_mma<6144, 5120, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; return launch_q5_mma_r64_c128; } diff --git a/src/ops/linear/q5/shapes/n7168_k5120.cu b/src/ops/linear/q5/shapes/n7168_k5120.cu index 235aeafbc0..f98b2b94cf 100644 --- a/src/ops/linear/q5/shapes/n7168_k5120.cu +++ b/src/ops/linear/q5/shapes/n7168_k5120.cu @@ -1,16 +1,11 @@ #include "ops/linear/q5/q5_shapes.h" -#include "ops/linear/q5/q5_ksplit_launch.cuh" +#include "ops/linear/q5/q5_ksplit_mma.cuh" namespace ninfer::ops::detail { Q5Launch select_q5_n7168_k5120(std::int32_t tokens) { if (tokens == 1) return launch_q5_split4_c1_k5120; - if (tokens <= 2) return launch_q5_ksplit<5120, 2, 4>; - if (tokens <= 3) return launch_q5_ksplit<5120, 3, 4>; - if (tokens <= 4) return launch_q5_ksplit<5120, 4, 4>; - if (tokens <= 5) return launch_q5_ksplit<5120, 5, 4>; - if (tokens <= 6) return launch_q5_ksplit<5120, 6, 4>; - if (tokens <= 11) return launch_q5_simt_r8_c4; + if (tokens <= 96) return launch_q5_ksplit_mma<7168, 5120, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; return launch_q5_mma_r64_c128; } diff --git a/src/ops/linear_add/q5/q5_linear_add_gemm_mma.cu b/src/ops/linear_add/q5/q5_linear_add_gemm_mma.cu index 870251f8f3..675e832422 100644 --- a/src/ops/linear_add/q5/q5_linear_add_gemm_mma.cu +++ b/src/ops/linear_add/q5/q5_linear_add_gemm_mma.cu @@ -13,12 +13,6 @@ namespace ninfer::ops::detail { namespace { -using MmaR64C16Schedule = - Q5RowSplitMmaGemmSchedule<64, 16, 64, 16, 8, 2, 3, Q5FragmentPipeline::Serial, Cache::cg, - Cache::cg, Q5ScaleLoad::Pair32>; -using MmaR64C24Schedule = - Q5RowSplitMmaGemmSchedule<64, 24, 64, 16, 8, 2, 2, Q5FragmentPipeline::Serial, Cache::cg, - Cache::cg, Q5ScaleLoad::Pair32>; using MmaR64C32S3Schedule = Q5RowSplitMmaGemmSchedule<64, 32, 64, 16, 16, 3, 2, Q5FragmentPipeline::Serial, Cache::cg, Cache::cg, Q5ScaleLoad::Pair32>; @@ -67,16 +61,6 @@ void launch_route(const Tensor& x, const Weight& w, Tensor& residual_out, cudaSt } // namespace -void q5_linear_add_mma_r64_c16_launch(const Tensor& x, const Weight& w, Tensor& residual_out, - cudaStream_t stream) { - launch_route(x, w, residual_out, stream); -} - -void q5_linear_add_mma_r64_c24_launch(const Tensor& x, const Weight& w, Tensor& residual_out, - cudaStream_t stream) { - launch_route(x, w, residual_out, stream); -} - void q5_linear_add_mma_r64_c32_s3_launch(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) { launch_route(x, w, residual_out, stream); diff --git a/src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu b/src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu deleted file mode 100644 index 79aab9df29..0000000000 --- a/src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu +++ /dev/null @@ -1,75 +0,0 @@ -#include "core/weight.h" -#include "ops/linear_add/q5/q5_linear_add_kernels.h" - -#include "core/device.h" -#include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh" - -#include - -#include -#include - -namespace ninfer::ops::detail { -namespace { - -template -void launch_split2(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) { - constexpr int kThreads = 2 * 32; - const dim3 grid(static_cast(residual_out.ne[0]), 1u, 1u); - q5_rowsplit_gemm_simt_split2_kernel<<>>( - static_cast(x.data), static_cast(w.qdata), - static_cast(w.qhigh), static_cast(w.scales), - static_cast<__nv_bfloat16*>(residual_out.data), residual_out.ne[0], x.ne[0], x.ne[1], - w.padded_shape[1], FullSlabs); -} - -template -void dispatch_shape(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) { - if (w.k == 6144) { - launch_split2(x, w, residual_out, stream); - } else if (w.k == 17408) { - launch_split2(x, w, residual_out, stream); - } else { - throw std::invalid_argument("q5 linear_add split2: unsupported exact K"); - } -} - -template -void dispatch_cols(std::int32_t cols, Launch&& launch) { - switch (cols) { -#define NINFER_Q5_LINEAR_ADD_EXACT(COLS) \ - case COLS: \ - launch.template operator()(); \ - return - NINFER_Q5_LINEAR_ADD_EXACT(1); - NINFER_Q5_LINEAR_ADD_EXACT(2); - NINFER_Q5_LINEAR_ADD_EXACT(3); - NINFER_Q5_LINEAR_ADD_EXACT(4); - NINFER_Q5_LINEAR_ADD_EXACT(5); - NINFER_Q5_LINEAR_ADD_EXACT(6); - NINFER_Q5_LINEAR_ADD_EXACT(7); - NINFER_Q5_LINEAR_ADD_EXACT(8); - NINFER_Q5_LINEAR_ADD_EXACT(9); - NINFER_Q5_LINEAR_ADD_EXACT(10); - NINFER_Q5_LINEAR_ADD_EXACT(11); - NINFER_Q5_LINEAR_ADD_EXACT(12); - NINFER_Q5_LINEAR_ADD_EXACT(13); - NINFER_Q5_LINEAR_ADD_EXACT(14); - NINFER_Q5_LINEAR_ADD_EXACT(15); - NINFER_Q5_LINEAR_ADD_EXACT(16); -#undef NINFER_Q5_LINEAR_ADD_EXACT - default: - throw std::invalid_argument("q5 linear_add split2: T must be in [1,16]"); - } -} - -} // namespace - -void q5_linear_add_split2_exact_launch(const Tensor& x, const Weight& w, Tensor& residual_out, - cudaStream_t stream) { - dispatch_cols(x.ne[1], [&]() { dispatch_shape(x, w, residual_out, stream); }); - CUDA_CHECK(cudaGetLastError()); -} - -} // namespace ninfer::ops::detail diff --git a/src/ops/linear_add/q5/q5_linear_add_kernels.h b/src/ops/linear_add/q5/q5_linear_add_kernels.h index 0c43fd6ef8..9ee80d7c53 100644 --- a/src/ops/linear_add/q5/q5_linear_add_kernels.h +++ b/src/ops/linear_add/q5/q5_linear_add_kernels.h @@ -7,12 +7,8 @@ namespace ninfer::ops::detail { -void q5_linear_add_split2_exact_launch(const Tensor& x, const Weight& w, Tensor& residual_out, - cudaStream_t stream); -void q5_linear_add_mma_r64_c16_launch(const Tensor& x, const Weight& w, Tensor& residual_out, - cudaStream_t stream); -void q5_linear_add_mma_r64_c24_launch(const Tensor& x, const Weight& w, Tensor& residual_out, - cudaStream_t stream); +void q5_linear_add_ksplit_mma_residual_launch(const Tensor& x, const Weight& w, + Tensor& residual_out, cudaStream_t stream); void q5_linear_add_mma_r64_c32_s3_launch(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream); void q5_linear_add_mma_r64_c32_s4_launch(const Tensor& x, const Weight& w, Tensor& residual_out, diff --git a/src/ops/linear_add/q5/q5_linear_add_ksplit_mma.cu b/src/ops/linear_add/q5/q5_linear_add_ksplit_mma.cu new file mode 100644 index 0000000000..a4a67a25fe --- /dev/null +++ b/src/ops/linear_add/q5/q5_linear_add_ksplit_mma.cu @@ -0,0 +1,51 @@ +#include "core/weight.h" +#include "ops/linear_add/q5/q5_linear_add_kernels.h" + +#include "core/device.h" +#include "ops/linear/q5/q5_ksplit_mma.cuh" + +#include +#include + +namespace ninfer::ops::detail { +namespace { + +constexpr int kRows = 5120; + +// Capacity tiles 8/16/24/32 serve T <= 32 with one column tile; T in (32, 64] runs two +// 32-column tiles along blockIdx.y. The residual form reads and rewrites residual_out. +template +void launch_shape(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) { + const std::int32_t cols = x.ne[1]; + if (cols <= 8) { + launch_q5_ksplit_mma(x, w, residual_out, stream); + } else if (cols <= 16) { + launch_q5_ksplit_mma(x, w, residual_out, stream); + } else if (cols <= 24) { + launch_q5_ksplit_mma(x, w, residual_out, stream); + } else if (cols <= 32) { + launch_q5_ksplit_mma(x, w, residual_out, stream); + } else if (cols <= 64) { + launch_q5_ksplit_mma(x, w, residual_out, stream); + } else { + throw std::invalid_argument("q5 linear_add K-split MMA: T must be in [1,64]"); + } +} + +} // namespace + +void q5_linear_add_ksplit_mma_residual_launch(const Tensor& x, const Weight& w, + Tensor& residual_out, cudaStream_t stream) { + if (residual_out.ne[0] != kRows) { + throw std::invalid_argument("q5 linear_add K-split MMA requires 5120 output rows"); + } + if (w.k == 6144) { + launch_shape<6144>(x, w, residual_out, stream); + } else if (w.k == 17408) { + launch_shape<17408>(x, w, residual_out, stream); + } else { + throw std::invalid_argument("q5 linear_add K-split MMA: unsupported exact K"); + } +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/linear_add/q5/q5_linear_add_plan.cpp b/src/ops/linear_add/q5/q5_linear_add_plan.cpp index 5a0d3a14b6..189a70546f 100644 --- a/src/ops/linear_add/q5/q5_linear_add_plan.cpp +++ b/src/ops/linear_add/q5/q5_linear_add_plan.cpp @@ -37,20 +37,19 @@ constexpr std::array kSupports{{ {5120, 17408, 17408}, }}; -constexpr std::array kK6144Routes{{ - {{1, 13}, Q5LinearAddScheduleId::Split2ExactResidual}, - {{14, 32}, Q5LinearAddScheduleId::MmaResidualR64C16}, - {{33, 48}, Q5LinearAddScheduleId::MmaResidualR64C24}, - {{49, 192}, Q5LinearAddScheduleId::MmaResidualR64C32S4}, +// The K-split MMA route streams the weights once per 32-column tile and stays ahead of the +// 64-row GEMM tiles through two column tiles (measured crossover: T=60 at K=6144, T=64 at +// K=17408); beyond that the wide tiles amortize better. +constexpr std::array kK6144Routes{{ + {{1, 60}, Q5LinearAddScheduleId::KSplitMmaResidual}, + {{61, 192}, Q5LinearAddScheduleId::MmaResidualR64C32S4}, {{193, 512}, Q5LinearAddScheduleId::MmaResidualR64C128}, {{513, kAnyCols}, Q5LinearAddScheduleId::MmaResidualR64C128Tail}, }}; -constexpr std::array kK17408Routes{{ - {{1, 16}, Q5LinearAddScheduleId::Split2ExactResidual}, - {{17, 32}, Q5LinearAddScheduleId::MmaResidualR64C16}, - {{33, 48}, Q5LinearAddScheduleId::MmaResidualR64C24}, - {{49, 192}, Q5LinearAddScheduleId::MmaResidualR64C32S3}, +constexpr std::array kK17408Routes{{ + {{1, 64}, Q5LinearAddScheduleId::KSplitMmaResidual}, + {{65, 192}, Q5LinearAddScheduleId::MmaResidualR64C32S3}, {{193, 512}, Q5LinearAddScheduleId::MmaResidualR64C128}, {{513, kAnyCols}, Q5LinearAddScheduleId::MmaResidualR64C128Tail}, }}; @@ -114,12 +113,8 @@ void launch_wide_with_narrow_tail(const Tensor& x, const Weight& w, Tensor& resi const char* q5_linear_add_schedule_name(Q5LinearAddScheduleId schedule) noexcept { switch (schedule) { - case Q5LinearAddScheduleId::Split2ExactResidual: - return "linear_add.q5.simt.split2.exact.residual"; - case Q5LinearAddScheduleId::MmaResidualR64C16: - return "linear_add.q5.mma.r64.c16.cta_collective_residual"; - case Q5LinearAddScheduleId::MmaResidualR64C24: - return "linear_add.q5.mma.r64.c24.cta_collective_residual"; + case Q5LinearAddScheduleId::KSplitMmaResidual: + return "linear_add.q5.mma.ksplit.residual"; case Q5LinearAddScheduleId::MmaResidualR64C32S3: return "linear_add.q5.mma.r64.c32.s3.cta_collective_residual"; case Q5LinearAddScheduleId::MmaResidualR64C32S4: @@ -172,14 +167,8 @@ void q5_linear_add_execute_plan(const Q5LinearAddPlan& plan, const Tensor& x, co (void)ws; switch (plan.schedule) { - case Q5LinearAddScheduleId::Split2ExactResidual: - q5_linear_add_split2_exact_launch(x, w, residual_out, stream); - return; - case Q5LinearAddScheduleId::MmaResidualR64C16: - q5_linear_add_mma_r64_c16_launch(x, w, residual_out, stream); - return; - case Q5LinearAddScheduleId::MmaResidualR64C24: - q5_linear_add_mma_r64_c24_launch(x, w, residual_out, stream); + case Q5LinearAddScheduleId::KSplitMmaResidual: + q5_linear_add_ksplit_mma_residual_launch(x, w, residual_out, stream); return; case Q5LinearAddScheduleId::MmaResidualR64C32S3: q5_linear_add_mma_r64_c32_s3_launch(x, w, residual_out, stream); diff --git a/src/ops/linear_add/q5/q5_linear_add_plan.h b/src/ops/linear_add/q5/q5_linear_add_plan.h index 055431e2bb..9642e9e9c2 100644 --- a/src/ops/linear_add/q5/q5_linear_add_plan.h +++ b/src/ops/linear_add/q5/q5_linear_add_plan.h @@ -11,9 +11,7 @@ namespace ninfer::ops::detail { enum class Q5LinearAddScheduleId { - Split2ExactResidual, - MmaResidualR64C16, - MmaResidualR64C24, + KSplitMmaResidual, MmaResidualR64C32S3, MmaResidualR64C32S4, MmaResidualR64C128, diff --git a/src/ops/linear_add/sources.cmake b/src/ops/linear_add/sources.cmake index a1414e9c71..8f19503912 100644 --- a/src/ops/linear_add/sources.cmake +++ b/src/ops/linear_add/sources.cmake @@ -13,7 +13,7 @@ target_sources(ninfer_ops PRIVATE "${CMAKE_CURRENT_LIST_DIR}/fp8/fp8_linear_add_plan.cpp" "${CMAKE_CURRENT_LIST_DIR}/q4/q4_linear_add.cu" "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_gemm_mma.cu" - "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_gemm_simt.cu" + "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_ksplit_mma.cu" "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_plan.cpp" "${CMAKE_CURRENT_LIST_DIR}/q8/q8_linear_add_gemm_mma.cu" "${CMAKE_CURRENT_LIST_DIR}/q8/q8_linear_add_gemm_simt.cu" diff --git a/tests/ops/linear/test_q5_a16.cpp b/tests/ops/linear/test_q5_a16.cpp index f1a9e6ccf3..bbcc7c2590 100644 --- a/tests/ops/linear/test_q5_a16.cpp +++ b/tests/ops/linear/test_q5_a16.cpp @@ -44,12 +44,22 @@ int q5_a16_conformance() { Invocation{4, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, Invocation{6, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, a16(7), + a16(8), + a16(9), a16(11), graph(12), + a16(15), + a16(16), + a16(17), a16(24), a16(25), + a16(31), + a16(32), + a16(33), a16(64), a16(65), + a16(96), + a16(97), a16(112), a16(113), a16(128), @@ -68,10 +78,20 @@ int q5_a16_conformance() { Invocation{4, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, Invocation{6, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, a16(7), + a16(8), + a16(9), a16(11), graph(12), + a16(15), a16(16), a16(17), + a16(24), + a16(25), + a16(31), + a16(32), + a16(33), + a16(96), + a16(97), a16(112), a16(113), a16(128), @@ -90,10 +110,18 @@ int q5_a16_conformance() { Invocation{4, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, Invocation{6, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, a16(7), + a16(8), + a16(9), a16(15), graph(16), + a16(17), a16(24), a16(25), + a16(31), + a16(32), + a16(33), + a16(96), + a16(97), a16(112), a16(113), graph(128), @@ -114,10 +142,18 @@ int q5_a16_conformance() { Invocation{4, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, Invocation{6, CallForm::Policy, ninfer::ops::LinearPolicy::A16Only, true}, a16(7), + a16(8), + a16(9), a16(15), graph(16), + a16(17), a16(24), a16(25), + a16(31), + a16(32), + a16(33), + a16(96), + a16(97), a16(112), a16(113), graph(128), diff --git a/tests/ops/linear_add/test_q5_a16.cpp b/tests/ops/linear_add/test_q5_a16.cpp index 8a940a7134..dcd210d8bb 100644 --- a/tests/ops/linear_add/test_q5_a16.cpp +++ b/tests/ops/linear_add/test_q5_a16.cpp @@ -10,32 +10,21 @@ using ninfer::test::linear_add::ShapeCase; using ninfer::test::linear_add::WeightFormat; int q5_a16_conformance() { - // Starts of the registered positive-T regions. run_shape checks b-1/b/b+1 for every start, - // plus one interior point for every region, through the public Op. The first region now - // reaches down to T=1 (T=1 is covered by the leading probe and its own interior point), and - // 513 starts the whole-wave + narrow-tail composite, whose interiors cover a short, a - // 128-column and a 256-column tail (the last one keeping the single wide launch). - // - // composite_offset re-probes the same starts one whole wave (512 columns) later, where the - // tail - not the leading launch - is the one resolving against them: 525/526/527, 544/545/546, - // 560/561/562 for k=6144 (528/529/530 for k=17408), 704/705/706 for the tail=192/193 - // composite-to-wide fallback, and 1024/1025/1026 for a zero, one and two column tail after the - // second whole wave. - constexpr std::array kInteriors{1, 2, 3, 8, 24, 40, 56, 64, + // Cover the k-split capacity tiles, the wide-route crossover, and the narrow-tail composite. + constexpr std::array kInteriors{1, 2, 3, 8, 24, 40, 56, 64, 96, 128, 129, 256, 640, 768, 1024}; - constexpr std::array kK6144RouteStarts{14, 33, 49, 193, 513}; + constexpr std::array kK6144RouteStarts{9, 17, 25, 33, 61, 193, 513}; constexpr std::array kK6144GraphTokens{513, 526, 545, 561, 705, 1025}; int failures = 0; failures += ninfer::test::linear_add::run_shape( "Q5_A16 LinearAdd", WeightFormat::Q5G64F16S, ShapeCase{5120, 6144, 401U, kK6144RouteStarts, kInteriors, kK6144GraphTokens, false, 512}); - constexpr std::array kK17408RouteStarts{17, 33, 49, 193, 513}; + constexpr std::array kK17408RouteStarts{9, 17, 25, 33, 65, 193, 513}; constexpr std::array kK17408GraphTokens{513, 529, 545, 561, 705, 1025}; - failures += - ninfer::test::linear_add::run_shape("Q5_A16 LinearAdd", WeightFormat::Q5G64F16S, - ShapeCase{5120, 17408, 409U, kK17408RouteStarts, - kInteriors, kK17408GraphTokens, false, 512}); + failures += ninfer::test::linear_add::run_shape( + "Q5_A16 LinearAdd", WeightFormat::Q5G64F16S, + ShapeCase{5120, 17408, 409U, kK17408RouteStarts, kInteriors, kK17408GraphTokens, false, 512}); return failures; } From 2b09c8f92dc0a27c96e5053a1550fc5134338ef6 Mon Sep 17 00:00:00 2001 From: giveen Date: Fri, 18 Sep 2026 22:03:03 -0600 Subject: [PATCH 3/4] perf(ops): keep the Q5 LinearAdd T=1 route on the split2 SIMT kernel The previous commit collapsed the small-T Q5 LinearAdd bands into a single {1,60}/{1,64} K-split MMA band whose crossover was fit against the batched end of the range. At T=1 that route leaves most of a 32-column tile idle, so ordinary (non-speculative) decode pays for tiles it cannot fill. Plain Q5 linear already carves T=1 out to a dedicated kernel; do the same here. Restores q5_linear_add_gemm_simt.cu and routes T=1 back to Split2ExactResidual, leaving T>=2 on the K-split MMA route. Speculative decode verifies at T>=2 and is untouched. Measured on qwen3.8-27b, RTX 5090, pp2048+tg512, r=5, both builds compiled back to back in one session: ordinary decode, bf16 KV: 76.12 -> 77.22 tok/s (+1.4%) ordinary decode, fp8 KV: 77.56 -> 78.16 tok/s (+0.8%) MTP3 decode, fp8 KV: 222.87 tok/s, acceptance 93.32% (unchanged) Note: an earlier report put the T=1 regression at ~4%. That magnitude did not reproduce here under either KV dtype; the recovered margin is ~1%, at roughly 3 sigma in the bf16 pair. The boundary itself is still unswept - T=2 and T=4 are untested candidates. Co-Authored-By: Claude Opus 5 --- .../linear_add/q5/q5_linear_add_gemm_simt.cu | 75 +++++++++++++++++++ src/ops/linear_add/q5/q5_linear_add_kernels.h | 2 + src/ops/linear_add/q5/q5_linear_add_plan.cpp | 24 ++++-- src/ops/linear_add/q5/q5_linear_add_plan.h | 1 + src/ops/linear_add/sources.cmake | 1 + 5 files changed, 96 insertions(+), 7 deletions(-) create mode 100644 src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu diff --git a/src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu b/src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu new file mode 100644 index 0000000000..79aab9df29 --- /dev/null +++ b/src/ops/linear_add/q5/q5_linear_add_gemm_simt.cu @@ -0,0 +1,75 @@ +#include "core/weight.h" +#include "ops/linear_add/q5/q5_linear_add_kernels.h" + +#include "core/device.h" +#include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh" + +#include + +#include +#include + +namespace ninfer::ops::detail { +namespace { + +template +void launch_split2(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) { + constexpr int kThreads = 2 * 32; + const dim3 grid(static_cast(residual_out.ne[0]), 1u, 1u); + q5_rowsplit_gemm_simt_split2_kernel<<>>( + static_cast(x.data), static_cast(w.qdata), + static_cast(w.qhigh), static_cast(w.scales), + static_cast<__nv_bfloat16*>(residual_out.data), residual_out.ne[0], x.ne[0], x.ne[1], + w.padded_shape[1], FullSlabs); +} + +template +void dispatch_shape(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream) { + if (w.k == 6144) { + launch_split2(x, w, residual_out, stream); + } else if (w.k == 17408) { + launch_split2(x, w, residual_out, stream); + } else { + throw std::invalid_argument("q5 linear_add split2: unsupported exact K"); + } +} + +template +void dispatch_cols(std::int32_t cols, Launch&& launch) { + switch (cols) { +#define NINFER_Q5_LINEAR_ADD_EXACT(COLS) \ + case COLS: \ + launch.template operator()(); \ + return + NINFER_Q5_LINEAR_ADD_EXACT(1); + NINFER_Q5_LINEAR_ADD_EXACT(2); + NINFER_Q5_LINEAR_ADD_EXACT(3); + NINFER_Q5_LINEAR_ADD_EXACT(4); + NINFER_Q5_LINEAR_ADD_EXACT(5); + NINFER_Q5_LINEAR_ADD_EXACT(6); + NINFER_Q5_LINEAR_ADD_EXACT(7); + NINFER_Q5_LINEAR_ADD_EXACT(8); + NINFER_Q5_LINEAR_ADD_EXACT(9); + NINFER_Q5_LINEAR_ADD_EXACT(10); + NINFER_Q5_LINEAR_ADD_EXACT(11); + NINFER_Q5_LINEAR_ADD_EXACT(12); + NINFER_Q5_LINEAR_ADD_EXACT(13); + NINFER_Q5_LINEAR_ADD_EXACT(14); + NINFER_Q5_LINEAR_ADD_EXACT(15); + NINFER_Q5_LINEAR_ADD_EXACT(16); +#undef NINFER_Q5_LINEAR_ADD_EXACT + default: + throw std::invalid_argument("q5 linear_add split2: T must be in [1,16]"); + } +} + +} // namespace + +void q5_linear_add_split2_exact_launch(const Tensor& x, const Weight& w, Tensor& residual_out, + cudaStream_t stream) { + dispatch_cols(x.ne[1], [&]() { dispatch_shape(x, w, residual_out, stream); }); + CUDA_CHECK(cudaGetLastError()); +} + +} // namespace ninfer::ops::detail diff --git a/src/ops/linear_add/q5/q5_linear_add_kernels.h b/src/ops/linear_add/q5/q5_linear_add_kernels.h index 9ee80d7c53..023f424824 100644 --- a/src/ops/linear_add/q5/q5_linear_add_kernels.h +++ b/src/ops/linear_add/q5/q5_linear_add_kernels.h @@ -7,6 +7,8 @@ namespace ninfer::ops::detail { +void q5_linear_add_split2_exact_launch(const Tensor& x, const Weight& w, Tensor& residual_out, + cudaStream_t stream); void q5_linear_add_ksplit_mma_residual_launch(const Tensor& x, const Weight& w, Tensor& residual_out, cudaStream_t stream); void q5_linear_add_mma_r64_c32_s3_launch(const Tensor& x, const Weight& w, Tensor& residual_out, diff --git a/src/ops/linear_add/q5/q5_linear_add_plan.cpp b/src/ops/linear_add/q5/q5_linear_add_plan.cpp index 189a70546f..f7bc66fd9e 100644 --- a/src/ops/linear_add/q5/q5_linear_add_plan.cpp +++ b/src/ops/linear_add/q5/q5_linear_add_plan.cpp @@ -37,18 +37,23 @@ constexpr std::array kSupports{{ {5120, 17408, 17408}, }}; -// The K-split MMA route streams the weights once per 32-column tile and stays ahead of the -// 64-row GEMM tiles through two column tiles (measured crossover: T=60 at K=6144, T=64 at -// K=17408); beyond that the wide tiles amortize better. -constexpr std::array kK6144Routes{{ - {{1, 60}, Q5LinearAddScheduleId::KSplitMmaResidual}, +// T=1 keeps the split2 SIMT kernel: the K-split MMA route wastes most of a 32-column tile on a +// single row, worth ~1% of ordinary (non-speculative) decode end to end on qwen3.8-27b/RTX 5090. +// Speculative decode verifies at T>=2 and is unaffected. From T=2 the K-split +// MMA route streams the weights once per 32-column tile and stays ahead of the 64-row GEMM tiles +// through two column tiles (measured crossover: T=60 at K=6144, T=64 at K=17408); beyond that the +// wide tiles amortize better. +constexpr std::array kK6144Routes{{ + {{1, 1}, Q5LinearAddScheduleId::Split2ExactResidual}, + {{2, 60}, Q5LinearAddScheduleId::KSplitMmaResidual}, {{61, 192}, Q5LinearAddScheduleId::MmaResidualR64C32S4}, {{193, 512}, Q5LinearAddScheduleId::MmaResidualR64C128}, {{513, kAnyCols}, Q5LinearAddScheduleId::MmaResidualR64C128Tail}, }}; -constexpr std::array kK17408Routes{{ - {{1, 64}, Q5LinearAddScheduleId::KSplitMmaResidual}, +constexpr std::array kK17408Routes{{ + {{1, 1}, Q5LinearAddScheduleId::Split2ExactResidual}, + {{2, 64}, Q5LinearAddScheduleId::KSplitMmaResidual}, {{65, 192}, Q5LinearAddScheduleId::MmaResidualR64C32S3}, {{193, 512}, Q5LinearAddScheduleId::MmaResidualR64C128}, {{513, kAnyCols}, Q5LinearAddScheduleId::MmaResidualR64C128Tail}, @@ -113,6 +118,8 @@ void launch_wide_with_narrow_tail(const Tensor& x, const Weight& w, Tensor& resi const char* q5_linear_add_schedule_name(Q5LinearAddScheduleId schedule) noexcept { switch (schedule) { + case Q5LinearAddScheduleId::Split2ExactResidual: + return "linear_add.q5.simt.split2.exact.residual"; case Q5LinearAddScheduleId::KSplitMmaResidual: return "linear_add.q5.mma.ksplit.residual"; case Q5LinearAddScheduleId::MmaResidualR64C32S3: @@ -167,6 +174,9 @@ void q5_linear_add_execute_plan(const Q5LinearAddPlan& plan, const Tensor& x, co (void)ws; switch (plan.schedule) { + case Q5LinearAddScheduleId::Split2ExactResidual: + q5_linear_add_split2_exact_launch(x, w, residual_out, stream); + return; case Q5LinearAddScheduleId::KSplitMmaResidual: q5_linear_add_ksplit_mma_residual_launch(x, w, residual_out, stream); return; diff --git a/src/ops/linear_add/q5/q5_linear_add_plan.h b/src/ops/linear_add/q5/q5_linear_add_plan.h index 9642e9e9c2..ed9297d0d3 100644 --- a/src/ops/linear_add/q5/q5_linear_add_plan.h +++ b/src/ops/linear_add/q5/q5_linear_add_plan.h @@ -11,6 +11,7 @@ namespace ninfer::ops::detail { enum class Q5LinearAddScheduleId { + Split2ExactResidual, KSplitMmaResidual, MmaResidualR64C32S3, MmaResidualR64C32S4, diff --git a/src/ops/linear_add/sources.cmake b/src/ops/linear_add/sources.cmake index 8f19503912..35d921af75 100644 --- a/src/ops/linear_add/sources.cmake +++ b/src/ops/linear_add/sources.cmake @@ -13,6 +13,7 @@ target_sources(ninfer_ops PRIVATE "${CMAKE_CURRENT_LIST_DIR}/fp8/fp8_linear_add_plan.cpp" "${CMAKE_CURRENT_LIST_DIR}/q4/q4_linear_add.cu" "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_gemm_mma.cu" + "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_gemm_simt.cu" "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_ksplit_mma.cu" "${CMAKE_CURRENT_LIST_DIR}/q5/q5_linear_add_plan.cpp" "${CMAKE_CURRENT_LIST_DIR}/q8/q8_linear_add_gemm_mma.cu" From 248ba3f6750bad538a3e47692f65b2ecce3bd1f4 Mon Sep 17 00:00:00 2001 From: giveen Date: Fri, 18 Sep 2026 22:32:32 -0600 Subject: [PATCH 4/4] perf(ops): ladder the Q5 linear K-split capacity to the token count The Q5 shape dispatchers routed all of T<=96 to a single Capacity=32 K-split MMA instance. Capacity is compile time: it sizes the accumulator fragments and kItemsPerSplit, so one wide instance makes small T stage and accumulate columns it never fills. The cost was flat and visible - at N=5120/K=6144, T=2 took the same 28.7us as T=17. Q4 has laddered its K-split rungs since it landed (4/8/16/24 per shape), and the Q5 linear_add route already instantiates the same template at 8/16/24/32. Only the Q5 plain-linear dispatchers were left on a single rung. This wires them to rungs 4/8/16, then the existing Capacity=32 column-tiled instance for T<=96. No new kernels. T=1 moves to the 4-rung only on the two N=5120 shapes, where it beats the dedicated split4_c1 GEMV; N=6144 and N=7168 keep that GEMV, which is still faster for them. A 24-rung was measured and rejected: it lost to the column-tiled Capacity=32 instance across T=17..24. Measured with ninfer_linear_bench, --repeat 100, cold cache, RTX 5090, min_us, against a build of the parent commit. T=17..32 is an unchanged control band and brackets per-run drift: shape T=1 T=2..16 T=17..32 (control) 5120x6144 -11.3% -27.5% +0.0% 5120x17408 -7.8% -26.5% -0.0% 6144x5120 -0.4% -29.4% -0.1% 7168x5120 -0.2% -33.4% +6.8% End-to-end on qwen3.8-27b, pp2048+tg512, fp8 KV, this is flat: ordinary decode 77.74 vs 78.16 tok/s and MTP3 221.08 vs 222.87 tok/s, both within a sigma of the parent commit. The op-level win is real and repeatable; this model's decode path just does not spend enough time in these four Q5 plain-linear shapes for it to surface. Kept because it is strictly faster per call at no cost, but it is not an end-to-end throughput change and should not be reported as one. launch_q5_split4_c1_k6144 and launch_q5_split4_c1_k17408 now have no caller. Left in place rather than deleted. Co-Authored-By: Claude Opus 5 --- src/ops/linear/q5/q5_ksplit_mma.cuh | 7 +++++++ src/ops/linear/q5/shapes/n5120_k17408.cu | 4 +++- src/ops/linear/q5/shapes/n5120_k6144.cu | 4 +++- src/ops/linear/q5/shapes/n6144_k5120.cu | 3 +++ src/ops/linear/q5/shapes/n7168_k5120.cu | 3 +++ 5 files changed, 19 insertions(+), 2 deletions(-) diff --git a/src/ops/linear/q5/q5_ksplit_mma.cuh b/src/ops/linear/q5/q5_ksplit_mma.cuh index ebb9b8ccc3..967ce28e54 100644 --- a/src/ops/linear/q5/q5_ksplit_mma.cuh +++ b/src/ops/linear/q5/q5_ksplit_mma.cuh @@ -369,6 +369,13 @@ __launch_bounds__(Q5KSplitMmaSchedule::kThreads, q5_ksplit_min_blocks_per_sm Capacity launches ceil(T / Capacity) column tiles // along blockIdx.y, streaming the weights once per tile; that stays far ahead of the wide GEMM // tiles, which compute padded columns, until T approaches a few tiles. +// +// Capacity is compile time and sizes the accumulator fragments and the staged item count, so a +// single wide instance makes small T pay for columns it never fills: one Capacity=32 instance +// serving all of T<=96 costs the same at T=2 as at T=17. Callers should ladder Capacity to T +// (4/8/16, then 32 with column tiling) the way the shape dispatchers and the linear_add route do; +// measured on an RTX 5090, the ladder is 20-40% faster than a lone Capacity=32 instance across +// T=2..16 on every Q5 shape. template void launch_q5_ksplit_mma(const Tensor& x, const Weight& weight, Tensor& out, diff --git a/src/ops/linear/q5/shapes/n5120_k17408.cu b/src/ops/linear/q5/shapes/n5120_k17408.cu index 0c10ca7bef..1d5a61e464 100644 --- a/src/ops/linear/q5/shapes/n5120_k17408.cu +++ b/src/ops/linear/q5/shapes/n5120_k17408.cu @@ -4,7 +4,9 @@ namespace ninfer::ops::detail { Q5Launch select_q5_n5120_k17408(std::int32_t tokens) { - if (tokens == 1) return launch_q5_split4_c1_k17408; + if (tokens <= 4) return launch_q5_ksplit_mma<5120, 17408, 4>; + if (tokens <= 8) return launch_q5_ksplit_mma<5120, 17408, 8>; + if (tokens <= 16) return launch_q5_ksplit_mma<5120, 17408, 16>; if (tokens <= 96) return launch_q5_ksplit_mma<5120, 17408, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; if (tokens <= 256) return launch_q5_mma_r32_c128; diff --git a/src/ops/linear/q5/shapes/n5120_k6144.cu b/src/ops/linear/q5/shapes/n5120_k6144.cu index d10289e2a1..1574074629 100644 --- a/src/ops/linear/q5/shapes/n5120_k6144.cu +++ b/src/ops/linear/q5/shapes/n5120_k6144.cu @@ -4,7 +4,9 @@ namespace ninfer::ops::detail { Q5Launch select_q5_n5120_k6144(std::int32_t tokens) { - if (tokens == 1) return launch_q5_split4_c1_k6144; + if (tokens <= 4) return launch_q5_ksplit_mma<5120, 6144, 4>; + if (tokens <= 8) return launch_q5_ksplit_mma<5120, 6144, 8>; + if (tokens <= 16) return launch_q5_ksplit_mma<5120, 6144, 16>; if (tokens <= 96) return launch_q5_ksplit_mma<5120, 6144, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; if (tokens <= 256) return launch_q5_mma_r32_c128; diff --git a/src/ops/linear/q5/shapes/n6144_k5120.cu b/src/ops/linear/q5/shapes/n6144_k5120.cu index b600db54d2..494787d344 100644 --- a/src/ops/linear/q5/shapes/n6144_k5120.cu +++ b/src/ops/linear/q5/shapes/n6144_k5120.cu @@ -7,6 +7,9 @@ namespace ninfer::ops::detail { Q5Launch select_q5_n6144_k5120(std::int32_t tokens) { if (tokens == 1) return launch_q5_split4_c1_k5120; if (tokens == 2) return launch_q5_ksplit<5120, 2, 4>; + if (tokens <= 4) return launch_q5_ksplit_mma<6144, 5120, 4>; + if (tokens <= 8) return launch_q5_ksplit_mma<6144, 5120, 8>; + if (tokens <= 16) return launch_q5_ksplit_mma<6144, 5120, 16>; if (tokens <= 96) return launch_q5_ksplit_mma<6144, 5120, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; return launch_q5_mma_r64_c128; diff --git a/src/ops/linear/q5/shapes/n7168_k5120.cu b/src/ops/linear/q5/shapes/n7168_k5120.cu index f98b2b94cf..85b8919158 100644 --- a/src/ops/linear/q5/shapes/n7168_k5120.cu +++ b/src/ops/linear/q5/shapes/n7168_k5120.cu @@ -5,6 +5,9 @@ namespace ninfer::ops::detail { Q5Launch select_q5_n7168_k5120(std::int32_t tokens) { if (tokens == 1) return launch_q5_split4_c1_k5120; + if (tokens <= 4) return launch_q5_ksplit_mma<7168, 5120, 4>; + if (tokens <= 8) return launch_q5_ksplit_mma<7168, 5120, 8>; + if (tokens <= 16) return launch_q5_ksplit_mma<7168, 5120, 16>; if (tokens <= 96) return launch_q5_ksplit_mma<7168, 5120, 32, 96>; if (tokens <= 112) return launch_q5_mma_r64_c32_s3; return launch_q5_mma_r64_c128;