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