From 0435e219b561067f9d45457c5cd36d126e4de345 Mon Sep 17 00:00:00 2001 From: Minnnn <6550671@qq.com> Date: Wed, 23 Sep 2026 21:34:31 +0800 Subject: [PATCH] perf(ops): one split4 Q5 parent for the q4/q5 small-column input projections The Q5 parent of both Q4/Q5 input projections now uses the split4 shape - one CTA owns one output row and its four warps split the K dimension, reducing their partial sums through shared memory - instead of the row-block shape at 7..8 and the c4 narrow tile at 9..12. Per-side projection route change: no new kernel, no change to the complete-op structure, no change to the column catalog or to the grouped bands. Complete public Op medians, B0 (master + #271) against this branch. GDN: B=1 W=7 87.81 -> 69.31 us (21.1%), W=8 91.39 -> 76.86 (15.9%), W=9 108.29 -> 97.57 (9.9%), W=10 109.86 -> 105.86 (3.6%); the aggregate-10 organisations B=2/W=5 and B=5/W=2 by 4-5%; Record at 7/8/9/10 by 21.6/17.4/9.5/3.7%. Attention: 18.8/14.6/5.4% at 7/8/9. Both band ends are measured crossovers between two legal shapes, per parent. The GDN parent is extended to 10: split4 wins at 10 in every organisation that exposes 10 aggregate columns and loses at 11 and 12, so the band stops there. The attention parent stays at 2..9: it ties at 10 and loses at 11/12. Q4 is untouched - T=4 keeps its launch_t4_pdl SIMT pair and the Q4 K-split MMA band is 7..12. Also here: the row-block experiment header is removed (it has no caller), and the production comments and the tests' descriptions state each shape's actual mechanism, the W=4 candidate's result without a causal claim, the Snapshot/Record workspace split, and each boundary number's timing boundary. The research notes that describe the same material ship with the review bundle rather than as a file in the repository. The fusion and batch candidates that were built and measured along the way (fused T=4, fused W=7/8, batched fused at aggregate 8) all lost; none of them is in this change. --- .../q4_q5/q4_q5_attn_input_small_t.cu | 50 ++-- .../q4_q5/q4_q5_gdn_input_independent.cu | 62 ++--- .../q5/q5_rowsplit_rowblock_small_t.cuh | 228 ------------------ tests/ops/test_attn_input_proj.cpp | 5 +- tests/ops/test_gdn_input_proj.cpp | 14 +- tests/ops/test_gdn_input_proj_conv_record.cpp | 10 +- .../ops/test_gdn_input_proj_conv_snapshot.cpp | 84 ++++++- 7 files changed, 140 insertions(+), 313 deletions(-) delete mode 100644 src/ops/linear/q5/q5_rowsplit_rowblock_small_t.cuh diff --git a/src/ops/attn_input_proj/q4_q5/q4_q5_attn_input_small_t.cu b/src/ops/attn_input_proj/q4_q5/q4_q5_attn_input_small_t.cu index 4ef2fc06f8..77a905577d 100644 --- a/src/ops/attn_input_proj/q4_q5/q4_q5_attn_input_small_t.cu +++ b/src/ops/attn_input_proj/q4_q5/q4_q5_attn_input_small_t.cu @@ -9,7 +9,6 @@ #include "ops/linear/q4/q4_rowsplit_gemv.cuh" #include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh" #include "ops/linear/q5/q5_rowsplit_gemv.cuh" -#include "ops/linear/q5/q5_rowsplit_rowblock_small_t.cuh" #include @@ -196,8 +195,17 @@ void launch_q5_split4_exact(const Tensor& x, const Weight& weight, Tensor& gate, case 6: launch_q5_split4<6>(x, weight, gate, value, stream); return; + case 7: + launch_q5_split4<7>(x, weight, gate, value, stream); + return; + case 8: + launch_q5_split4<8>(x, weight, gate, value, stream); + return; + case 9: + launch_q5_split4<9>(x, weight, gate, value, stream); + return; default: - throw std::invalid_argument("attention Q5 split4 requires T in [2,6]"); + throw std::invalid_argument("attention Q5 split4 requires T in [2,9]"); } } @@ -220,45 +228,25 @@ void launch_q5_simt(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& CUDA_CHECK(cudaGetLastError()); } -void launch_q5_rowblock(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value, - cudaStream_t stream) { - constexpr int kColsPerTile = 8; - constexpr int kRowsPerBlock = 8; - constexpr int kStages = 2; - constexpr int kThreads = kRowsPerBlock * 32; - const std::int32_t cols = x.ne[1]; - const dim3 grid(static_cast(div_up(kParentRows, kRowsPerBlock)), - static_cast(div_up(cols, kColsPerTile)), 1u); - q5_rowsplit_rowblock_small_t_kernel - <<>>( - static_cast(x.data), - static_cast(weight.qdata), - static_cast(weight.qhigh), - static_cast(weight.scales), - static_cast<__nv_bfloat16*>(gate.data), static_cast<__nv_bfloat16*>(value.data), - kParentRows, gate.ne[0], kHidden, cols, weight.padded_shape[1], kHidden / 1024); - CUDA_CHECK(cudaGetLastError()); -} - void launch_q5(const Tensor& x, const Weight& weight, Tensor& gate, Tensor& value, cudaStream_t stream) { if (x.ne[1] == 1) { launch_q5_gemv(x, weight, gate, value, stream); return; } - if (x.ne[1] <= 6) { + if (x.ne[1] <= 9) { + // Split4: one CTA owns one output row, its four warps split the K dimension and reduce + // their partial sums through shared memory, with the column count as a compile-time + // template argument so the kernel covers exactly the live columns. The fused projections + // use the same shape up to 6, so the Q5 parent has one mechanism from 2 to 9. launch_q5_split4_exact(x, weight, gate, value, stream); return; } - if (x.ne[1] <= 8) { - // See the GDN sibling: at T=7/8 this side is bound by repeated activation loads, and staging - // the activation slab per block measured 42.2 us against 64.8 us at T=7 - the best of the - // candidates tried here. - launch_q5_rowblock(x, weight, gate, value, stream); - return; - } if (x.ne[1] <= 12) { + // c4 SIMT: one output row per warp, up to four columns per column tile, with the quantized + // weight planes staged in shared memory and activations read from the input tensor. + // Retained in this interval on complete-Op measurements; both shapes are legal at every + // count in [2,12]. launch_q5_simt<4>(x, weight, gate, value, stream); return; } diff --git a/src/ops/gdn_input_proj/q4_q5/q4_q5_gdn_input_independent.cu b/src/ops/gdn_input_proj/q4_q5/q4_q5_gdn_input_independent.cu index 0e3ffbf17f..a807feba73 100644 --- a/src/ops/gdn_input_proj/q4_q5/q4_q5_gdn_input_independent.cu +++ b/src/ops/gdn_input_proj/q4_q5/q4_q5_gdn_input_independent.cu @@ -10,7 +10,6 @@ #include "ops/linear/q4/q4_rowsplit_gemv.cuh" #include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh" #include "ops/linear/q5/q5_rowsplit_gemv.cuh" -#include "ops/linear/q5/q5_rowsplit_rowblock_small_t.cuh" #include @@ -147,28 +146,6 @@ void launch_q4(const Tensor& x, const Weight& weight, Tensor& out, cudaStream_t } } -void launch_q5_rowblock(const Tensor& x, const Weight& weight, Tensor& value, Tensor& z, - cudaStream_t stream) { - constexpr int kColsPerTile = 8; - constexpr int kRowsPerBlock = 8; - constexpr int kStages = 2; - constexpr int kThreads = kRowsPerBlock * 32; - const std::int32_t cols = x.ne[1]; - const std::int32_t out_ld = static_cast(value.nb[1] / sizeof(__nv_bfloat16)); - const dim3 grid(static_cast(div_up(kValueZRows, kRowsPerBlock)), - static_cast(div_up(cols, kColsPerTile)), 1u); - q5_rowsplit_rowblock_small_t_kernel - <<>>( - static_cast(x.data), - static_cast(weight.qdata), - static_cast(weight.qhigh), - static_cast(weight.scales), - static_cast<__nv_bfloat16*>(value.data), static_cast<__nv_bfloat16*>(z.data), - kValueZRows, out_ld, kHidden, cols, weight.padded_shape[1], kHidden / 1024); - CUDA_CHECK(cudaGetLastError()); -} - void launch_q5_gemv(const Tensor& x, const Weight& weight, Tensor& value, Tensor& z, cudaStream_t stream) { constexpr int kRowsPerBlock = 16; @@ -218,8 +195,20 @@ void launch_q5_split4_exact(const Tensor& x, const Weight& weight, Tensor& value case 6: launch_q5_split4<6>(x, weight, value, z, stream); return; + case 7: + launch_q5_split4<7>(x, weight, value, z, stream); + return; + case 8: + launch_q5_split4<8>(x, weight, value, z, stream); + return; + case 9: + launch_q5_split4<9>(x, weight, value, z, stream); + return; + case 10: + launch_q5_split4<10>(x, weight, value, z, stream); + return; default: - throw std::invalid_argument("GDN Q5 split4 requires T in [2,6]"); + throw std::invalid_argument("GDN Q5 split4 requires T in [2,10]"); } } @@ -249,25 +238,20 @@ void launch_q5(const Tensor& x, const Weight& weight, Tensor& value, Tensor& z, launch_q5_gemv(x, weight, value, z, stream); return; } - if (x.ne[1] <= 6) { + if (x.ne[1] <= 10) { + // Split4: one CTA owns one output row, its four warps split the K dimension and reduce + // their partial sums through shared memory. The column count is a compile-time template + // argument, so the kernel covers exactly the live columns. The fused projections below + // (T=2..6) use the same shape, so the Q5 parent has one mechanism from 2 to 10; the band + // end is the measured crossover against the c4 tile at 11..12, not a shape limit. launch_q5_split4_exact(x, weight, value, z, stream); return; } - if (x.ne[1] <= 8) { - // T=7/8: the row-block kernel stages one 1024-value activation slab per block in shared - // memory and lets all kRowsPerBlock warps read it, so the activation traffic drops by - // kRowsPerBlock. At T=8 that moved this side from 93.4 us (row-split SIMT, activation bound - // by repeated activation traffic) to 62.7 us. - launch_q5_rowblock(x, weight, value, z, stream); - return; - } if (x.ne[1] <= 12) { - // T=9..12: this side stays on the narrow-column SIMT tile while the Q4 parent moves to the - // K-split, which is what makes the split form faster than the grouped kernel here. Measured - // as a complete op at T=9..12, a 4-column tile gives 101.6-105.7 us against 163.1-165.1 us for - // an 8-column tile and 109.8-126.7 us for the row-block shape. The 4-column tile is the - // fastest of the three that were measured; the measurement does not single out one cause for - // the difference between them. + // c4 SIMT: one output row per warp, up to four columns per column tile, with the quantized + // weight planes staged in shared memory and activations read from the input tensor. + // Retained in this interval on complete-Op measurements; both shapes are legal at every + // count in [2,15]. launch_q5_simt_cols<4>(x, weight, value, z, stream); return; } diff --git a/src/ops/linear/q5/q5_rowsplit_rowblock_small_t.cuh b/src/ops/linear/q5/q5_rowsplit_rowblock_small_t.cuh deleted file mode 100644 index f41805573d..0000000000 --- a/src/ops/linear/q5/q5_rowsplit_rowblock_small_t.cuh +++ /dev/null @@ -1,228 +0,0 @@ -#pragma once - -// Row-blocked Q5 small-T kernel for the fused Q4/Q5 input projections. -// -// Why this exists. The warp-per-row small-T kernel next door streams the weights of one output row -// per warp and reads the activations straight from global (L2) memory. That is the right shape when -// the weights dominate, and at T=1 they do. From about T=4 upward it stops being true: a warp-per-row -// block moves `rows * T * K * 2` bytes of activation for one K pass while the weights are only -// `rows * K * 5/8`, so at T=8 the activation traffic is 25x the weight traffic in bytes. Fitting the -// measured small-T times gives a *logical* effective activation-load bandwidth of about 12.7 TB/s -// for the GDN Q5 parent (12288x5120) -- a derived figure from bytes over time, not a counter -// reading, and L1/L2 hits mean it is not the throughput of one physical interface. Every row -// re-reads the same activation tile, so the fix is not a smaller tile but reuse: stage one -// activation slab in shared memory per block and let all kRowsPerBlock warps read it. That divides -// the repeated activation loads by kRowsPerBlock, and at T=7/8 this shape measured the fastest of -// the candidates tried (62.7 us at T=8 against 93.4 us for the row-split SIMT). -// -// Structure (deliberately the sibling kernel's, plus the staged slab): -// - one warp owns one output row; blockIdx.y selects a tile of kTt activation columns; -// - the weights of every row in the block are staged through shared memory with a cp.async -// pipeline, exactly as the warp-per-row kernel stages them; -// - one 1024-value slab of the activation tile is staged in shared memory by the whole block and -// read by every warp for every row, so the global activation reads are amortized over -// kRowsPerBlock rows; -// - dequantization is the same fp16-mantissa trick as the sibling kernel, and the group scale is -// applied per 8-value chunk; -// - every thread reaches every barrier: a row outside the grid stores nothing but still joins the -// pipeline, so the block-wide barriers are uniform. -// -// Tails: full 1024-value slabs require k % 8 == 0 (16-byte aligned activation columns) and cover -// [0, k/1024*1024); the remaining groups use the sibling kernel's scalar per-pair path reading global -// memory directly, masked at the k boundary. Weights in [k, padded_k) are never used. - -#include "ops/linear/q5/q5_rowsplit_gemm_simt.cuh" - -#include - -namespace ninfer::ops::detail { - -// Store epilogue for the fused input projections: two output tensors with their own row strides and -// a compile-time seam, plus the live column count (a kTt-wide tile may hold fewer valid columns). -// `col0` is the tile's first column in the output; the tile's columns are not necessarily the whole -// token range, so the output column is col0 + token. -struct Q5RowBlockStoreEpilogue { - template - __device__ __forceinline__ void - operator()(__nv_bfloat16* out, __nv_bfloat16* out_tail, std::int32_t n, std::int32_t out_ld, - std::int32_t row, std::int32_t col0, std::int32_t ncols, - const float (&values)[Tokens]) const { -#pragma unroll - for (int token = 0; token < Tokens; ++token) { - if (token >= ncols) { continue; } - const std::int64_t col = col0 + token; - if constexpr (SplitOutput) { - if (row < SplitRow) { - out[col * out_ld + row] = __float2bfloat16(values[token]); - } else { - out_tail[col * (n - SplitRow) + row - SplitRow] = - __float2bfloat16(values[token]); - } - } else { - out[col * out_ld + row] = __float2bfloat16(values[token]); - } - } - } -}; - -template -__launch_bounds__(kRowsPerBlock * 32) __global__ - void q5_rowsplit_rowblock_small_t_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, __nv_bfloat16* __restrict__ out_tail, std::int32_t n, - std::int32_t out_ld, std::int32_t k, std::int32_t t, std::int32_t padded_k, - std::int32_t full_slabs, Epilogue epilogue = {}) { - using Codec = typename SC::Codec; - constexpr int kPrefetch = kStages - 1; - constexpr int kHighU4Alloc = SC::kHighU4 > 0 ? SC::kHighU4 : 1; - constexpr int kSlabK = 1024; - constexpr int kVecsPerCol = kSlabK / 8; - constexpr int kThreads = kRowsPerBlock * 32; - static_assert(kTt >= 1 && kTt <= 32, "row-block small-T requires a bounded column tile"); - static_assert(kRowsPerBlock >= 1 && kRowsPerBlock <= 32, "row-block small-T block size"); - static_assert(kStages >= 2, "row-block small-T requires a cp.async pipeline"); - static_assert(!SplitOutput || SplitRow > 0, - "split-output row-block small-T requires a positive compile-time seam"); - - __shared__ __align__(16) __nv_bfloat16 x_sh[kTt * kSlabK]; - __shared__ __align__(16) uint4 s_nib[kRowsPerBlock][kStages][SC::kNibU4]; - __shared__ __align__(16) uint4 s_hi[kRowsPerBlock][kStages][kHighU4Alloc]; - __shared__ __align__(16) std::uint32_t s_sc[kRowsPerBlock][kStages][SC::kScaleU32]; - - const int lane = static_cast(threadIdx.x) & 31; - const int warp = static_cast(threadIdx.x) >> 5; - const int row = static_cast(blockIdx.x) * kRowsPerBlock + warp; - const int col0 = static_cast(blockIdx.y) * kTt; - const int ncols = t - col0 < kTt ? t - col0 : kTt; - // A row outside the grid still joins every barrier; only its loads and stores are dropped. - const bool live = row < n && ncols > 0; - const int safe_row = live ? row : 0; - - const int kg_padded = padded_k / Codec::kGroupK; - const std::uint8_t* code_row = codes + static_cast(safe_row) * kg_padded * 32; - const std::uint8_t* high_row = - SC::kHighBytesPerGroup > 0 - ? high + static_cast(safe_row) * kg_padded * SC::kHighBytesPerGroup - : nullptr; - const std::uint8_t* scale_row = scales + static_cast(safe_row) * kg_padded * 2; - const __nv_bfloat16* x0 = x + static_cast(col0) * k; - - float acc[kTt]; -#pragma unroll - for (int i = 0; i < kTt; ++i) { acc[i] = 0.0f; } - - const auto issue_weights = [&](int slab) { - if (live) { - q5_simt_issue_slab(s_nib[warp][slab % kStages], s_hi[warp][slab % kStages], - s_sc[warp][slab % kStages], code_row, high_row, scale_row, slab, - lane); - } else { - pipe_commit(); - } - }; - -#pragma unroll - for (int p = 0; p < kPrefetch; ++p) { - if (p < full_slabs) { - issue_weights(p); - } else { - pipe_commit(); - } - } - -#pragma unroll 1 - for (int s = 0; s < full_slabs; ++s) { - const int fetch = s + kPrefetch; - if (fetch < full_slabs) { - issue_weights(fetch); - } else { - pipe_commit(); - } - - // Stage this slab's activation tile once for the whole block. The barrier before it guarantees - // no warp is still reading the previous slab; the barrier after it publishes the stores. - __syncthreads(); - { - const auto* src = reinterpret_cast(x0); - auto* dst = reinterpret_cast(x_sh); - const int limit = ncols * kVecsPerCol; - for (int i = static_cast(threadIdx.x); i < limit; i += kThreads) { - const int col = i / kVecsPerCol; - const int vec = i - col * kVecsPerCol; - dst[i] = src[static_cast(col) * (k / 8) + s * kVecsPerCol + vec]; - } - } - pipe_wait(); - __syncthreads(); - -#pragma unroll - for (int c = 0; c < 4; ++c) { - float w[8]; - SC::dequant_chunk(s_nib[warp][s % kStages], s_hi[warp][s % kStages], - s_sc[warp][s % kStages], c, lane, w); - const std::int64_t xoff = static_cast(c) * 256 + lane * 8; -#pragma unroll - for (int tt = 0; tt < kTt; ++tt) { - if (tt < ncols) { - const uint4 xv = *reinterpret_cast(&x_sh[tt * kSlabK + xoff]); - const float2 f0 = bf16x2_bits_to_float2(xv.x); - const float2 f1 = bf16x2_bits_to_float2(xv.y); - const float2 f2 = bf16x2_bits_to_float2(xv.z); - const float2 f3 = bf16x2_bits_to_float2(xv.w); - acc[tt] = fmaf(w[0], f0.x, acc[tt]); - acc[tt] = fmaf(w[1], f0.y, acc[tt]); - acc[tt] = fmaf(w[2], f1.x, acc[tt]); - acc[tt] = fmaf(w[3], f1.y, acc[tt]); - acc[tt] = fmaf(w[4], f2.x, acc[tt]); - acc[tt] = fmaf(w[5], f2.y, acc[tt]); - acc[tt] = fmaf(w[6], f3.x, acc[tt]); - acc[tt] = fmaf(w[7], f3.y, acc[tt]); - } - } - } - // The next iteration's issue rewrites the slot this slab was read from, and a 16-byte - // cp.async copy issued by one lane fills words another lane reads, so the warp has to be - // converged before that copy is issued. The warp-per-row kernel carries the same fence. - __syncwarp(); - } - - // Scalar tail: remaining groups read global memory directly, masked at k. Identical to the - // warp-per-row kernel's tail, so a shape with k % 8 != 0 takes the same path in both. - const int g0 = (full_slabs * kSlabK) / Codec::kGroupK; - const int kg_used = div_up(k, Codec::kGroupK); - if (live) { - for (int g = g0; g < kg_used; ++g) { - const int kk = g * Codec::kGroupK + lane * 2; - if (kk >= k) { continue; } - float w0 = 0.0f; - float w1 = 0.0f; - Codec::load_pair(codes, high, scales, - static_cast(row) * kg_padded + g, lane, w0, w1); -#pragma unroll - for (int tt = 0; tt < kTt; ++tt) { - if (tt < ncols) { - const std::int64_t xb = static_cast(tt) * k + kk; - acc[tt] = fmaf(w0, __bfloat162float(x0[xb]), acc[tt]); - if (kk + 1 < k) { acc[tt] = fmaf(w1, __bfloat162float(x0[xb + 1]), acc[tt]); } - } - } - } - } - - if (!live) { return; } - float values[kTt]; -#pragma unroll - for (int tt = 0; tt < kTt; ++tt) { - float a = acc[tt]; - a = warp_reduce_sum(a); - values[tt] = a; - } - if (lane == 0) { - epilogue.template operator()(out, out_tail, n, out_ld, row, col0, - ncols, values); - } -} - -} // namespace ninfer::ops::detail diff --git a/tests/ops/test_attn_input_proj.cpp b/tests/ops/test_attn_input_proj.cpp index 2911f8d558..bfa967bbbe 100644 --- a/tests/ops/test_attn_input_proj.cpp +++ b/tests/ops/test_attn_input_proj.cpp @@ -167,7 +167,10 @@ int run_q4_q5() { for (int t : {129, 144, 145, 160, 161, 192, 193, 256, 257, 1024}) failures += run_target_projection_case(query_key, &gate_value, t, ops::LinearPolicy::A16Only); - for (int t : {1, 8, 12, 13, 16, 32, 63, 64, 65, 96, 104, 105, 127, 128, 129, 192, 193}) + // The replayed set covers the Q5 split4 band's new counts (7 and 9) next to the ones already + // there, so the instances this change re-routes are replayed with a re-poisoned output and a + // changed activation at the captured address. + for (int t : {1, 7, 8, 9, 12, 13, 16, 32, 63, 64, 65, 96, 104, 105, 127, 128, 129, 192, 193}) failures += run_target_projection_case(query_key, &gate_value, t, ops::LinearPolicy::A16Only, true); return failures; diff --git a/tests/ops/test_gdn_input_proj.cpp b/tests/ops/test_gdn_input_proj.cpp index d424451688..b5167a3c53 100644 --- a/tests/ops/test_gdn_input_proj.cpp +++ b/tests/ops/test_gdn_input_proj.cpp @@ -156,15 +156,17 @@ int run_q4_q5() { DevicePackedWeight value_z_weight( quantized_weight::make_patterned_weight(QType::Q5_G64_FP16, 12288, kHidden, 419U)); int failures = 0; - // Every route boundary and both of its neighbours: the Q4/Q5 column catalog now hands 1..12 to - // the per-side small-column kernels (the K-split Q4 parent with the row-block Q5 side at 7/8 and - // the c4 SIMT Q5 side at 9..12), 13..32 to the 32x32 tile, 33..64 to the 32x64 tile, and 65 - // upward to the 64x128 tile that also supplies the 128-column tail slices. + // Every route boundary and both of its neighbours: the Q4/Q5 column catalog hands 1..12 to the + // per-side small-column kernels (the K-split Q4 parent with the split4 Q5 side at 2..10 and the c4 + // SIMT Q5 side at 11..12), 13..32 to the 32x32 tile, 33..64 to the 32x64 tile, and 65 upward to the + // 64x128 tile that also supplies the 128-column tail slices. The Q4 K-split band (7..12) is not + // changed by this work; its two ends are covered by the same list. for (const std::int32_t tokens : {1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 15, 16, 17, 32, 33, 64, 65, 127, 128, 129, 193}) { failures += run_q4_q5_case(query_key, value_z_weight, tokens); } - // One captured replay per route, including a 128-column tail slice (129 = 128 + 1). - for (const std::int32_t tokens : {7, 8, 9, 12, 13, 33, 65, 129}) { + // One captured replay per route, including both ends of the split4 band and a 128-column tail + // slice (129 = 128 + 1). + for (const std::int32_t tokens : {7, 8, 9, 10, 12, 13, 33, 65, 129}) { failures += run_q4_q5_graph_case(query_key, value_z_weight, tokens); } return failures; diff --git a/tests/ops/test_gdn_input_proj_conv_record.cpp b/tests/ops/test_gdn_input_proj_conv_record.cpp index 7510311bde..59999f1a06 100644 --- a/tests/ops/test_gdn_input_proj_conv_record.cpp +++ b/tests/ops/test_gdn_input_proj_conv_record.cpp @@ -173,7 +173,9 @@ int run_case(std::string_view label, std::int32_t hidden, std::int32_t value_row launch(); CUDA_CHECK(cudaStreamSynchronize(stream)); auto activation_bits = bf16_bits(activation); - if (width == 2 || width == 9 || width == 16) { + if (width == 2 || width == 9 || width == 10 || width == 16) { + // 2 is the bottom of the Q5 parent's split4 band, 9 the count that used to be its top, 10 the + // count that is its top now, and 16 a grouped representative extent. cudaGraph_t graph; cudaGraphExec_t executable; CUDA_CHECK(cudaStreamBeginCapture(stream, cudaStreamCaptureModeGlobal)); @@ -293,6 +295,12 @@ int run_q4_q5() { } failures += run(5, 3, {5, 3, 1}, 1491U); failures += run(4, 4, {4, 3, 2, 1}, 1492U); + // R7-review trial: the batched organisation whose aggregate column count is 8, which is inside the + // range the fused template covers when the request is read from the flattened token axis. The dense + // case checks the history reload at the request boundary, and the masked one additionally puts an + // invalid tail in the second request, so the boundary is checked next to a zeroed column. + failures += run(4, 2, {}, 1493U); + failures += run(4, 2, {4, 2}, 1494U); failures += qk.verify_preserved("Q4 record qk weight"); failures += value_z.verify_preserved("Q5 record value/z weight"); return failures; diff --git a/tests/ops/test_gdn_input_proj_conv_snapshot.cpp b/tests/ops/test_gdn_input_proj_conv_snapshot.cpp index 1d5cfac082..5ef71f26ae 100644 --- a/tests/ops/test_gdn_input_proj_conv_snapshot.cpp +++ b/tests/ops/test_gdn_input_proj_conv_snapshot.cpp @@ -555,13 +555,12 @@ int run_q4_q5() { const std::int32_t initial_slot = tokens == 5 ? 0 : tokens + 1; failures += run_q4_q5_case(query_key, value_z_weight, tokens, initial_slot); } - // The column counts that reach a projection mechanism this change touched: W=7 is the first - // extent with a tail column, W=8 the first full row-block extent on the Q5 side, W=9 the first - // extent of the 9..12 band (K-split Q4 parent with the narrow c4 SIMT Q5 tile, which is a - // different Q5 mechanism from the row-block used at 7/8). These are compared against the complete - // FP64 formula for every output row and every state channel, so a mechanism that is wrong away - // from the sampled rows cannot pass. - for (const std::int32_t tokens : {7, 8, 9}) { + // The column counts that reach a projection mechanism this change touched: 7 and 8 are the first + // counts the Q5 parent's split4 shape takes over (they were the row-block band, and the first two + // counts of the narrow-SIMT band), 9 the count that used to be the top of the split4 band, and 10 + // the top of it now. These are compared against the complete FP64 formula for every output row and + // every state channel, so a mechanism that is wrong away from the sampled rows cannot pass. + for (const std::int32_t tokens : {7, 8, 9, 10}) { failures += run_q4_q5_case(query_key, value_z_weight, tokens, tokens + 1, true); } // The rest of the admitted range keeps the sampled comparison: the same formula, a smaller @@ -600,7 +599,78 @@ int run_q4_q5() { state, valid, initial, snapshot_base, q, k, v, z, workspace, nullptr); }); + // The batched organisations whose aggregate column count is 10 reach the same Q5 parent mechanism + // as the dense B=1 W=10 case above, through the seam of a multi-request call: B=2/W=5 here, which + // covers the per-request history split and the aggregate route. Two forms of the same aggregate + // count are covered because the per-side band end moved, and the aggregate count is what selects + // the route. + constexpr std::int32_t kAggTenWidth = 5; + constexpr std::int32_t kAggTenBatch = 2; + const std::size_t agg_ten_workspace_bytes = + ops::gdn_input_proj_conv_snapshot_workspace_capacity_bytes(kQueryRows, kKeyRows, kValueRows, + kAggTenBatch, kAggTenWidth, + kAggTenWidth); + failures += run_batched_case( + "Q4/Q5 A16 B=2 W=5", kHidden, kValueRows, kZRows, kAggTenWidth, kAggTenBatch, {}, + conv_weight, agg_ten_workspace_bytes, kGdnInputProjConvSnapshotA16Tolerance, + [&](std::int32_t row, std::int32_t flat_column, const std::vector& activation) { + const float* column = + activation.data() + static_cast(flat_column) * kHidden; + if (row < kQueryRows + kKeyRows) { + return quantized_weight::dot_fp64(query_key.host, row, column, kHidden); + } + return quantized_weight::dot_fp64(value_z_weight.host, row - kQueryRows - kKeyRows, + column, kHidden); + }, + [&](std::int32_t row, std::int32_t flat_column, const std::vector& activation) { + return quantized_weight::dot_fp64( + value_z_weight.host, kValueRows + row, + activation.data() + static_cast(flat_column) * kHidden, kHidden); + }, + [&](const Tensor& x, const Tensor& conv, Tensor& state, const Tensor& valid, + const Tensor& initial, const Tensor& snapshot_base, Tensor& q, Tensor& k, Tensor& v, + Tensor& z, WorkspaceArena& workspace) { + ops::gdn_input_proj_conv_snapshot(x, query_key.view(), value_z_weight.view(), conv, + state, valid, initial, snapshot_base, q, k, v, z, + workspace, nullptr); + }); + // One batched case with mixed valid prefixes whose aggregate column count is 8 (4 x 2): two live + // prefixes (4 and 2 of 4) over the same projection route, plus the exact-zero invalid tails and the + // per-request history split. The projection is routed once from the aggregate count, not per + // request, so this is a batched organisation of the aggregate-8 configuration and not a claim + // about where any single-side band ends. + constexpr std::int32_t kNarrowWidth = 4; + constexpr std::int32_t kNarrowBatch = 2; + const std::vector narrow_valid{4, 2}; + const std::size_t narrow_workspace_bytes = + ops::gdn_input_proj_conv_snapshot_workspace_capacity_bytes(kQueryRows, kKeyRows, kValueRows, + kNarrowBatch, kNarrowWidth, + kNarrowWidth); + failures += run_batched_case( + "Q4/Q5 A16 B=2 W=4 masked", kHidden, kValueRows, kZRows, kNarrowWidth, kNarrowBatch, + narrow_valid, conv_weight, narrow_workspace_bytes, kGdnInputProjConvSnapshotA16Tolerance, + [&](std::int32_t row, std::int32_t flat_column, const std::vector& activation) { + const float* column = + activation.data() + static_cast(flat_column) * kHidden; + if (row < kQueryRows + kKeyRows) { + return quantized_weight::dot_fp64(query_key.host, row, column, kHidden); + } + return quantized_weight::dot_fp64(value_z_weight.host, row - kQueryRows - kKeyRows, + column, kHidden); + }, + [&](std::int32_t row, std::int32_t flat_column, const std::vector& activation) { + return quantized_weight::dot_fp64( + value_z_weight.host, kValueRows + row, + activation.data() + static_cast(flat_column) * kHidden, kHidden); + }, + [&](const Tensor& x, const Tensor& conv, Tensor& state, const Tensor& valid, + const Tensor& initial, const Tensor& snapshot_base, Tensor& q, Tensor& k, Tensor& v, + Tensor& z, WorkspaceArena& workspace) { + ops::gdn_input_proj_conv_snapshot(x, query_key.view(), value_z_weight.view(), conv, + state, valid, initial, snapshot_base, q, k, v, z, + workspace, nullptr); + }); // One batched case with mixed valid prefixes whose aggregate column count is 48 (16 x 3). The // projection is routed once for the whole call from that aggregate column count - so this case // exercises the mixed valid prefixes and the exact-zero invalid tails, not three different