diff --git a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu index 0d41ca96a3..77f911cb40 100644 --- a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu +++ b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cu @@ -30,6 +30,16 @@ template void launch_route(const Tensor& x, const Weight& w, Tensor& out, cudaStream_t stream) { const bool full = (w.n % Schedule::BM) == 0 && (x.ne[1] % Schedule::BN) == 0 && w.k == w.padded_shape[1] && (w.k % Schedule::BK) == 0; + // cp.async.cg requires a 16 B-aligned source. On the predicated path the activation source + // steps by 8 bf16 and the scale source by kg = padded_k / 32 two-byte scales, so the encoding + // is legal only when both are multiples of 8. The full path inherits this from (w.k % BK) == 0 + // above; a schedule that asks for cg on the predicated path has to be checked here. + if constexpr (Schedule::kPredicatedCache == Cache::cg) { + if (!full && ((w.k % 8) != 0 || ((w.padded_shape[1] / 32) % 8) != 0)) { + throw std::invalid_argument( + "w8 rowsplit mma: predicated cg needs k and the group count to be multiples of 8"); + } + } for_each_token_slice(x.ne[1], Schedule::BN, [&](std::int32_t offset, std::int32_t count) { const Tensor x_slice = x.slice(1, offset, count); Tensor out_slice = out.slice(1, offset, count); @@ -77,7 +87,10 @@ using MmaR48C112 = W8RowSplitMmaGemmSchedule<48, 112, 48, 16, 2>; using MmaR48C128 = W8RowSplitMmaGemmSchedule<48, 128, 48, 16, 2>; using MmaR64C96 = W8RowSplitMmaGemmSchedule<64, 96, 64, 16, 2>; using MmaR64C112 = W8RowSplitMmaGemmSchedule<64, 112, 64, 16, 2>; -using MmaR64C128 = W8RowSplitMmaGemmSchedule<64, 128, 64, 16, 2, 2>; +// The wide route of the linear family is the one schedule measured to want cg on the predicated +// path; every other schedule keeps the inherited ca. +using MmaR64C128 = + W8RowSplitMmaGemmSchedule<64, 128, 64, 16, 2, 2>::with_predicated_cache; using MmaR96C96 = W8RowSplitMmaGemmSchedule<96, 96, 48, 16, 2>; using MmaR128C64 = W8RowSplitMmaGemmSchedule<128, 64, 64, 16, 2>; using MmaR128C80 = W8RowSplitMmaGemmSchedule<128, 80, 64, 16, 2>; diff --git a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh index 13e555e1b1..154c275f6f 100644 --- a/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh +++ b/src/ops/linear/w8/w8_rowsplit_gemm_mma.cuh @@ -9,6 +9,7 @@ // m16n8k16 BF16 MMA with FP32 accumulation. #include "ops/common/mma.cuh" +#include "ops/common/memory.cuh" #include "ops/common/math.cuh" #include "ops/linear/w8/w8_rowsplit_output.cuh" @@ -26,8 +27,21 @@ union alignas(16) W8Bf16x8Bits { static_assert(sizeof(W8Bf16x8Bits) == 16); +// The predicated loads below inherited Cache::ca from cp_async_zfill's default, while the full path +// a few lines down spells cg. This parameter makes that a choice. It defaults to ca, so adding it +// changes no instantiation, and it governs the predicated branch only - the full branch keeps its +// own cg - which is why it is named for it. +// +// ca stays the default because whether cg pays is not a property of the schedule. On the two +// BM = 16 schedules in this tree - src/ops/linear/w8/w8_feature.cu and +// src/ops/attn_input_proj/w8/w8_dflash2_attn_input.cu - m / BM CTAs read the same activation tile +// (m is the weight-row count, the kernel's first size argument) and four of them share an SM; L1 +// serves 35.6 % of their sectors, and forcing cg there costs +16.8 % and +21.0 % of the operation. +// Most other schedules gain from cg, but not all, and not on every shape: MmaR64x16C48K128A1, a +// BM = 64 tile, gains 3.3 % on [34816, 5120] and loses 1.1 % on [248320, 5120]. So the policy is +// set per schedule and only where it has been measured across the shapes that reach it. template + int ACTIVATION_STAGES_ = STAGES_, Cache PredicatedCache_ = Cache::ca> struct W8RowSplitMmaGemmSchedule { static constexpr int BM = BM_; static constexpr int BN = BN_; @@ -48,6 +62,16 @@ struct W8RowSplitMmaGemmSchedule { static constexpr int SMEM_BYTES = BM * BK * 2 + ACTIVATION_STAGES * BN * BK * 2 + BM * BK + BM * SCALE_CACHE_BYTES; + // Cache policy for the predicated loads only; see the note above the template. + static constexpr Cache kPredicatedCache = PredicatedCache_; + + // Restate this schedule with a different predicated policy, leaving every other parameter where + // it is rather than respelling it - and its default with it - at the point of use. + template + using with_predicated_cache = + W8RowSplitMmaGemmSchedule; + static_assert(BM % WM == 0 && BN % WN == 0); static_assert(WM % 16 == 0 && WN % 8 == 0); static_assert(THREADS <= 1024); @@ -132,7 +156,11 @@ __global__ __launch_bounds__(Cfg::THREADS, Cfg::MIN_BLOCKS) void w8_rowsplit_gem cp_async<16, Cache::cg>(dst, &x[static_cast(nn) * k + kk]); } else { const int valid = (nn < n && kk < k) ? min(8, k - kk) * 2 : 0; - ninfer::ops::cp_async_zfill<16>( + // cp.async.cg needs the source naturally aligned to 16 B. kk steps by 8 and x is + // bf16, so that holds exactly when k is a multiple of 8. The full path gets this + // from its own (k % BK) == 0 test; the predicated path has no such guarantee, so + // launch_route checks it for any schedule that asks for cg. + ninfer::ops::cp_async_zfill<16, Cfg::kPredicatedCache>( dst, &x[static_cast(nn < n ? nn : 0) * k + (kk < k ? kk : 0)], valid); } @@ -156,8 +184,8 @@ __global__ __launch_bounds__(Cfg::THREADS, Cfg::MIN_BLOCKS) void w8_rowsplit_gem } else { const bool valid_row = output_tile.valid(grow, m); const std::int64_t gi = static_cast(valid_row ? grow : 0) * kg + g0; - ninfer::ops::cp_async_zfill<16>(dst, &codes[gi * 32 + chunk * 16], - valid_row ? 16 : 0); + ninfer::ops::cp_async_zfill<16, Cfg::kPredicatedCache>( + dst, &codes[gi * 32 + chunk * 16], valid_row ? 16 : 0); } } if ((kt % SCALE_CACHE_TILES) == 0) { @@ -173,7 +201,10 @@ __global__ __launch_bounds__(Cfg::THREADS, Cfg::MIN_BLOCKS) void w8_rowsplit_gem const int valid_scales = valid_row && g0 < kg ? min(8, kg - g0) : 0; const std::int64_t gi = static_cast(valid_row ? grow : 0) * kg + min(g0, kg - 1); - ninfer::ops::cp_async_zfill<16>(dst, &scales[gi * 2], valid_scales * 2); + // gi steps by kg, and scales are 2 B, so 16 B alignment here needs kg to be + // a multiple of 8. launch_route checks that alongside k. + ninfer::ops::cp_async_zfill<16, Cfg::kPredicatedCache>(dst, &scales[gi * 2], + valid_scales * 2); } } }