| 1 | /* |
| 2 | * Modified by Neural Magic |
| 3 | * Copyright (C) Marlin.2024 Elias Frantar |
| 4 | * |
| 5 | * Licensed under the Apache License, Version 2.0 (the "License"); |
| 6 | * you may not use this file except in compliance with the License. |
| 7 | * You may obtain a copy of the License at |
| 8 | * |
| 9 | * http://www.apache.org/licenses/LICENSE-2.0 |
| 10 | * |
| 11 | * Unless required by applicable law or agreed to in writing, software |
| 12 | * distributed under the License is distributed on an "AS IS" BASIS, |
| 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 14 | * See the License for the specific language governing permissions and |
| 15 | * limitations under the License. |
| 16 | */ |
| 17 | |
| 18 | /* |
| 19 | * Adapted from https://github.com/IST-DASLab/marlin |
| 20 | */ |
| 21 | |
| 22 | #ifndef MARLIN_NAMESPACE_NAME |
| 23 | #define MARLIN_NAMESPACE_NAME marlin |
| 24 | #endif |
| 25 | |
| 26 | #include "marlin.cuh" |
| 27 | #include "marlin_dtypes.cuh" |
| 28 | #include "dequant.h" |
| 29 | #include "marlin_mma.h" |
| 30 | #include "core/scalar_type.hpp" |
| 31 | |
| 32 | #define STATIC_ASSERT_SCALAR_TYPE_VALID(scalar_t) \ |
| 33 | static_assert(std::is_same<scalar_t, half>::value || \ |
| 34 | std::is_same<scalar_t, nv_bfloat16>::value, \ |
| 35 | "only float16 and bfloat16 is supported"); |
| 36 | |
| 37 | namespace MARLIN_NAMESPACE_NAME { |
| 38 | |
| 39 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 750 |
| 40 | |
| 41 | template <typename scalar_t, // compute dtype, half or nv_float16 |
| 42 | const vllm::ScalarTypeId b_type_id, // weight MarlinScalarType id |
| 43 | const vllm::ScalarTypeId s_type_id, // weight scale ScalarType id |
| 44 | const int threads, // number of threads in a threadblock |
| 45 | const int thread_m_blocks, // number of 16x16 blocks in the m |
| 46 | // dimension (batchsize) of the |
| 47 | // threadblock |
| 48 | const int thread_n_blocks, // same for n dimension (output) |
| 49 | const int thread_k_blocks, // same for k dimension (reduction) |
| 50 | const bool m_block_size_8, // whether m_block_size == 8 |
| 51 | // only works when thread_m_blocks == 1 |
| 52 | const int stages, // number of stages for the async global->shared |
| 53 | // fetch pipeline |
| 54 | const int group_blocks, // number of consecutive 16x16 blocks |
| 55 | // with a separate quantization scale |
| 56 | const bool is_zp_float // is zero point of float16 type? |
| 57 | > |
| 58 | __global__ void Marlin( |
| 59 | const int4* __restrict__ A, // fp16 input matrix of shape mxk |
| 60 | const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn |
| 61 | int4* __restrict__ C, // fp16 output buffer of shape mxn |
| 62 | int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) |
| 63 | const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape |
| 64 | // (k/groupsize)xn |
| 65 | int prob_m, // batch dimension m |
| 66 | int prob_n, // output dimension n |
| 67 | int prob_k, // reduction dimension k |
| 68 | int* locks, // extra global storage for barrier synchronization |
| 69 | bool use_fp32_reduce // whether to use fp32 global reduce |
| 70 | ) {} |
| 71 | |
| 72 | } // namespace marlin |
| 73 | |
| 74 | #else |
| 75 | |
| 76 | // Instruction for loading a full 16x16 matrix fragment of operand A from shared |
| 77 | // memory, directly in tensor core layout. |
| 78 | template <int count, vllm::ScalarTypeId type_id> |
| 79 | __device__ inline void ldsm(typename MarlinScalarType<type_id>::FragA& frag_a, |
| 80 | const void* smem_ptr) { |
| 81 | uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a); |
| 82 | uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr)); |
| 83 | if constexpr (count == 4) { |
| 84 | asm volatile( |
| 85 | "ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n" |
| 86 | : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3]) |
| 87 | : "r"(smem)); |
| 88 | } else if constexpr (count == 2) { |
| 89 | asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0,%1}, [%2];\n" |
| 90 | : "=r"(a[0]), "=r"(a[1]) |
| 91 | : "r"(smem)); |
| 92 | } else if constexpr (count == 1) { |
| 93 | asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n" |
| 94 | : "=r"(a[0]) |
| 95 | : "r"(smem)); |
| 96 | } else { |
| 97 | static_assert(count == 1 || count == 2 || count == 4, "invalid count"); |
| 98 | } |
| 99 | } |
| 100 | |
| 101 | // Multiply dequantized values by the corresponding quantization scale; used |
| 102 | // only for grouped quantization. |
| 103 | template <vllm::ScalarTypeId type_id> |
| 104 | __device__ inline void scale(typename MarlinScalarType<type_id>::FragB& frag_b, |
| 105 | typename MarlinScalarType<type_id>::FragS& frag_s, |
| 106 | int i) { |
| 107 | using scalar_t = typename MarlinScalarType<type_id>::scalar_t; |
| 108 | using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2; |
| 109 | scalar_t2 s = MarlinScalarType<type_id>::num2num2( |
| 110 | reinterpret_cast<scalar_t*>(&frag_s)[i]); |
| 111 | frag_b[0] = __hmul2(frag_b[0], s); |
| 112 | frag_b[1] = __hmul2(frag_b[1], s); |
| 113 | } |
| 114 | |
| 115 | template <vllm::ScalarTypeId type_id> |
| 116 | __device__ inline void scale_and_sub( |
| 117 | typename MarlinScalarType<type_id>::FragB& frag_b, |
| 118 | typename MarlinScalarType<type_id>::scalar_t s, |
| 119 | typename MarlinScalarType<type_id>::scalar_t zp) { |
| 120 | using scalar_t = typename MarlinScalarType<type_id>::scalar_t; |
| 121 | using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2; |
| 122 | scalar_t2 s2 = MarlinScalarType<type_id>::num2num2(s); |
| 123 | scalar_t2 zp2 = MarlinScalarType<type_id>::num2num2(zp); |
| 124 | frag_b[0] = __hfma2(frag_b[0], s2, __hneg2(zp2)); |
| 125 | frag_b[1] = __hfma2(frag_b[1], s2, __hneg2(zp2)); |
| 126 | } |
| 127 | |
| 128 | template <vllm::ScalarTypeId type_id> |
| 129 | __device__ inline void sub_zp( |
| 130 | typename MarlinScalarType<type_id>::FragB& frag_b, |
| 131 | typename MarlinScalarType<type_id>::scalar_t2& frag_zp, int i) { |
| 132 | using scalar_t = typename MarlinScalarType<type_id>::scalar_t; |
| 133 | using scalar_t2 = typename MarlinScalarType<type_id>::scalar_t2; |
| 134 | scalar_t2 zp = MarlinScalarType<type_id>::num2num2( |
| 135 | reinterpret_cast<scalar_t*>(&frag_zp)[i]); |
| 136 | frag_b[0] = __hsub2(frag_b[0], zp); |
| 137 | frag_b[1] = __hsub2(frag_b[1], zp); |
| 138 | } |
| 139 | |
| 140 | // Given 2 floats multiply by 2 scales (halves) |
| 141 | template <vllm::ScalarTypeId type_id> |
| 142 | __device__ inline void scale_float( |
| 143 | float* c, typename MarlinScalarType<type_id>::FragS& s) { |
| 144 | using scalar_t = typename MarlinScalarType<type_id>::scalar_t; |
| 145 | scalar_t* s_ptr = reinterpret_cast<scalar_t*>(&s); |
| 146 | c[0] = __fmul_rn(c[0], MarlinScalarType<type_id>::num2float(s_ptr[0])); |
| 147 | c[1] = __fmul_rn(c[1], MarlinScalarType<type_id>::num2float(s_ptr[1])); |
| 148 | } |
| 149 | |
| 150 | // Wait until barrier reaches `count`, then lock for current threadblock. |
| 151 | __device__ inline void barrier_acquire(int* lock, int count) { |
| 152 | if (threadIdx.x == 0) { |
| 153 | int state = -1; |
| 154 | do |
| 155 | // Guarantee that subsequent writes by this threadblock will be visible |
| 156 | // globally. |
| 157 | asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" |
| 158 | : "=r"(state) |
| 159 | : "l"(lock)); |
| 160 | while (state != count); |
| 161 | } |
| 162 | __syncthreads(); |
| 163 | } |
| 164 | |
| 165 | // Release barrier and increment visitation count. |
| 166 | __device__ inline void barrier_release(int* lock, bool reset = false) { |
| 167 | __syncthreads(); |
| 168 | if (threadIdx.x == 0) { |
| 169 | if (reset) { |
| 170 | lock[0] = 0; |
| 171 | return; |
| 172 | } |
| 173 | int val = 1; |
| 174 | // Make sure that all writes since acquiring this barrier are visible |
| 175 | // globally, while releasing the barrier. |
| 176 | asm volatile("fence.acq_rel.gpu;\n"); |
| 177 | asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n" |
| 178 | : |
| 179 | : "l"(lock), "r"(val)); |
| 180 | } |
| 181 | } |
| 182 | |
| 183 | // Wait until value of lock to be negative, and then add 1 |
| 184 | __device__ inline void wait_negative_and_add(int* lock) { |
| 185 | if (threadIdx.x == 0) { |
| 186 | int state = 0; |
| 187 | do |
| 188 | // Guarantee that subsequent writes by this threadblock will be visible |
| 189 | // globally. |
| 190 | asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n" |
| 191 | : "=r"(state) |
| 192 | : "l"(lock)); |
| 193 | while (state >= 0); |
| 194 | atomicAdd(lock, 1); |
| 195 | } |
| 196 | __syncthreads(); |
| 197 | } |
| 198 | |
| 199 | template <const vllm::ScalarTypeId a_type_id, // A ScalarType id |
| 200 | const vllm::ScalarTypeId b_type_id, // B ScalarType id |
| 201 | const vllm::ScalarTypeId c_type_id, // C ScalarType id |
| 202 | const vllm::ScalarTypeId s_type_id, // B_SCALE ScalarType id |
| 203 | const int threads, // number of threads in a threadblock |
| 204 | const int thread_m_blocks, // number of 16x16 blocks in the m |
| 205 | // dimension (batchsize) of the |
| 206 | // threadblock |
| 207 | const int thread_n_blocks, // same for n dimension (output) |
| 208 | const int thread_k_blocks, // same for k dimension (reduction) |
| 209 | const bool m_block_size_8, // whether m_block_size == 8 |
| 210 | // only works when thread_m_blocks == 1 |
| 211 | const int stages, // number of stages for the async global->shared |
| 212 | // fetch pipeline |
| 213 | const int group_blocks, // number of consecutive 16x16 blocks |
| 214 | // with a separate quantization scale |
| 215 | const bool is_zp_float // is zero point of float16 type? |
| 216 | > |
| 217 | __global__ void Marlin( |
| 218 | const int4* __restrict__ A0, // fp16 input matrix of shape mxk |
| 219 | const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn |
| 220 | int4* __restrict__ C0, // fp16 output buffer of shape mxn |
| 221 | int4* __restrict__ C_tmp, // fp32 tmp output buffer (for reduce) |
| 222 | const int4* __restrict__ b_bias_ptr, |
| 223 | // float scales of input matrix, only used when is_a_8bit == true. |
| 224 | // shape (m,) |
| 225 | const float* __restrict__ a_scales_ptr, |
| 226 | // fp16 quantization scales. shape (k/groupsize, n) |
| 227 | const int4* __restrict__ scales_ptr, |
| 228 | // float global scale (for nvfp4// only) |
| 229 | const float* __restrict__ global_scale_ptr, |
| 230 | // 4bit packed zero-points of shape |
| 231 | // (k/groupsize, n/pack_factor) |
| 232 | const int4* __restrict__ zp_ptr, |
| 233 | int prob_m, // batch dimension m |
| 234 | int prob_n, // output dimension n |
| 235 | int prob_k, // reduction dimension k |
| 236 | int lda, // A.stride(0), equal to prob_k is A is contiguous |
| 237 | int* locks, // extra global storage for barrier synchronization |
| 238 | bool has_bias, |
| 239 | bool use_atomic_add, // whether to use atomic add to reduce |
| 240 | bool use_fp32_reduce, // whether to use fp32 global reduce |
| 241 | int max_shared_mem) { |
| 242 | // Each threadblock processes one "stripe" of the B matrix with (roughly) the |
| 243 | // same size, which might involve multiple column "slices" (of width 16 * |
| 244 | // `thread_n_blocks`). Stripes are defined as shown in the 3x3 matrix 5 SM |
| 245 | // example: |
| 246 | // 0 1 3 |
| 247 | // 0 2 3 |
| 248 | // 1 2 4 |
| 249 | // While this kind of partitioning makes things somewhat more complicated, it |
| 250 | // ensures good utilization of all SMs for many kinds of shape and GPU |
| 251 | // configurations, while requiring as few slow global cross-threadblock |
| 252 | // reductions as possible. |
| 253 | |
| 254 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 890 |
| 255 | // FP8 computation is only supported for Ada Lovelace or newer architectures. |
| 256 | if constexpr (a_type_id == vllm::kFE4M3fn.id()) return; |
| 257 | #endif |
| 258 | |
| 259 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 |
| 260 | // Turing TensorCore only supports fp16 and int8 |
| 261 | if constexpr (a_type_id != vllm::kFloat16.id() && a_type_id != vllm::kS8.id()) |
| 262 | return; |
| 263 | #endif |
| 264 | |
| 265 | #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ == 750 |
| 266 | constexpr auto num_bits = vllm::ScalarType::from_id(b_type_id).size_bits(); |
| 267 | // Disable use_fp16_accum for NVFP4 and cases when group_size == -1 && |
| 268 | // num_bits == 4 |
| 269 | constexpr bool use_fp16_accum = |
| 270 | a_type_id == vllm::kFloat16.id() && |
| 271 | (!(b_type_id == vllm::kFE2M1f.id() && s_type_id == vllm::kFE4M3fn.id()) && |
| 272 | !(group_blocks == -1 && num_bits == 4)); |
| 273 | #else |
| 274 | constexpr bool use_fp16_accum = false; |
| 275 | #endif |
| 276 | using Adtype = MarlinScalarType<a_type_id>; |
| 277 | using Cdtype = MarlinScalarType<c_type_id>; |
| 278 | const int4* A = A0; |
| 279 | int4* C = C0; |
| 280 | |
| 281 | using scalar_t = typename MarlinScalarType<a_type_id>::scalar_t; |
| 282 | using scalar_t2 = typename MarlinScalarType<a_type_id>::scalar_t2; |
| 283 | using scalar_32bit_t = typename MarlinScalarType<a_type_id>::scalar_32bit_t; |
| 284 | |
| 285 | using c_scalar_t = typename MarlinScalarType<c_type_id>::scalar_t; |
| 286 | using c_scalar_t2 = typename MarlinScalarType<c_type_id>::scalar_t2; |
| 287 | |
| 288 | using FragA = typename MarlinScalarType<a_type_id>::FragA; |
| 289 | using FragB = typename MarlinScalarType<a_type_id>::FragB; |
| 290 | using FragC = typename MarlinScalarType<a_type_id>::FragC; |
| 291 | using FragS = typename MarlinScalarType<c_type_id>::FragS; |
| 292 | using FragZP = typename MarlinScalarType<c_type_id>::FragZP; |
| 293 | |
| 294 | static constexpr auto a_type = vllm::ScalarType::from_id(a_type_id); |
| 295 | static constexpr auto b_type = vllm::ScalarType::from_id(b_type_id); |
| 296 | static constexpr auto c_type = vllm::ScalarType::from_id(c_type_id); |
| 297 | static constexpr auto s_type = vllm::ScalarType::from_id(s_type_id); |
| 298 | if constexpr (b_type == vllm::kFE2M1f) { |
| 299 | static_assert(s_type == vllm::kFE4M3fn && group_blocks == 1 || |
| 300 | s_type == vllm::kFE8M0fnu && group_blocks == 2); |
| 301 | } else if constexpr (s_type == vllm::kFE8M0fnu) { |
| 302 | // MXFP8: FP8 weights with e8m0 microscaling block scales |
| 303 | static_assert(b_type == vllm::kFE4M3fn && group_blocks == 2); |
| 304 | } else if constexpr (std::is_same<scalar_t, nv_bfloat16>::value) { |
| 305 | static_assert(s_type == vllm::kBFloat16); |
| 306 | } else if constexpr (std::is_same<scalar_t, half>::value) { |
| 307 | static_assert(s_type == vllm::kFloat16); |
| 308 | } |
| 309 | |
| 310 | constexpr bool is_a_8bit = a_type.size_bits() == 8; |
| 311 | constexpr bool is_8bit_scale = s_type.size_bits() == 8; |
| 312 | if constexpr (!is_a_8bit) { |
| 313 | static_assert(std::is_same<scalar_t, c_scalar_t>::value); |
| 314 | } |
| 315 | constexpr bool has_zp = b_type == vllm::kU4 || b_type == vllm::kU8; |
| 316 | constexpr bool is_int_type = b_type == vllm::kU4 || b_type == vllm::kU8 || |
| 317 | b_type == vllm::kS4 || b_type == vllm::kS8 || |
| 318 | b_type == vllm::kU4B8 || b_type == vllm::kU8B128; |
| 319 | // see comments of dequant.h for more details |
| 320 | constexpr bool dequant_skip_flop = |
| 321 | is_a_8bit || (b_type == vllm::kFE4M3fn && !(s_type == vllm::kFE8M0fnu)) || |
| 322 | b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn || |
| 323 | has_zp && !is_zp_float && !std::is_same<scalar_t, nv_bfloat16>::value || |
| 324 | has_zp && !is_zp_float && !(b_type == vllm::kU8); |
| 325 | |
| 326 | float global_scale_f32 = 1.0f; |
| 327 | |
| 328 | if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { |
| 329 | global_scale_f32 = global_scale_ptr[0]; |
| 330 | } |
| 331 | |
| 332 | static_assert(group_blocks != 0, |
| 333 | "runtime activation ordering is not supported"); |
| 334 | constexpr int m_block_size = m_block_size_8 ? 8 : (16 * thread_m_blocks); |
| 335 | |
| 336 | extern __shared__ int4 sh[]; |
| 337 | float* sh_a_s = reinterpret_cast<float*>(sh); |
| 338 | int4* sh_new = sh + (is_a_8bit ? (4 * thread_m_blocks) : 0); |
| 339 | constexpr int pack_factor = 32 / b_type.size_bits(); |
| 340 | static_assert(thread_m_blocks == 1 || !m_block_size_8); |
| 341 | |
| 342 | // For larger GEMMs we run multiple batchsize 64 versions in parallel for a |
| 343 | // better partitioning with less reductions |
| 344 | int parallel = 1; |
| 345 | if (prob_m > m_block_size) { |
| 346 | parallel = prob_m / m_block_size; |
| 347 | prob_m = m_block_size; |
| 348 | } |
| 349 | |
| 350 | int k_tiles = prob_k / 16 / thread_k_blocks; |
| 351 | int n_tiles = prob_n / 16 / thread_n_blocks; |
| 352 | |
| 353 | int global_mn_tiles = parallel * n_tiles; |
| 354 | int part2_mn_tiles = global_mn_tiles; |
| 355 | int part1_mn_iters = 0; |
| 356 | bool in_part2 = false; |
| 357 | |
| 358 | if (global_mn_tiles > gridDim.x) { |
| 359 | part2_mn_tiles = global_mn_tiles % gridDim.x; |
| 360 | if (part2_mn_tiles * 3 <= gridDim.x) part2_mn_tiles += gridDim.x; |
| 361 | part1_mn_iters = (global_mn_tiles - part2_mn_tiles) / gridDim.x; |
| 362 | } |
| 363 | |
| 364 | int iters = div_ceil(k_tiles * part2_mn_tiles, gridDim.x); |
| 365 | |
| 366 | if constexpr (group_blocks != -1) { |
| 367 | if (group_blocks >= thread_k_blocks) { |
| 368 | // Ensure that the number of tiles in each stripe is a multiple of the |
| 369 | // groupsize; this avoids an annoying special case where a stripe starts |
| 370 | // in the middle of group. |
| 371 | iters = (group_blocks / thread_k_blocks) * |
| 372 | div_ceil(iters, (group_blocks / thread_k_blocks)); |
| 373 | } |
| 374 | } |
| 375 | |
| 376 | int slice_row = 0; |
| 377 | int slice_col_par = blockIdx.x; |
| 378 | int slice_col; |
| 379 | int slice_iters = |
| 380 | k_tiles; // number of threadblock tiles in the current slice |
| 381 | // total number of active threadblocks in the current slice |
| 382 | int slice_count = 1; |
| 383 | // index of threadblock in current slice; numbered bottom to top |
| 384 | int slice_idx = 0; |
| 385 | |
| 386 | int par_id = 0; |
| 387 | int locks_off = 0; |
| 388 | |
| 389 | if (part2_mn_tiles >= gridDim.x) { |
| 390 | // when part2_mn_tiles >= sms |
| 391 | // then there are at most $sms$ conflict tile blocks |
| 392 | locks_off = blockIdx.x; |
| 393 | } else { |
| 394 | locks_off = (iters * blockIdx.x) / k_tiles - 1; |
| 395 | } |
| 396 | |
| 397 | // Compute all information about the current slice which is required for |
| 398 | // synchronization. |
| 399 | bool first_init = true; |
| 400 | auto init_part2_slice = [&]() { |
| 401 | slice_iters = |
| 402 | iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row); |
| 403 | if (slice_iters < 0 || slice_col_par >= part2_mn_tiles) slice_iters = 0; |
| 404 | if (slice_iters == 0) return; |
| 405 | if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row; |
| 406 | slice_count = 1; |
| 407 | slice_idx = 0; |
| 408 | int col_first = iters * div_ceil(k_tiles * slice_col_par, iters); |
| 409 | if (col_first <= k_tiles * (slice_col_par + 1)) { |
| 410 | int col_off = col_first - k_tiles * slice_col_par; |
| 411 | slice_count = div_ceil(k_tiles - col_off, iters); |
| 412 | if (col_off > 0) slice_count++; |
| 413 | int delta_first = iters * blockIdx.x - col_first; |
| 414 | if (delta_first < 0 || (col_off == 0 && delta_first == 0)) |
| 415 | slice_idx = slice_count - 1; |
| 416 | else { |
| 417 | slice_idx = slice_count - 1 - delta_first / iters; |
| 418 | if (col_off > 0) slice_idx--; |
| 419 | } |
| 420 | } |
| 421 | if (part2_mn_tiles >= gridDim.x) { |
| 422 | if (slice_count > 1 && slice_idx == slice_count - 1) { |
| 423 | locks_off++; |
| 424 | } |
| 425 | } else { |
| 426 | locks_off++; |
| 427 | } |
| 428 | |
| 429 | if (first_init && use_atomic_add && slice_count > 1 && slice_idx == 0) { |
| 430 | constexpr int threads_per_m = 16 * thread_n_blocks / 8; |
| 431 | int m_per_thread = |
| 432 | div_ceil(thread_m_blocks * 16, threads / threads_per_m); |
| 433 | if (m_block_size_8) m_per_thread = div_ceil(8, threads / threads_per_m); |
| 434 | for (int i = 0; i < m_per_thread; i++) { |
| 435 | int row = threads / threads_per_m * i + threadIdx.x / threads_per_m; |
| 436 | if (row < prob_m) { |
| 437 | int col = slice_col * 16 * thread_n_blocks / 8 + |
| 438 | threadIdx.x % threads_per_m; |
| 439 | C[row * prob_n / 8 + col] = {0, 0, 0, 0}; |
| 440 | } |
| 441 | } |
| 442 | // After write zero to output, write a negative value to lock. |
| 443 | // Every SM that processes the same slice would wait for |
| 444 | // the negative value, and then atomicAdd 1 to it. |
| 445 | // After all SMs are processed, the lock value would back to 0 again. |
| 446 | __syncthreads(); |
| 447 | if (threadIdx.x == 0) locks[locks_off] = 1 - slice_count; |
| 448 | } |
| 449 | |
| 450 | if (slice_col == n_tiles) { |
| 451 | A += 16 * thread_m_blocks * lda / (is_a_8bit ? 16 : 8); |
| 452 | C += 16 * thread_m_blocks * prob_n / 8; |
| 453 | slice_col = 0; |
| 454 | par_id++; |
| 455 | } |
| 456 | if (is_a_8bit && (first_init || slice_col == 0)) { |
| 457 | __syncthreads(); |
| 458 | int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; |
| 459 | cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], |
| 460 | threadIdx.x < prob_m); |
| 461 | } |
| 462 | }; |
| 463 | |
| 464 | auto init_part1_slice = [&]() { |
| 465 | if (part1_mn_iters) { |
| 466 | part1_mn_iters--; |
| 467 | par_id = slice_col_par / n_tiles; |
| 468 | slice_col = slice_col_par % n_tiles; |
| 469 | slice_iters = k_tiles; |
| 470 | A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; |
| 471 | C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; |
| 472 | if (is_a_8bit) { |
| 473 | __syncthreads(); |
| 474 | int a_s_gl_rd = par_id * 16 * thread_m_blocks + threadIdx.x; |
| 475 | cp_async1_ca_pred(&sh_a_s[threadIdx.x], &a_scales_ptr[a_s_gl_rd], |
| 476 | threadIdx.x < prob_m); |
| 477 | } |
| 478 | } |
| 479 | }; |
| 480 | |
| 481 | auto init_slice = [&]() { |
| 482 | if (!in_part2 && !part1_mn_iters) { |
| 483 | in_part2 = true; |
| 484 | slice_col_par = (iters * blockIdx.x) / k_tiles; |
| 485 | slice_row = (iters * blockIdx.x) % k_tiles; |
| 486 | slice_col = (slice_col_par + global_mn_tiles - part2_mn_tiles) % n_tiles; |
| 487 | par_id = (slice_col_par + global_mn_tiles - part2_mn_tiles) / n_tiles; |
| 488 | A = A0 + 16 * thread_m_blocks / (is_a_8bit ? 16 : 8) * par_id * lda; |
| 489 | C = C0 + 16 * thread_m_blocks / 8 * par_id * prob_n; |
| 490 | } |
| 491 | if (!in_part2) { |
| 492 | init_part1_slice(); |
| 493 | } else { |
| 494 | init_part2_slice(); |
| 495 | first_init = false; |
| 496 | } |
| 497 | }; |
| 498 | |
| 499 | init_slice(); |
| 500 | |
| 501 | // A sizes/strides |
| 502 | |
| 503 | // stride of the A matrix in global memory |
| 504 | int a_gl_stride = lda / (is_a_8bit ? 16 : 8); |
| 505 | // stride of an A matrix tile in shared memory |
| 506 | constexpr int a_sh_stride = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); |
| 507 | // delta between subsequent A tiles in global memory |
| 508 | constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / (is_a_8bit ? 16 : 8); |
| 509 | // between subsequent accesses within a tile |
| 510 | int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o); |
| 511 | // between shared memory writes |
| 512 | constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o); |
| 513 | // within a shared memory tile |
| 514 | constexpr int a_sh_rd_delta_i = a_sh_stride * 16; |
| 515 | // overall size of a tile |
| 516 | constexpr int a_sh_stage = a_sh_stride * m_block_size; |
| 517 | // number of shared write iterations for a tile |
| 518 | constexpr int a_sh_wr_iters = div_ceil(a_sh_stage, a_sh_wr_delta); |
| 519 | |
| 520 | // B sizes/strides |
| 521 | int b_gl_stride = 16 * prob_n / (pack_factor * (is_a_8bit ? 2 : 4)); |
| 522 | constexpr int b_sh_stride = |
| 523 | ((thread_n_blocks * 16) * 16 / pack_factor) / (is_a_8bit ? 2 : 4); |
| 524 | constexpr int b_thread_vecs = b_type.size_bits() == 4 ? 1 : 2; |
| 525 | constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs; |
| 526 | |
| 527 | int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); |
| 528 | constexpr int b_sh_wr_delta = threads * b_thread_vecs; |
| 529 | constexpr int b_sh_stage = |
| 530 | b_sh_stride * thread_k_blocks / (is_a_8bit ? 2 : 1); |
| 531 | constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta; |
| 532 | |
| 533 | // Scale sizes/strides |
| 534 | int s_gl_stride = prob_n / (is_8bit_scale ? 16 : 8); |
| 535 | constexpr int s_sh_stride = 16 * thread_n_blocks / (is_8bit_scale ? 16 : 8); |
| 536 | constexpr int s_tb_groups = |
| 537 | group_blocks != -1 && group_blocks < thread_k_blocks |
| 538 | ? thread_k_blocks / group_blocks |
| 539 | : 1; |
| 540 | constexpr int s_sh_stage = s_tb_groups * s_sh_stride; |
| 541 | int s_gl_rd_delta = s_gl_stride; |
| 542 | |
| 543 | constexpr int tb_n_warps = thread_n_blocks / (is_a_8bit ? 2 : 4); |
| 544 | |
| 545 | // Zero-points sizes/strides |
| 546 | int zp_gl_stride = is_zp_float ? prob_n / 8 : (prob_n / pack_factor) / 4; |
| 547 | constexpr int zp_sh_stride = is_zp_float |
| 548 | ? 16 * thread_n_blocks / 8 |
| 549 | : ((16 * thread_n_blocks) / pack_factor) / 4; |
| 550 | constexpr int zp_tb_groups = s_tb_groups; |
| 551 | constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0; |
| 552 | int zp_gl_rd_delta = zp_gl_stride; |
| 553 | |
| 554 | // Global A read index of current thread. |
| 555 | int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + |
| 556 | (threadIdx.x % a_gl_rd_delta_o); |
| 557 | a_gl_rd += a_gl_rd_delta_o * slice_row; |
| 558 | // Shared write index of current thread. |
| 559 | int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) + |
| 560 | (threadIdx.x % a_gl_rd_delta_o); |
| 561 | // Shared read index. |
| 562 | int a_sh_rd = |
| 563 | a_sh_stride * ((threadIdx.x % 32) % (16 / (m_block_size_8 ? 2 : 1))) + |
| 564 | (threadIdx.x % 32) / (16 / (m_block_size_8 ? 2 : 1)); |
| 565 | a_sh_rd += 2 * ((threadIdx.x / 32) / tb_n_warps) * b_sh_wr_iters; |
| 566 | |
| 567 | int b_gl_rd; |
| 568 | if (threads <= b_sh_stride) { |
| 569 | b_gl_rd = threadIdx.x; |
| 570 | } else { |
| 571 | b_gl_rd = |
| 572 | b_gl_stride * (threadIdx.x / b_sh_stride) + (threadIdx.x % b_sh_stride); |
| 573 | } |
| 574 | |
| 575 | b_gl_rd += b_sh_stride * slice_col; |
| 576 | b_gl_rd += b_gl_rd_delta_o * slice_row; |
| 577 | auto b_sh_rd = threadIdx.x * b_thread_vecs; |
| 578 | b_sh_rd += b_sh_rd / b_sh_stride * (b_sh_stride * (b_sh_wr_iters - 1)); |
| 579 | |
| 580 | int s_gl_rd; |
| 581 | if constexpr (group_blocks == -1) { |
| 582 | s_gl_rd = s_sh_stride * slice_col + threadIdx.x; |
| 583 | } else if constexpr (group_blocks >= thread_k_blocks) { |
| 584 | s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + |
| 585 | s_sh_stride * slice_col + threadIdx.x; |
| 586 | } else { |
| 587 | s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + |
| 588 | threadIdx.x / s_sh_stride) + |
| 589 | s_sh_stride * slice_col + threadIdx.x % s_sh_stride; |
| 590 | } |
| 591 | auto s_sh_wr = threadIdx.x; |
| 592 | bool s_sh_wr_pred = threadIdx.x < s_sh_stage; |
| 593 | |
| 594 | // Zero-points |
| 595 | int zp_gl_rd; |
| 596 | if constexpr (has_zp) { |
| 597 | if constexpr (group_blocks == -1) { |
| 598 | zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; |
| 599 | } else if constexpr (group_blocks >= thread_k_blocks) { |
| 600 | zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + |
| 601 | zp_sh_stride * slice_col + threadIdx.x; |
| 602 | } else { |
| 603 | zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + |
| 604 | threadIdx.x / zp_sh_stride) + |
| 605 | zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; |
| 606 | } |
| 607 | } |
| 608 | auto zp_sh_wr = threadIdx.x; |
| 609 | bool zp_sh_wr_pred = zp_sh_stage > 0 && threadIdx.x < zp_sh_stage; |
| 610 | |
| 611 | // We use a different scale layout for grouped and column-wise quantization as |
| 612 | // we scale a `half2` tile in column-major layout in the former and in |
| 613 | // row-major in the latter case. |
| 614 | int s_sh_rd; |
| 615 | if constexpr (is_a_8bit) { |
| 616 | s_sh_rd = 4 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 4); |
| 617 | } else if constexpr (group_blocks != -1) |
| 618 | s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; |
| 619 | else if constexpr (group_blocks == -1 && |
| 620 | (m_block_size_8 || (has_zp && !dequant_skip_flop))) |
| 621 | s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; |
| 622 | else |
| 623 | s_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) % 4; |
| 624 | |
| 625 | int bias_sh_rd; |
| 626 | if constexpr (m_block_size_8) { |
| 627 | bias_sh_rd = 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 8; |
| 628 | } else { |
| 629 | bias_sh_rd = (is_a_8bit ? 4 : 8) * ((threadIdx.x / 32) % tb_n_warps) + |
| 630 | (threadIdx.x % 32) % 4; |
| 631 | } |
| 632 | |
| 633 | int bias_sh_wr = threadIdx.x; |
| 634 | int bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; |
| 635 | |
| 636 | // Zero-points have the same read layout as the scales |
| 637 | // (without column-wise case) |
| 638 | constexpr int num_col_threads = 8; |
| 639 | constexpr int num_row_threads = 4; |
| 640 | constexpr int num_ints_per_thread = 8 / pack_factor; |
| 641 | int zp_sh_rd; |
| 642 | if constexpr (has_zp) { |
| 643 | if constexpr (is_zp_float) { |
| 644 | if constexpr (group_blocks != -1) { |
| 645 | zp_sh_rd = |
| 646 | 8 * ((threadIdx.x / 32) % tb_n_warps) + (threadIdx.x % 32) / 4; |
| 647 | } |
| 648 | } else if (is_a_8bit) { |
| 649 | zp_sh_rd = num_ints_per_thread * num_col_threads * |
| 650 | ((threadIdx.x / 32) % tb_n_warps / 2) + |
| 651 | num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); |
| 652 | } else { |
| 653 | zp_sh_rd = num_ints_per_thread * num_col_threads * |
| 654 | ((threadIdx.x / 32) % tb_n_warps) + |
| 655 | num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads); |
| 656 | } |
| 657 | } |
| 658 | |
| 659 | // Precompute which thread should not read memory in which iterations; this is |
| 660 | // needed if there are more threads than required for a certain tilesize or |
| 661 | // when the batchsize is not a multiple of 16. |
| 662 | bool a_sh_wr_pred[a_sh_wr_iters]; |
| 663 | #pragma unroll |
| 664 | for (int i = 0; i < a_sh_wr_iters; i++) |
| 665 | a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m; |
| 666 | |
| 667 | // To ensure that writing and reading A tiles to/from shared memory, the |
| 668 | // latter in fragment format, is fully bank conflict free, we need to use a |
| 669 | // rather fancy XOR-based layout. The key here is that neither reads nor |
| 670 | // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the |
| 671 | // same shared memory banks. Further, it seems (based on NSight-Compute) that |
| 672 | // each warp must also write a consecutive memory segment? |
| 673 | auto transform_a = [&](int i) { |
| 674 | int row = i / a_gl_rd_delta_o; |
| 675 | return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ (row % 8); |
| 676 | }; |
| 677 | // Since the computation of this remapping is non-trivial and, due to our main |
| 678 | // loop unrolls, all shared memory accesses are static, we simply precompute |
| 679 | // both transformed reads and writes. |
| 680 | int a_sh_wr_trans[a_sh_wr_iters]; |
| 681 | #pragma unroll |
| 682 | for (int i = 0; i < a_sh_wr_iters; i++) |
| 683 | a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr); |
| 684 | int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks]; |
| 685 | #pragma unroll |
| 686 | for (int i = 0; i < b_sh_wr_iters; i++) { |
| 687 | #pragma unroll |
| 688 | for (int j = 0; j < thread_m_blocks; j++) |
| 689 | a_sh_rd_trans[i][j] = transform_a(2 * i + a_sh_rd_delta_i * j + a_sh_rd); |
| 690 | } |
| 691 | |
| 692 | // Since B-accesses have non-constant stride they have to be computed at |
| 693 | // runtime; we break dependencies between subsequent accesses with a tile by |
| 694 | // maintining multiple pointers (we have enough registers), a tiny |
| 695 | // optimization. |
| 696 | |
| 697 | // Shared memory storage for global fetch pipelines. |
| 698 | constexpr int sh_red_size = (2 * thread_n_blocks + 1) * 16 * thread_m_blocks; |
| 699 | constexpr int sh_b_size = stages * b_sh_stage; |
| 700 | int4* sh_b = sh_new; |
| 701 | int4* sh_red = sh_new; |
| 702 | constexpr int sh_size_b_red_min = |
| 703 | (sh_red_size < sh_b_size ? sh_red_size : sh_b_size); |
| 704 | constexpr int sh_size_b_red_max = |
| 705 | (sh_red_size > sh_b_size ? sh_red_size : sh_b_size); |
| 706 | constexpr int sh_bias_size = (thread_n_blocks * 16 / 8); |
| 707 | constexpr int sh_b_red_bias_size = |
| 708 | sh_size_b_red_max > (sh_size_b_red_min + sh_bias_size) |
| 709 | ? sh_size_b_red_max |
| 710 | : (sh_size_b_red_min + sh_bias_size); |
| 711 | |
| 712 | int4* sh_bias = sh_new + sh_size_b_red_min; |
| 713 | int4* sh_zp = sh_new + sh_b_red_bias_size; |
| 714 | constexpr int sh_s_size = stages * s_sh_stage; |
| 715 | int4* sh_s = sh_zp + (stages * zp_sh_stage); |
| 716 | int4* sh_a = sh_s + sh_s_size; |
| 717 | |
| 718 | // Register storage for double buffer of shared memory reads. |
| 719 | FragA frag_a[2][thread_m_blocks]; |
| 720 | I4 frag_b_quant[2][b_thread_vecs]; |
| 721 | FragC frag_c[thread_m_blocks][is_a_8bit ? 2 : 4][2]; |
| 722 | FragC frag_c_tmp[thread_m_blocks][is_a_8bit ? 2 : 4][2]; |
| 723 | FragS frag_s[2][4]; |
| 724 | FragS frag_bias[2][4]; |
| 725 | int frag_qzp[2][num_ints_per_thread]; // Zero-points |
| 726 | FragZP frag_zp; // Zero-points in fp16 |
| 727 | FragZP frag_zpf[2]; // Zero-points in fp16 in HQQ |
| 728 | |
| 729 | if constexpr (is_a_8bit) { |
| 730 | #pragma unroll |
| 731 | for (int j = 0; j < 2; j++) { |
| 732 | #pragma unroll |
| 733 | for (int i = 0; i < thread_m_blocks; i++) { |
| 734 | #pragma unroll |
| 735 | for (int g = 0; g < 4; g++) { |
| 736 | frag_c_tmp[i][j][0][g] = 0.0f; |
| 737 | } |
| 738 | |
| 739 | #pragma unroll |
| 740 | for (int g = 0; g < 4; g++) { |
| 741 | frag_c_tmp[i][j][1][g] = 0.0f; |
| 742 | } |
| 743 | } |
| 744 | } |
| 745 | } |
| 746 | |
| 747 | // Zero accumulators. |
| 748 | auto zero_accums = [&]() { |
| 749 | #pragma unroll |
| 750 | for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++) |
| 751 | reinterpret_cast<float*>(frag_c)[i] = 0; |
| 752 | }; |
| 753 | |
| 754 | // Asynchronously fetch the next A, B and s tile from global to the next |
| 755 | // shared memory pipeline location. |
| 756 | auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) { |
| 757 | if (pred) { |
| 758 | int4* sh_a_stage = sh_a + a_sh_stage * pipe; |
| 759 | #pragma unroll |
| 760 | for (int i = 0; i < a_sh_wr_iters; i++) { |
| 761 | cp_async4_pred( |
| 762 | &sh_a_stage[a_sh_wr_trans[i]], |
| 763 | &A[a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off], |
| 764 | a_sh_wr_pred[i]); |
| 765 | } |
| 766 | int4* sh_b_stage = sh_b + b_sh_stage * pipe; |
| 767 | #pragma unroll |
| 768 | for (int i = 0; i < (b_sh_wr_iters * b_thread_vecs); i++) { |
| 769 | constexpr int count = div_ceil(b_sh_stride, threads); |
| 770 | int b_gl_idx = |
| 771 | b_gl_rd + (i % count) * threads + |
| 772 | b_gl_stride * (i / count) * div_ceil(threads, b_sh_stride); |
| 773 | |
| 774 | cp_async4(&sh_b_stage[threads * i + threadIdx.x], &B[b_gl_idx]); |
| 775 | } |
| 776 | |
| 777 | b_gl_rd += b_gl_rd_delta_o; |
| 778 | |
| 779 | if constexpr (group_blocks != -1) { |
| 780 | int4* sh_s_stage = sh_s + s_sh_stage * pipe; |
| 781 | |
| 782 | // Only fetch scales if this tile starts a new group |
| 783 | if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { |
| 784 | if (s_sh_wr_pred) { |
| 785 | cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]); |
| 786 | } |
| 787 | s_gl_rd += s_gl_rd_delta * s_tb_groups; |
| 788 | } |
| 789 | } |
| 790 | |
| 791 | if constexpr (has_zp && group_blocks != -1) { |
| 792 | int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; |
| 793 | |
| 794 | // Only fetch zero points if this tile starts a new group |
| 795 | if (pipe % div_ceil(group_blocks, thread_k_blocks) == 0) { |
| 796 | if (zp_sh_wr_pred) { |
| 797 | cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]); |
| 798 | } |
| 799 | zp_gl_rd += zp_gl_rd_delta * zp_tb_groups; |
| 800 | } |
| 801 | } |
| 802 | } |
| 803 | // Insert a fence even when we are winding down the pipeline to ensure that |
| 804 | // waiting is also correct at this point. |
| 805 | cp_async_fence(); |
| 806 | }; |
| 807 | |
| 808 | auto fetch_col_zp_to_shared = [&]() { |
| 809 | if (zp_sh_wr_pred) { |
| 810 | cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]); |
| 811 | } |
| 812 | }; |
| 813 | |
| 814 | auto fetch_col_scale_to_shared = [&]() { |
| 815 | if (s_sh_wr_pred) { |
| 816 | cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); |
| 817 | } |
| 818 | }; |
| 819 | |
| 820 | // Wait until the next thread tile has been loaded to shared memory. |
| 821 | auto wait_for_stage = [&]() { |
| 822 | // We only have `stages - 2` active fetches since we are double buffering |
| 823 | // and can only issue the next fetch when it is guaranteed that the previous |
| 824 | // shared memory load is fully complete (as it may otherwise be |
| 825 | // overwritten). |
| 826 | cp_async_wait<stages - 2>(); |
| 827 | __syncthreads(); |
| 828 | }; |
| 829 | |
| 830 | // Load the next sub-tile from the current location in the shared memory pipe |
| 831 | // into the current register buffer. |
| 832 | auto fetch_to_registers = [&](int k, int pipe) { |
| 833 | int4* sh_a_stage = sh_a + a_sh_stage * pipe; |
| 834 | #pragma unroll |
| 835 | for (int i = 0; i < thread_m_blocks; i++) |
| 836 | ldsm<m_block_size_8 ? 2 : 4, a_type_id>( |
| 837 | frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]); |
| 838 | int4* sh_b_stage = sh_b + b_sh_stage * pipe; |
| 839 | |
| 840 | #pragma unroll |
| 841 | for (int i = 0; i < b_thread_vecs; i++) { |
| 842 | frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>( |
| 843 | &sh_b_stage[b_sh_stride * (k % b_sh_wr_iters) + b_sh_rd + i]); |
| 844 | } |
| 845 | }; |
| 846 | |
| 847 | auto fetch_scales_to_registers = [&](int k, int full_pipe) { |
| 848 | int pipe = full_pipe % stages; |
| 849 | using IT1 = typename std::conditional_t<is_a_8bit, int2, int4>; |
| 850 | using IT0 = typename std::conditional_t<is_a_8bit, int, int2>; |
| 851 | constexpr int group_blocks2 = div_ceil(group_blocks, is_a_8bit ? 2 : 1); |
| 852 | |
| 853 | if constexpr (group_blocks == -1) { |
| 854 | // Load only when starting a new slice. |
| 855 | if (k == 0 && full_pipe == 0 && dequant_skip_flop) { |
| 856 | reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd]; |
| 857 | reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4]; |
| 858 | } |
| 859 | } else if constexpr (group_blocks >= thread_k_blocks) { |
| 860 | constexpr int g = group_blocks / thread_k_blocks; |
| 861 | if (pipe % g == 0) { |
| 862 | if (k % b_sh_wr_iters == 0) { |
| 863 | int4* sh_s_stage = sh_s + s_sh_stage * (g * (pipe / g)); |
| 864 | reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd]; |
| 865 | } else { |
| 866 | reinterpret_cast<int4*>(&frag_s[1])[0] = |
| 867 | reinterpret_cast<int4*>(&frag_s[0])[0]; |
| 868 | } |
| 869 | } |
| 870 | } else if (group_blocks2 < b_sh_wr_iters || k % b_sh_wr_iters == 0) { |
| 871 | auto warp_id = threadIdx.x / 32; |
| 872 | int warp_row = warp_id / tb_n_warps; |
| 873 | |
| 874 | int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; |
| 875 | int cur_group_id = k_blocks / group_blocks2; |
| 876 | |
| 877 | int4* sh_s_stage = sh_s + s_sh_stage * pipe; |
| 878 | |
| 879 | if constexpr (!is_8bit_scale) { |
| 880 | reinterpret_cast<int4*>(&frag_s[k % 2])[0] = |
| 881 | sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride]; |
| 882 | } else { |
| 883 | reinterpret_cast<int2*>(&frag_s[k % 2])[0] = reinterpret_cast<int2*>( |
| 884 | sh_s_stage)[s_sh_rd + cur_group_id * (2 * s_sh_stride)]; |
| 885 | } |
| 886 | } else if (group_blocks >= b_sh_wr_iters) { |
| 887 | if constexpr (!is_8bit_scale) { |
| 888 | reinterpret_cast<int4*>(&frag_s[1])[0] = |
| 889 | reinterpret_cast<int4*>(&frag_s[0])[0]; |
| 890 | } else { |
| 891 | reinterpret_cast<int2*>(&frag_s[1])[0] = |
| 892 | reinterpret_cast<int2*>(&frag_s[0])[0]; |
| 893 | } |
| 894 | } |
| 895 | }; |
| 896 | |
| 897 | auto fetch_zp_to_registers = [&](int k, int full_pipe) { |
| 898 | static_assert(!has_zp || group_blocks != 0); |
| 899 | |
| 900 | if constexpr (has_zp && !is_zp_float) { |
| 901 | int pipe = full_pipe % stages; |
| 902 | |
| 903 | if constexpr (group_blocks == -1) { |
| 904 | // load only when starting a new slice |
| 905 | if (k == 0 && full_pipe == 0 || is_a_8bit) { |
| 906 | #pragma unroll |
| 907 | for (int i = 0; i < num_ints_per_thread; i++) { |
| 908 | frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i]; |
| 909 | } |
| 910 | } |
| 911 | } else if constexpr (group_blocks >= thread_k_blocks) { |
| 912 | constexpr int g = group_blocks / thread_k_blocks; |
| 913 | if (pipe % g == 0 && k % b_sh_wr_iters == 0 || is_a_8bit) { |
| 914 | int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); |
| 915 | #pragma unroll |
| 916 | for (int i = 0; i < num_ints_per_thread; i++) { |
| 917 | frag_qzp[k % 2][i] = |
| 918 | (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i]; |
| 919 | } |
| 920 | } |
| 921 | } else { |
| 922 | auto warp_id = threadIdx.x / 32; |
| 923 | |
| 924 | int warp_row = warp_id / tb_n_warps; |
| 925 | |
| 926 | int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; |
| 927 | int cur_group_id = k_blocks / div_ceil(group_blocks, is_a_8bit ? 2 : 1); |
| 928 | |
| 929 | int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; |
| 930 | |
| 931 | sh_zp_stage += cur_group_id * zp_sh_stride; |
| 932 | |
| 933 | #pragma unroll |
| 934 | for (int i = 0; i < num_ints_per_thread; i++) { |
| 935 | frag_qzp[k % 2][i] = |
| 936 | (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i]; |
| 937 | } |
| 938 | } |
| 939 | } |
| 940 | |
| 941 | else if constexpr (has_zp && is_zp_float) { |
| 942 | int pipe = full_pipe % stages; |
| 943 | |
| 944 | if constexpr (group_blocks != -1) { |
| 945 | if constexpr (group_blocks >= thread_k_blocks) { |
| 946 | constexpr int g = group_blocks / thread_k_blocks; |
| 947 | if (pipe % g == 0 && k % b_sh_wr_iters == 0) { |
| 948 | int4* sh_zp_stage = sh_zp + zp_sh_stage * (g * (pipe / g)); |
| 949 | reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] = |
| 950 | sh_zp_stage[zp_sh_rd]; |
| 951 | } |
| 952 | } else if (group_blocks < b_sh_wr_iters || k % b_sh_wr_iters == 0) { |
| 953 | auto warp_id = threadIdx.x / 32; |
| 954 | |
| 955 | int warp_row = warp_id / tb_n_warps; |
| 956 | int k_blocks = b_sh_wr_iters * warp_row + k % b_sh_wr_iters; |
| 957 | int cur_group_id = k_blocks / group_blocks; |
| 958 | |
| 959 | int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe; |
| 960 | |
| 961 | reinterpret_cast<int4*>(&frag_zpf[k % 2])[0] = |
| 962 | sh_zp_stage[zp_sh_rd + cur_group_id * zp_sh_stride]; |
| 963 | } |
| 964 | } |
| 965 | } |
| 966 | }; |
| 967 | |
| 968 | auto dequant_data = [&](int q, scalar_32bit_t* frag_b_ptr, int zp = 0) { |
| 969 | if constexpr (a_type.size_bits() != b_type.size_bits()) { |
| 970 | if constexpr (is_a_8bit && has_zp) { |
| 971 | sub_zp_and_dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>( |
| 972 | q, frag_b_ptr, zp); |
| 973 | } else { |
| 974 | dequant<scalar_32bit_t, b_type_id, dequant_skip_flop>(q, frag_b_ptr); |
| 975 | } |
| 976 | } |
| 977 | }; |
| 978 | |
| 979 | // Execute the actual tensor core matmul of a sub-tile. |
| 980 | bool is_first_matmul_in_slice = true; |
| 981 | auto matmul = [&](int k, int pipe) { |
| 982 | if (is_a_8bit) return; |
| 983 | int k2 = k % 2; |
| 984 | constexpr int g = |
| 985 | group_blocks > 0 ? div_ceil(group_blocks, thread_k_blocks) : 1; |
| 986 | const bool is_new_zp = |
| 987 | (group_blocks > 0 && (group_blocks < b_sh_wr_iters || k == 0)) && |
| 988 | (pipe % g == 0) || |
| 989 | (group_blocks == -1 && is_first_matmul_in_slice); |
| 990 | if constexpr (has_zp && !is_zp_float) { |
| 991 | if (is_new_zp) { |
| 992 | if constexpr (group_blocks == -1) is_first_matmul_in_slice = false; |
| 993 | int zp_quant_0, zp_quant_1; |
| 994 | |
| 995 | if constexpr (b_type.size_bits() == 4) { |
| 996 | zp_quant_0 = frag_qzp[k2][0]; |
| 997 | zp_quant_1 = zp_quant_0 >> 8; |
| 998 | } else { |
| 999 | static_assert(b_type.size_bits() == 8); |
| 1000 | zp_quant_0 = frag_qzp[k2][0]; |
| 1001 | zp_quant_1 = frag_qzp[k2][1]; |
| 1002 | } |
| 1003 | |
| 1004 | dequant_data(zp_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_zp)); |
| 1005 | dequant_data(zp_quant_1, |
| 1006 | reinterpret_cast<scalar_32bit_t*>(&frag_zp) + 2); |
| 1007 | } |
| 1008 | } |
| 1009 | if constexpr (!dequant_skip_flop && has_zp && is_zp_float) { |
| 1010 | if (is_new_zp) { |
| 1011 | reinterpret_cast<int4*>(&frag_zp)[0] = |
| 1012 | reinterpret_cast<int4*>(&frag_zpf[k2])[0]; |
| 1013 | } |
| 1014 | } |
| 1015 | |
| 1016 | if constexpr (s_type == vllm::kFE4M3fn || s_type == vllm::kFE8M0fnu) { |
| 1017 | int s_quant_0 = reinterpret_cast<int*>(frag_s[k2])[0]; |
| 1018 | int s_quant_1 = reinterpret_cast<int*>(frag_s[k2])[1]; |
| 1019 | |
| 1020 | dequant_fp8_scales<c_scalar_t2, s_type_id>( |
| 1021 | s_quant_0, reinterpret_cast<c_scalar_t2*>(&frag_s[k2])); |
| 1022 | dequant_fp8_scales<c_scalar_t2, s_type_id>( |
| 1023 | s_quant_1, reinterpret_cast<c_scalar_t2*>(&frag_s[k2]) + 2); |
| 1024 | } |
| 1025 | |
| 1026 | // We have the m dimension as the inner loop in order to encourage overlapping |
| 1027 | // dequantization and matmul operations. |
| 1028 | #pragma unroll |
| 1029 | for (int j = 0; j < 4; j++) { |
| 1030 | FragB frag_b0; |
| 1031 | FragB frag_b1; |
| 1032 | int b_quant_0, b_quant_1; |
| 1033 | |
| 1034 | if constexpr (b_type_id == vllm::kFE2M1f.id()) { |
| 1035 | b_quant_1 = frag_b_quant[k2][0][j]; |
| 1036 | b_quant_0 = b_quant_1 << 8; |
| 1037 | } else if constexpr (b_type.size_bits() == 4) { |
| 1038 | b_quant_0 = frag_b_quant[k2][0][j]; |
| 1039 | b_quant_1 = b_quant_0 >> 8; |
| 1040 | } else { |
| 1041 | static_assert(b_type.size_bits() == 8); |
| 1042 | int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k2]); |
| 1043 | b_quant_0 = frag_b_quant_ptr[j * 2 + 0]; |
| 1044 | b_quant_1 = frag_b_quant_ptr[j * 2 + 1]; |
| 1045 | } |
| 1046 | |
| 1047 | dequant_data(b_quant_0, reinterpret_cast<scalar_32bit_t*>(&frag_b0)); |
| 1048 | dequant_data(b_quant_1, reinterpret_cast<scalar_32bit_t*>(&frag_b1)); |
| 1049 | |
| 1050 | if constexpr (dequant_skip_flop && has_zp && !is_zp_float && !is_a_8bit) { |
| 1051 | sub_zp<a_type_id>(frag_b0, frag_zp[j], 0); |
| 1052 | sub_zp<a_type_id>(frag_b1, frag_zp[j], 1); |
| 1053 | } |
| 1054 | |
| 1055 | if constexpr (!dequant_skip_flop && has_zp && !is_zp_float && |
| 1056 | group_blocks == -1 && !is_a_8bit) { |
| 1057 | int idx = (threadIdx.x / 4) % 2; |
| 1058 | scalar_t2 s2 = Adtype::nums2num2( |
| 1059 | reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 0])[idx], |
| 1060 | reinterpret_cast<scalar_t*>(&frag_s[j / 2][j % 2 * 2 + 1])[idx]); |
| 1061 | if (is_new_zp) frag_zp[j] = __hmul2(frag_zp[j], s2); |
| 1062 | scale_and_sub<a_type_id>(frag_b0, s2.x, frag_zp[j].x); |
| 1063 | scale_and_sub<a_type_id>(frag_b1, s2.y, frag_zp[j].y); |
| 1064 | } else if constexpr (!dequant_skip_flop && has_zp && group_blocks != -1 && |
| 1065 | !is_a_8bit) { |
| 1066 | if (is_new_zp) |
| 1067 | frag_zp[j] = __hmul2(frag_zp[j], |
| 1068 | *reinterpret_cast<scalar_t2*>(&frag_s[k2][j])); |
| 1069 | scale_and_sub<a_type_id>(frag_b0, frag_s[k2][j][0].x, frag_zp[j].x); |
| 1070 | scale_and_sub<a_type_id>(frag_b1, frag_s[k2][j][0].y, frag_zp[j].y); |
| 1071 | } else if constexpr (group_blocks != -1 && !is_a_8bit) { |
| 1072 | scale<a_type_id>(frag_b0, frag_s[k2][j], 0); |
| 1073 | scale<a_type_id>(frag_b1, frag_s[k2][j], 1); |
| 1074 | } |
| 1075 | |
| 1076 | #pragma unroll |
| 1077 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1078 | if constexpr (m_block_size_8) { |
| 1079 | mma_trans<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0, frag_b1, |
| 1080 | frag_c[i][j][0]); |
| 1081 | } else { |
| 1082 | mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b0, |
| 1083 | frag_c[i][j][0]); |
| 1084 | mma<a_type_id, use_fp16_accum>(frag_a[k2][i], frag_b1, |
| 1085 | frag_c[i][j][1]); |
| 1086 | } |
| 1087 | } |
| 1088 | } |
| 1089 | }; |
| 1090 | |
| 1091 | auto matmul_a8 = [&](int k) { |
| 1092 | int k2 = k % 2; |
| 1093 | #pragma unroll |
| 1094 | for (int j = 0; j < 2; j++) { |
| 1095 | FragB frag_b[2]; |
| 1096 | |
| 1097 | if (is_a_8bit && b_type.size_bits() == 4 && !has_zp) { |
| 1098 | dequant_data(frag_b_quant[k2][0][j * 2], |
| 1099 | reinterpret_cast<scalar_32bit_t*>(&frag_b)); |
| 1100 | dequant_data(frag_b_quant[k2][0][j * 2 + 1], |
| 1101 | reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2); |
| 1102 | } else if (is_a_8bit && b_type.size_bits() == 4 && has_zp) { |
| 1103 | int off = (threadIdx.x / 32) % 2 * 2 + j; |
| 1104 | int zp = (frag_qzp[k2][0] >> (off * 8)) & 0xF; |
| 1105 | dequant_data(frag_b_quant[k2][0][j * 2], |
| 1106 | reinterpret_cast<scalar_32bit_t*>(&frag_b), zp); |
| 1107 | zp = (frag_qzp[k2][0] >> (off * 8 + 4)) & 0xF; |
| 1108 | dequant_data(frag_b_quant[k2][0][j * 2 + 1], |
| 1109 | reinterpret_cast<scalar_32bit_t*>(&frag_b) + 2, zp); |
| 1110 | } else { |
| 1111 | reinterpret_cast<int2*>(&frag_b)[0] = |
| 1112 | reinterpret_cast<int2*>(&frag_b_quant[k2][j])[0]; |
| 1113 | reinterpret_cast<int2*>(&frag_b)[1] = |
| 1114 | reinterpret_cast<int2*>(&frag_b_quant[k2][j])[1]; |
| 1115 | } |
| 1116 | |
| 1117 | #pragma unroll |
| 1118 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1119 | mma<a_type_id, false, 32>( |
| 1120 | frag_a[k2][i], frag_b[0], |
| 1121 | (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][0]); |
| 1122 | mma<a_type_id, false, 32>( |
| 1123 | frag_a[k2][i], frag_b[1], |
| 1124 | (group_blocks == -1 ? frag_c : frag_c_tmp)[i][j][1]); |
| 1125 | } |
| 1126 | |
| 1127 | if constexpr (group_blocks != -1) { |
| 1128 | if (group_blocks == 2 || k == 1) { |
| 1129 | if constexpr (a_type == vllm::kS8) { |
| 1130 | int2 s_vals[2]; |
| 1131 | s_vals[0] = { |
| 1132 | (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[0], |
| 1133 | (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2][0])[1]}; |
| 1134 | s_vals[1] = { |
| 1135 | (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[0], |
| 1136 | (int)reinterpret_cast<uint16_t*>(&frag_s[k2][j * 2 + 1][0])[1]}; |
| 1137 | |
| 1138 | #pragma unroll |
| 1139 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1140 | #pragma unroll |
| 1141 | for (int g = 0; g < 4; g++) { |
| 1142 | int scale = reinterpret_cast<int*>(&s_vals[0])[g % 2]; |
| 1143 | *reinterpret_cast<int32_t*>(&frag_c[i][j][0][g]) += |
| 1144 | *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][0][g]) * |
| 1145 | scale; |
| 1146 | frag_c_tmp[i][j][0][g] = 0.0f; |
| 1147 | } |
| 1148 | |
| 1149 | #pragma unroll |
| 1150 | for (int g = 0; g < 4; g++) { |
| 1151 | int scale = reinterpret_cast<int*>(&s_vals[1])[g % 2]; |
| 1152 | *reinterpret_cast<int32_t*>(&frag_c[i][j][1][g]) += |
| 1153 | *reinterpret_cast<int32_t*>(&frag_c_tmp[i][j][1][g]) * |
| 1154 | scale; |
| 1155 | frag_c_tmp[i][j][1][g] = 0.0f; |
| 1156 | } |
| 1157 | } |
| 1158 | } else { |
| 1159 | float2 s_vals[2]; |
| 1160 | if constexpr (s_type_id != vllm::kFE8M0fnu.id()) { |
| 1161 | static_assert(a_type.size_bits() == 16 || |
| 1162 | s_type.size_bits() == 16); |
| 1163 | s_vals[0] = Cdtype::num22float2(frag_s[k2][j * 2][0]); |
| 1164 | s_vals[1] = Cdtype::num22float2(frag_s[k2][j * 2 + 1][0]); |
| 1165 | } else { |
| 1166 | int32_t* s_vals_int = reinterpret_cast<int32_t*>(&s_vals[0]); |
| 1167 | int32_t s_vals_e8m0 = |
| 1168 | *reinterpret_cast<int32_t*>(&frag_s[k2][j][0]); |
| 1169 | |
| 1170 | s_vals_int[0] = (s_vals_e8m0 & 0xFF) << 23; |
| 1171 | s_vals_int[1] = (s_vals_e8m0 & 0xFF00) << 15; |
| 1172 | s_vals_int[2] = (s_vals_e8m0 & 0xFF0000) << 7; |
| 1173 | s_vals_int[3] = (s_vals_e8m0 & 0xFF000000) >> 1; |
| 1174 | } |
| 1175 | |
| 1176 | #pragma unroll |
| 1177 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1178 | #pragma unroll |
| 1179 | for (int g = 0; g < 4; g++) { |
| 1180 | float scale = reinterpret_cast<float*>(&s_vals[0])[g % 2]; |
| 1181 | frag_c[i][j][0][g] += frag_c_tmp[i][j][0][g] * scale; |
| 1182 | frag_c_tmp[i][j][0][g] = 0.0f; |
| 1183 | } |
| 1184 | |
| 1185 | #pragma unroll |
| 1186 | for (int g = 0; g < 4; g++) { |
| 1187 | float scale = reinterpret_cast<float*>(&s_vals[1])[g % 2]; |
| 1188 | frag_c[i][j][1][g] += frag_c_tmp[i][j][1][g] * scale; |
| 1189 | frag_c_tmp[i][j][1][g] = 0.0f; |
| 1190 | } |
| 1191 | } |
| 1192 | } |
| 1193 | } |
| 1194 | } |
| 1195 | } |
| 1196 | }; |
| 1197 | |
| 1198 | // Since we slice across the k dimension of a tile in order to increase the |
| 1199 | // number of warps while keeping the n dimension of a tile reasonable, we have |
| 1200 | // multiple warps that accumulate their partial sums of the same output |
| 1201 | // location; which we have to reduce over in the end. We do in shared memory. |
| 1202 | auto thread_block_reduce = [&]() { |
| 1203 | constexpr int red_off = threads / b_sh_stride_threads / 2; |
| 1204 | if (red_off >= 1) { |
| 1205 | auto red_idx = threadIdx.x / b_sh_stride_threads; |
| 1206 | constexpr int red_sh_stride = |
| 1207 | b_sh_stride_threads * (is_a_8bit ? 2 : 4) * 2; |
| 1208 | constexpr int red_sh_delta = b_sh_stride_threads; |
| 1209 | int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) + |
| 1210 | (threadIdx.x % b_sh_stride_threads); |
| 1211 | |
| 1212 | // Parallel logarithmic shared memory reduction. We make sure to avoid any |
| 1213 | // unnecessary read or write iterations, e.g., for two warps we write only |
| 1214 | // once by warp 1 and read only once by warp 0. |
| 1215 | |
| 1216 | #pragma unroll |
| 1217 | for (int m_block = 0; m_block < thread_m_blocks; m_block++) { |
| 1218 | #pragma unroll |
| 1219 | for (int i = red_off; i > 0; i /= 2) { |
| 1220 | if (i <= red_idx && red_idx < 2 * i) { |
| 1221 | #pragma unroll |
| 1222 | for (int j = 0; j < (is_a_8bit ? 2 : 4) * 2; |
| 1223 | j += (m_block_size_8 ? 2 : 1)) { |
| 1224 | int red_sh_wr = |
| 1225 | red_sh_delta * j + (red_sh_rd - red_sh_stride * i); |
| 1226 | if (i < red_off) { |
| 1227 | float* c_rd = reinterpret_cast<float*>( |
| 1228 | &sh_red[red_sh_delta * j + red_sh_rd]); |
| 1229 | float* c_wr = reinterpret_cast<float*>(&sh_red[red_sh_wr]); |
| 1230 | #pragma unroll |
| 1231 | for (int k = 0; k < 4; k++) |
| 1232 | reinterpret_cast<FragC*>( |
| 1233 | frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j][k] += |
| 1234 | c_rd[k] + c_wr[k]; |
| 1235 | } |
| 1236 | sh_red[red_sh_wr] = reinterpret_cast<int4*>( |
| 1237 | &frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + j]; |
| 1238 | } |
| 1239 | } |
| 1240 | __syncthreads(); |
| 1241 | } |
| 1242 | if (red_idx == 0) { |
| 1243 | #pragma unroll |
| 1244 | for (int i = 0; i < (is_a_8bit ? 2 : 4) * 2; |
| 1245 | i += (m_block_size_8 ? 2 : 1)) { |
| 1246 | float* c_rd = |
| 1247 | reinterpret_cast<float*>(&sh_red[red_sh_delta * i + red_sh_rd]); |
| 1248 | #pragma unroll |
| 1249 | for (int j = 0; j < 4; j++) |
| 1250 | reinterpret_cast<FragC*>( |
| 1251 | frag_c)[(is_a_8bit ? 2 : 4) * 2 * m_block + i][j] += c_rd[j]; |
| 1252 | } |
| 1253 | } |
| 1254 | __syncthreads(); |
| 1255 | } |
| 1256 | } |
| 1257 | }; |
| 1258 | |
| 1259 | // Since multiple threadblocks may process parts of the same column slice, we |
| 1260 | // finally have to globally reduce over the results. As the striped |
| 1261 | // partitioning minimizes the number of such reductions and our outputs are |
| 1262 | // usually rather small, we perform this reduction serially in L2 cache. |
| 1263 | auto global_reduce_fp16 = [&](bool first = false, bool last = false) { |
| 1264 | // We are very careful here to reduce directly in the output buffer to |
| 1265 | // maximize L2 cache utilization in this step. To do this, we write out |
| 1266 | // results in FP16 (but still reduce with FP32 compute). |
| 1267 | constexpr int active_threads = 32 * tb_n_warps; |
| 1268 | if (threadIdx.x < active_threads) { |
| 1269 | int c_gl_stride = prob_n / 8; |
| 1270 | int c_gl_wr_delta_o = 8 * c_gl_stride * (is_a_8bit ? 2 : 1); |
| 1271 | int c_gl_wr_delta_i = 4 * (active_threads / 32); |
| 1272 | int c_gl_wr; |
| 1273 | if constexpr (m_block_size_8) { |
| 1274 | c_gl_wr = c_gl_stride * ((threadIdx.x % 4) * 2) + |
| 1275 | 4 * (threadIdx.x / 32) + (threadIdx.x % 32) / 8; |
| 1276 | c_gl_wr += (2 * thread_n_blocks) * slice_col; |
| 1277 | } else { |
| 1278 | c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) * (is_a_8bit ? 2 : 1) + |
| 1279 | 4 * (threadIdx.x / 32) + threadIdx.x % 4; |
| 1280 | c_gl_wr += (2 * thread_n_blocks) * slice_col * (is_a_8bit ? 2 : 1); |
| 1281 | } |
| 1282 | constexpr int c_sh_wr_delta = active_threads; |
| 1283 | auto c_sh_wr = threadIdx.x; |
| 1284 | |
| 1285 | int row = (threadIdx.x % 32) / 4; |
| 1286 | |
| 1287 | if (!first) { |
| 1288 | // Interestingly, doing direct global accesses here really seems to mess up |
| 1289 | // the compiler and lead to slowdowns, hence we also use async-copies even |
| 1290 | // though these fetches are not actually asynchronous. |
| 1291 | #pragma unroll |
| 1292 | for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { |
| 1293 | if constexpr (m_block_size_8) { |
| 1294 | cp_async4_pred(&sh_red[c_sh_wr + c_sh_wr_delta * i], |
| 1295 | &C[c_gl_wr + i * c_gl_stride + |
| 1296 | (threadIdx.x % 8) / 4 * c_gl_wr_delta_i], |
| 1297 | (threadIdx.x % 4) * 2 + i < prob_m); |
| 1298 | } else if constexpr (is_a_8bit) { |
| 1299 | int2* sh_red_int2 = reinterpret_cast<int2*>(sh_red); |
| 1300 | int2* c_int2 = reinterpret_cast<int2*>(C); |
| 1301 | cp_async2_ca_pred( |
| 1302 | &sh_red_int2[c_sh_wr + c_sh_wr_delta * i], |
| 1303 | &c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + |
| 1304 | c_gl_wr_delta_i * (i % 2)], |
| 1305 | i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); |
| 1306 | } else { |
| 1307 | cp_async4_pred( |
| 1308 | &sh_red[c_sh_wr + c_sh_wr_delta * i], |
| 1309 | &C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + |
| 1310 | c_gl_wr_delta_i * (i % 2)], |
| 1311 | i < (thread_m_blocks - 1) * 4 || 8 * (i / 2) + row < prob_m); |
| 1312 | } |
| 1313 | } |
| 1314 | cp_async_fence(); |
| 1315 | cp_async_wait<0>(); |
| 1316 | } |
| 1317 | |
| 1318 | #pragma unroll |
| 1319 | for (int i = 0; i < (m_block_size_8 ? 2 : thread_m_blocks * 4); i++) { |
| 1320 | bool mask = (!m_block_size_8) && (i < (thread_m_blocks - 1) * 4 || |
| 1321 | 8 * (i / 2) + row < prob_m) || |
| 1322 | (m_block_size_8) && ((threadIdx.x % 4) * 2 + i < prob_m); |
| 1323 | if (mask) { |
| 1324 | if (!first) { |
| 1325 | c_scalar_t* c_red_f16; |
| 1326 | if constexpr (is_a_8bit) { |
| 1327 | int2 tmp = |
| 1328 | reinterpret_cast<int2*>(sh_red)[c_sh_wr + i * c_sh_wr_delta]; |
| 1329 | c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp); |
| 1330 | } else { |
| 1331 | int4 tmp = sh_red[c_sh_wr + i * c_sh_wr_delta]; |
| 1332 | c_red_f16 = reinterpret_cast<c_scalar_t*>(&tmp); |
| 1333 | } |
| 1334 | #pragma unroll |
| 1335 | for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { |
| 1336 | int delta = 0; |
| 1337 | if constexpr (m_block_size_8) { |
| 1338 | delta = j % 2 == 1 ? -2 : 0; |
| 1339 | } |
| 1340 | reinterpret_cast<float*>( |
| 1341 | &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + |
| 1342 | (i % 4) + delta] += Cdtype::num2float(c_red_f16[j]); |
| 1343 | } |
| 1344 | } |
| 1345 | if (!last) { |
| 1346 | c_scalar_t c_f16[is_a_8bit ? 4 : 8]; |
| 1347 | #pragma unroll |
| 1348 | for (int j = 0; j < 2 * (is_a_8bit ? 2 : 4); j++) { |
| 1349 | int delta = 0; |
| 1350 | if constexpr (m_block_size_8) { |
| 1351 | delta = j % 2 == 1 ? -2 : 0; |
| 1352 | } |
| 1353 | c_f16[j] = Cdtype::float2num(reinterpret_cast<float*>( |
| 1354 | &frag_c)[(is_a_8bit ? 2 : 4) * 2 * 4 * (i / 4) + 4 * j + |
| 1355 | (i % 4) + delta]); |
| 1356 | } |
| 1357 | if constexpr (m_block_size_8) { |
| 1358 | C[c_gl_wr + i * c_gl_stride + |
| 1359 | (threadIdx.x % 8) / 4 * c_gl_wr_delta_i] = |
| 1360 | *reinterpret_cast<int4*>(c_f16); |
| 1361 | } else if constexpr (is_a_8bit) { |
| 1362 | int2* c_int2 = reinterpret_cast<int2*>(C); |
| 1363 | c_int2[c_gl_wr + c_gl_wr_delta_o * (i / 2) + |
| 1364 | c_gl_wr_delta_i * (i % 2)] = |
| 1365 | *reinterpret_cast<int2*>(c_f16); |
| 1366 | } else { |
| 1367 | C[c_gl_wr + c_gl_wr_delta_o * (i / 2) + |
| 1368 | c_gl_wr_delta_i * (i % 2)] = *reinterpret_cast<int4*>(c_f16); |
| 1369 | } |
| 1370 | } |
| 1371 | } |
| 1372 | } |
| 1373 | } |
| 1374 | }; |
| 1375 | |
| 1376 | // Globally reduce over threadblocks that compute the same column block. |
| 1377 | // We use a tmp C buffer to reduce in full fp32 precision. |
| 1378 | auto global_reduce_fp32 = [&](bool first = false, bool last = false) { |
| 1379 | constexpr int tb_m = thread_m_blocks * 16; |
| 1380 | constexpr int tb_n = thread_n_blocks * 16; |
| 1381 | |
| 1382 | constexpr int c_size = tb_m * tb_n * sizeof(float) / 16; |
| 1383 | |
| 1384 | constexpr int active_threads = 32 * tb_n_warps; |
| 1385 | bool is_th_active = threadIdx.x < active_threads; |
| 1386 | |
| 1387 | constexpr int num_floats = thread_m_blocks * (is_a_8bit ? 2 : 4) * 2 * 4; |
| 1388 | constexpr int th_size = num_floats * sizeof(float) / 16; |
| 1389 | |
| 1390 | int c_cur_offset = locks_off * c_size; |
| 1391 | |
| 1392 | if (!is_th_active) { |
| 1393 | return; |
| 1394 | } |
| 1395 | |
| 1396 | if (!first) { |
| 1397 | float* frag_c_ptr = reinterpret_cast<float*>(&frag_c); |
| 1398 | #pragma unroll |
| 1399 | for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { |
| 1400 | sh_red[threadIdx.x] = |
| 1401 | C_tmp[c_cur_offset + active_threads * k + threadIdx.x]; |
| 1402 | |
| 1403 | float* sh_c_ptr = reinterpret_cast<float*>(&sh_red[threadIdx.x]); |
| 1404 | #pragma unroll |
| 1405 | for (int f = 0; f < 4; f++) { |
| 1406 | frag_c_ptr[k * 4 + f] += sh_c_ptr[f]; |
| 1407 | } |
| 1408 | } |
| 1409 | } |
| 1410 | |
| 1411 | if (!last) { |
| 1412 | int4* frag_c_ptr = reinterpret_cast<int4*>(&frag_c); |
| 1413 | #pragma unroll |
| 1414 | for (int k = 0; k < th_size; k += (m_block_size_8 ? 2 : 1)) { |
| 1415 | C_tmp[c_cur_offset + active_threads * k + threadIdx.x] = frag_c_ptr[k]; |
| 1416 | } |
| 1417 | } |
| 1418 | }; |
| 1419 | |
| 1420 | // Write out the reduce final result in the correct layout. We only actually |
| 1421 | // reshuffle matrix fragments in this step, the reduction above is performed |
| 1422 | // in fragment layout. |
| 1423 | auto write_result = [&](bool last) { |
| 1424 | int c_gl_stride = prob_n / 8; |
| 1425 | constexpr int c_sh_stride = 2 * thread_n_blocks + 1; |
| 1426 | int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks)); |
| 1427 | constexpr int c_sh_rd_delta = |
| 1428 | c_sh_stride * (threads / (2 * thread_n_blocks)); |
| 1429 | |
| 1430 | int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) + |
| 1431 | (threadIdx.x % (2 * thread_n_blocks)); |
| 1432 | c_gl_wr += (2 * thread_n_blocks) * slice_col; |
| 1433 | int c_sh_wr; |
| 1434 | if constexpr (m_block_size_8) { |
| 1435 | c_sh_wr = (8 * c_sh_stride) * ((threadIdx.x % 32) % 4 * 2) + |
| 1436 | (threadIdx.x % 32) / 4; |
| 1437 | c_sh_wr += 64 * (threadIdx.x / 32); |
| 1438 | } else { |
| 1439 | c_sh_wr = |
| 1440 | (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4; |
| 1441 | c_sh_wr += (is_a_8bit ? 16 : 32) * (threadIdx.x / 32); |
| 1442 | } |
| 1443 | |
| 1444 | int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) + |
| 1445 | (threadIdx.x % (2 * thread_n_blocks)); |
| 1446 | |
| 1447 | int c_gl_wr_end = c_gl_stride * prob_m; |
| 1448 | // We first reorder in shared memory to guarantee the most efficient final |
| 1449 | // global write patterns |
| 1450 | auto write = [&](int idx, float c0, float c1, FragS& s, FragS& b_bias) { |
| 1451 | if constexpr (b_type == vllm::kFE2M1f && s_type == vllm::kFE4M3fn) { |
| 1452 | c0 *= global_scale_f32; |
| 1453 | c1 *= global_scale_f32; |
| 1454 | } |
| 1455 | c_scalar_t2 res = |
| 1456 | Cdtype::nums2num2(Cdtype::float2num(c0), Cdtype::float2num(c1)); |
| 1457 | |
| 1458 | // For per-column quantization we finally apply the scale here (only for |
| 1459 | // 4-bit) |
| 1460 | if constexpr (group_blocks == -1 && !is_a_8bit && |
| 1461 | b_type.size_bits() == 4 && |
| 1462 | (has_zp && dequant_skip_flop || !has_zp)) { |
| 1463 | c_scalar_t2 tmp_scale = s[0]; |
| 1464 | if constexpr (m_block_size_8) { |
| 1465 | tmp_scale = Cdtype::num2num2( |
| 1466 | reinterpret_cast<scalar_t*>(&s[0])[(threadIdx.x % 8) / 4]); |
| 1467 | } |
| 1468 | res = __hmul2(res, tmp_scale); |
| 1469 | } |
| 1470 | if (has_bias && last) { |
| 1471 | c_scalar_t2 tmp_bias = b_bias[0]; |
| 1472 | if constexpr (m_block_size_8) { |
| 1473 | tmp_bias = Cdtype::num2num2( |
| 1474 | reinterpret_cast<scalar_t*>(&b_bias[0])[(threadIdx.x % 8) / 4]); |
| 1475 | } |
| 1476 | res = __hadd2(res, tmp_bias); |
| 1477 | } |
| 1478 | |
| 1479 | if constexpr (m_block_size_8) { |
| 1480 | ((c_scalar_t*)sh_red)[idx] = res.x; |
| 1481 | ((c_scalar_t*)sh_red)[idx + 8 * c_sh_stride] = res.y; |
| 1482 | } else { |
| 1483 | ((c_scalar_t2*)sh_red)[idx] = res; |
| 1484 | } |
| 1485 | }; |
| 1486 | |
| 1487 | if (threadIdx.x / 32 < tb_n_warps) { |
| 1488 | #pragma unroll |
| 1489 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1490 | #pragma unroll |
| 1491 | for (int j = 0; j < (is_a_8bit ? 2 : 4); j++) { |
| 1492 | if constexpr (m_block_size_8) { |
| 1493 | int wr = c_sh_wr + 16 * j; |
| 1494 | write(wr, frag_c[i][j][0][0], frag_c[i][j][0][1], |
| 1495 | frag_s[j / 2][2 * (j % 2) + 0], |
| 1496 | frag_bias[j / 2][2 * (j % 2) + 0]); |
| 1497 | write(wr + 8, frag_c[i][j][0][2], frag_c[i][j][0][3], |
| 1498 | frag_s[j / 2][2 * (j % 2) + 1], |
| 1499 | frag_bias[j / 2][2 * (j % 2) + 1]); |
| 1500 | } else { |
| 1501 | int wr = c_sh_wr + 8 * j; |
| 1502 | write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0], |
| 1503 | frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0], |
| 1504 | frag_bias[j / 2][2 * (j % 2) + 0]); |
| 1505 | write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2], |
| 1506 | frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0], |
| 1507 | frag_bias[j / 2][2 * (j % 2) + 0]); |
| 1508 | write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0], |
| 1509 | frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1], |
| 1510 | frag_bias[j / 2][2 * (j % 2) + 1]); |
| 1511 | write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2], |
| 1512 | frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1], |
| 1513 | frag_bias[j / 2][2 * (j % 2) + 1]); |
| 1514 | } |
| 1515 | } |
| 1516 | c_sh_wr += 16 * (4 * c_sh_stride); |
| 1517 | } |
| 1518 | } |
| 1519 | __syncthreads(); |
| 1520 | |
| 1521 | #pragma unroll |
| 1522 | for (int i = 0; |
| 1523 | i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks)); |
| 1524 | i++) { |
| 1525 | if (c_gl_wr < c_gl_wr_end) { |
| 1526 | if (use_atomic_add && slice_count > 1) { |
| 1527 | c_scalar_t2* C_half2 = reinterpret_cast<c_scalar_t2*>(&C[c_gl_wr]); |
| 1528 | c_scalar_t2* sh_red_half2 = |
| 1529 | reinterpret_cast<c_scalar_t2*>(&sh_red[c_sh_rd]); |
| 1530 | #pragma unroll |
| 1531 | for (int a = 0; a < 4; a++) { |
| 1532 | atomicAdd(&C_half2[a], sh_red_half2[a]); |
| 1533 | } |
| 1534 | } else { |
| 1535 | C[c_gl_wr] = sh_red[c_sh_rd]; |
| 1536 | } |
| 1537 | c_gl_wr += c_gl_wr_delta; |
| 1538 | c_sh_rd += c_sh_rd_delta; |
| 1539 | } |
| 1540 | } |
| 1541 | __syncthreads(); |
| 1542 | }; |
| 1543 | |
| 1544 | // Start global fetch and register load pipelines. |
| 1545 | auto start_pipes = [&]() { |
| 1546 | |
| 1547 | #pragma unroll |
| 1548 | for (int i = 0; i < stages - 1; i++) { |
| 1549 | if constexpr (has_zp && !is_zp_float && group_blocks == -1) { |
| 1550 | if (i == 0) { |
| 1551 | fetch_col_zp_to_shared(); |
| 1552 | if constexpr (!dequant_skip_flop) { |
| 1553 | fetch_col_scale_to_shared(); |
| 1554 | } |
| 1555 | } |
| 1556 | } |
| 1557 | fetch_to_shared(i, i, i < slice_iters); |
| 1558 | } |
| 1559 | |
| 1560 | zero_accums(); |
| 1561 | wait_for_stage(); |
| 1562 | fetch_to_registers(0, 0); |
| 1563 | fetch_scales_to_registers(0, 0); |
| 1564 | fetch_zp_to_registers(0, 0); |
| 1565 | a_gl_rd += a_gl_rd_delta_o * (stages - 1); |
| 1566 | }; |
| 1567 | if (slice_iters) { |
| 1568 | start_pipes(); |
| 1569 | } |
| 1570 | |
| 1571 | // Main loop. |
| 1572 | while (slice_iters) { |
| 1573 | // We unroll over both the global fetch and the register load pipeline to |
| 1574 | // ensure all shared memory accesses are static. Note that both pipelines |
| 1575 | // have even length meaning that the next iteration will always start at |
| 1576 | // index 0. |
| 1577 | |
| 1578 | #pragma unroll |
| 1579 | for (int pipe = 0; pipe < stages;) { |
| 1580 | #pragma unroll |
| 1581 | for (int k = 0; k < b_sh_wr_iters; k++) { |
| 1582 | fetch_to_registers(k + 1, pipe % stages); |
| 1583 | fetch_scales_to_registers(k + 1, pipe); |
| 1584 | fetch_zp_to_registers(k + 1, pipe); |
| 1585 | if (k == b_sh_wr_iters - 2) { |
| 1586 | fetch_to_shared((pipe + stages - 1) % stages, pipe, |
| 1587 | slice_iters >= stages); |
| 1588 | pipe++; |
| 1589 | wait_for_stage(); |
| 1590 | } |
| 1591 | |
| 1592 | if constexpr (!is_a_8bit) { |
| 1593 | matmul(k, pipe - (k >= b_sh_wr_iters - 2 ? 1 : 0)); |
| 1594 | } else { |
| 1595 | static_assert(group_blocks != 0 && group_blocks != 1); |
| 1596 | matmul_a8(k); |
| 1597 | } |
| 1598 | } |
| 1599 | slice_iters--; |
| 1600 | if (slice_iters == 0) { |
| 1601 | break; |
| 1602 | } |
| 1603 | } |
| 1604 | |
| 1605 | a_gl_rd += a_gl_rd_delta_o * stages; |
| 1606 | |
| 1607 | // Process results and, if necessary, proceed to the next column slice. |
| 1608 | // While this pattern may not be the most readable, other ways of writing |
| 1609 | // the loop seemed to noticeably worse performance after compilation. |
| 1610 | if (slice_iters == 0) { |
| 1611 | // convert fp16 accum to fp32 for reduction |
| 1612 | if constexpr (use_fp16_accum) { |
| 1613 | #pragma unroll |
| 1614 | for (int i = 0; i < (thread_m_blocks * (is_a_8bit ? 2 : 4) * 2); i++) { |
| 1615 | float* frag_c_part_float = reinterpret_cast<float*>(frag_c) + i * 4; |
| 1616 | scalar_t* frag_c_part_half = |
| 1617 | reinterpret_cast<scalar_t*>(frag_c_part_float); |
| 1618 | |
| 1619 | #pragma unroll |
| 1620 | for (int i = 3; i >= 0; i--) { |
| 1621 | frag_c_part_float[i] = Cdtype::num2float(frag_c_part_half[i]); |
| 1622 | } |
| 1623 | } |
| 1624 | } |
| 1625 | |
| 1626 | if constexpr (is_a_8bit) { |
| 1627 | float frag_a_s[2 * thread_m_blocks]; |
| 1628 | |
| 1629 | for (int i = 0; i < 2 * thread_m_blocks; i++) |
| 1630 | frag_a_s[i] = sh_a_s[i * 8 + (threadIdx.x % 32) / 4]; |
| 1631 | |
| 1632 | #pragma unroll |
| 1633 | for (int j = 0; j < 2; j++) { |
| 1634 | #pragma unroll |
| 1635 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1636 | #pragma unroll |
| 1637 | for (int g = 0; g < 4; g++) { |
| 1638 | float c_val = frag_c[i][j][0][g]; |
| 1639 | |
| 1640 | if constexpr (a_type == vllm::kS8) { |
| 1641 | c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val)); |
| 1642 | } |
| 1643 | float s_val = frag_a_s[i * 2 + g / 2]; |
| 1644 | frag_c[i][j][0][g] = c_val * s_val; |
| 1645 | } |
| 1646 | #pragma unroll |
| 1647 | for (int g = 0; g < 4; g++) { |
| 1648 | float c_val = frag_c[i][j][1][g]; |
| 1649 | |
| 1650 | if constexpr (a_type == vllm::kS8) { |
| 1651 | c_val = __int2float_rn(*reinterpret_cast<int32_t*>(&c_val)); |
| 1652 | } |
| 1653 | float s_val = frag_a_s[i * 2 + g / 2]; |
| 1654 | frag_c[i][j][1][g] = c_val * s_val; |
| 1655 | } |
| 1656 | } |
| 1657 | } |
| 1658 | } |
| 1659 | |
| 1660 | cp_async_wait<0>(); |
| 1661 | bool last = slice_idx == slice_count - 1; |
| 1662 | // For per-column scales, we only fetch them here in the final step before |
| 1663 | // write-out |
| 1664 | if constexpr (group_blocks == -1 && |
| 1665 | (has_zp && dequant_skip_flop || !has_zp)) { |
| 1666 | if (b_type.size_bits() == 8 || (last || use_atomic_add) || is_a_8bit) { |
| 1667 | if (s_sh_wr_pred) { |
| 1668 | cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]); |
| 1669 | } |
| 1670 | cp_async_fence(); |
| 1671 | } |
| 1672 | } |
| 1673 | |
| 1674 | thread_block_reduce(); |
| 1675 | |
| 1676 | if (has_bias && last) { |
| 1677 | __syncthreads(); |
| 1678 | cp_async4_pred(&sh_bias[bias_sh_wr], &b_bias_ptr[bias_gl_rd], |
| 1679 | threadIdx.x < 16 * thread_n_blocks / 8); |
| 1680 | cp_async_fence(); |
| 1681 | } |
| 1682 | |
| 1683 | if constexpr (group_blocks == -1 && |
| 1684 | (has_zp && dequant_skip_flop || !has_zp || is_a_8bit)) { |
| 1685 | if constexpr (is_a_8bit) { |
| 1686 | cp_async_wait<0>(); |
| 1687 | __syncthreads(); |
| 1688 | if (threadIdx.x / 32 < tb_n_warps) { |
| 1689 | reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0]; |
| 1690 | } |
| 1691 | } else if (b_type.size_bits() == 8 || (last || use_atomic_add)) { |
| 1692 | cp_async_wait<0>(); |
| 1693 | __syncthreads(); |
| 1694 | if (threadIdx.x / 32 < tb_n_warps) { |
| 1695 | reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0]; |
| 1696 | reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4]; |
| 1697 | if constexpr (m_block_size_8) { |
| 1698 | int idx = (threadIdx.x / 4) % 2; |
| 1699 | c_scalar_t2* frag_s_half2 = |
| 1700 | reinterpret_cast<c_scalar_t2*>(frag_s); |
| 1701 | #pragma unroll |
| 1702 | for (int i = 0; i < 8; i++) { |
| 1703 | frag_s_half2[i] = Cdtype::num2num2( |
| 1704 | reinterpret_cast<c_scalar_t*>(&frag_s_half2[i])[idx]); |
| 1705 | } |
| 1706 | } |
| 1707 | } |
| 1708 | } |
| 1709 | } |
| 1710 | |
| 1711 | // For 8-bit channelwise, we apply the scale before the global reduction |
| 1712 | // that converts the fp32 results to fp16 (so that we avoid possible |
| 1713 | // overflow in fp16) |
| 1714 | if constexpr (group_blocks == -1 && is_a_8bit) { |
| 1715 | #pragma unroll |
| 1716 | for (int j = 0; j < 2; j++) { |
| 1717 | float2 aa[2]; |
| 1718 | aa[0] = Cdtype::num22float2(frag_s[0][j * 2][0]); |
| 1719 | aa[1] = Cdtype::num22float2(frag_s[0][j * 2 + 1][0]); |
| 1720 | |
| 1721 | #pragma unroll |
| 1722 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1723 | #pragma unroll |
| 1724 | for (int g = 0; g < 4; g++) { |
| 1725 | float scale = reinterpret_cast<float*>(&aa[0])[g % 2]; |
| 1726 | frag_c[i][j][0][g] *= scale; |
| 1727 | } |
| 1728 | |
| 1729 | #pragma unroll |
| 1730 | for (int g = 0; g < 4; g++) { |
| 1731 | float scale = reinterpret_cast<float*>(&aa[1])[g % 2]; |
| 1732 | frag_c[i][j][1][g] *= scale; |
| 1733 | } |
| 1734 | } |
| 1735 | } |
| 1736 | } else if (group_blocks == -1 && b_type.size_bits() == 8 && |
| 1737 | (has_zp && dequant_skip_flop || !has_zp)) { |
| 1738 | if (threadIdx.x / 32 < tb_n_warps) { |
| 1739 | #pragma unroll |
| 1740 | for (int i = 0; i < thread_m_blocks; i++) { |
| 1741 | #pragma unroll |
| 1742 | for (int j = 0; j < 4; j++) { |
| 1743 | scale_float<c_type_id>( |
| 1744 | reinterpret_cast<float*>(&frag_c[i][j][0][0]), |
| 1745 | frag_s[j / 2][2 * (j % 2) + 0]); |
| 1746 | scale_float<c_type_id>( |
| 1747 | reinterpret_cast<float*>(&frag_c[i][j][0][2]), |
| 1748 | frag_s[j / 2][2 * (j % 2) + (m_block_size_8 ? 1 : 0)]); |
| 1749 | |
| 1750 | if constexpr (!m_block_size_8) { |
| 1751 | scale_float<c_type_id>( |
| 1752 | reinterpret_cast<float*>(&frag_c[i][j][1][0]), |
| 1753 | frag_s[j / 2][2 * (j % 2) + 1]); |
| 1754 | scale_float<c_type_id>( |
| 1755 | reinterpret_cast<float*>(&frag_c[i][j][1][2]), |
| 1756 | frag_s[j / 2][2 * (j % 2) + 1]); |
| 1757 | } |
| 1758 | } |
| 1759 | } |
| 1760 | } |
| 1761 | } |
| 1762 | |
| 1763 | if (slice_count > 1 && !use_atomic_add) { |
| 1764 | // only globally reduce if there is more than one block in a slice |
| 1765 | barrier_acquire(&locks[locks_off], slice_idx); |
| 1766 | if (use_fp32_reduce) { |
| 1767 | global_reduce_fp32(slice_idx == 0, last); |
| 1768 | } else { |
| 1769 | global_reduce_fp16(slice_idx == 0, last); |
| 1770 | } |
| 1771 | barrier_release(&locks[locks_off], last); |
| 1772 | } |
| 1773 | |
| 1774 | if (has_bias && last) { |
| 1775 | cp_async_wait<0>(); |
| 1776 | __syncthreads(); |
| 1777 | reinterpret_cast<int4*>(&frag_bias)[0] = sh_bias[bias_sh_rd]; |
| 1778 | if constexpr (!is_a_8bit) |
| 1779 | reinterpret_cast<int4*>(&frag_bias)[1] = sh_bias[bias_sh_rd + 4]; |
| 1780 | __syncthreads(); |
| 1781 | } |
| 1782 | |
| 1783 | if (use_atomic_add && slice_count > 1 && slice_idx != 0) |
| 1784 | wait_negative_and_add(&locks[locks_off]); |
| 1785 | if (last || use_atomic_add) |
| 1786 | // only the last block in a slice actually writes the result |
| 1787 | write_result(last); |
| 1788 | slice_row = 0; |
| 1789 | if (!in_part2) { |
| 1790 | slice_col_par += gridDim.x; |
| 1791 | } else { |
| 1792 | slice_col_par++; |
| 1793 | slice_col++; |
| 1794 | } |
| 1795 | is_first_matmul_in_slice = true; |
| 1796 | init_slice(); |
| 1797 | |
| 1798 | if (slice_iters) { |
| 1799 | a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) + |
| 1800 | (threadIdx.x % a_gl_rd_delta_o); |
| 1801 | a_gl_rd += a_gl_rd_delta_o * slice_row; |
| 1802 | b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride) + |
| 1803 | (threadIdx.x % b_sh_stride); |
| 1804 | b_gl_rd += b_sh_stride * slice_col + b_gl_rd_delta_o * slice_row; |
| 1805 | |
| 1806 | bias_gl_rd = (thread_n_blocks * 16 / 8) * slice_col + threadIdx.x; |
| 1807 | if constexpr (group_blocks == -1) { |
| 1808 | s_gl_rd = s_sh_stride * slice_col + threadIdx.x; |
| 1809 | zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x; |
| 1810 | } else if constexpr (group_blocks >= thread_k_blocks) { |
| 1811 | s_gl_rd = |
| 1812 | s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + |
| 1813 | s_sh_stride * slice_col + threadIdx.x; |
| 1814 | zp_gl_rd = |
| 1815 | zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) + |
| 1816 | zp_sh_stride * slice_col + threadIdx.x; |
| 1817 | } else { |
| 1818 | s_gl_rd = |
| 1819 | s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + |
| 1820 | threadIdx.x / s_sh_stride) + |
| 1821 | s_sh_stride * slice_col + threadIdx.x % s_sh_stride; |
| 1822 | zp_gl_rd = |
| 1823 | zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks + |
| 1824 | threadIdx.x / zp_sh_stride) + |
| 1825 | zp_sh_stride * slice_col + threadIdx.x % zp_sh_stride; |
| 1826 | } |
| 1827 | start_pipes(); |
| 1828 | } |
| 1829 | } |
| 1830 | } |
| 1831 | } |
| 1832 | |
| 1833 | } // namespace MARLIN_NAMESPACE_NAME |
| 1834 | |
| 1835 | #endif |