@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a17cef70ea

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 11:00 UTC

Removed unused marlin cuda code (#17684)

Signed-off-by: mgoin <mgoin64@gmail.com>

Michael Goin committed May 6, 2025 at 20:59 UTC a17cef70eacba81577e1eaa91f2b5dd18624e5d5

9 files changed -3666

csrc/moe/marlin_kernels/marlin_moe_kernel.h deleted

-1616

@@ -1,1616 +0,0 @@

1 -#pragma once

2 -

3 -#include <torch/all.h>

4 -

5 -#include <ATen/cuda/CUDAContext.h>

6 -#include <c10/cuda/CUDAGuard.h>

7 -#include <cuda.h>

8 -#include <cuda_fp16.h>

9 -#include <cuda_runtime.h>

10 -

11 -#include <iostream>

12 -

13 -#include "core/scalar_type.hpp"

14 -

15 -namespace marlin_moe {

16 -

17 -constexpr int ceildiv(int a, int b) { return (a + b - 1) / b; }

18 -

19 -#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800

20 -

21 -// Instances of `Vec` are used to organize groups of >>registers<<, as needed

22 -// for instance as inputs to tensor core operations. Consequently, all

23 -// corresponding index accesses must be compile-time constants, which is why we

24 -// extensively use `#pragma unroll` throughout the kernel code to guarantee

25 -// this.

26 -template <typename T, int n>

27 -struct Vec {

28 - T elems[n];

29 - __device__ T& operator[](int i) { return elems[i]; }

30 -};

31 -

32 -using I4 = Vec<int, 4>;

33 -

34 -// Matrix fragments for tensor core instructions; their precise layout is

35 -// documented here:

36 -// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-fragments-for-mma-m16n8k16-with-floating-point-type

37 -using FragA = Vec<half2, 4>;

38 -using FragB = Vec<half2, 2>;

39 -using FragC = Vec<float, 4>;

40 -using FragS = Vec<half2, 1>; // quantization scales

41 -using FragZP = Vec<half2, 4>;

42 -

43 -// Predicated asynchronous global->shared copy; used for inputs A where we apply

44 -// predication to handle batchsizes that are not multiples of 16.

45 -__device__ inline void cp_async4_pred(void* smem_ptr, const void* glob_ptr,

46 - bool pred = true) {

47 - const int BYTES = 16;

48 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));

49 - asm volatile(

50 - "{\n"

51 - " .reg .pred p;\n"

52 - " setp.ne.b32 p, %0, 0;\n"

53 - " @p cp.async.cg.shared.global [%1], [%2], %3;\n"

54 - "}\n" ::"r"((int)pred),

55 - "r"(smem), "l"(glob_ptr), "n"(BYTES));

56 -}

57 -

58 -// Asynchronous global->shared copy

59 -__device__ inline void cp_async4(void* smem_ptr, const void* glob_ptr) {

60 - const int BYTES = 16;

61 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));

62 - asm volatile(

63 - "{\n"

64 - " cp.async.cg.shared.global [%0], [%1], %2;\n"

65 - "}\n" ::"r"(smem),

66 - "l"(glob_ptr), "n"(BYTES));

67 -}

68 -

69 -// Async copy fence.

70 -__device__ inline void cp_async_fence() {

71 - asm volatile("cp.async.commit_group;\n" ::);

72 -}

73 -

74 -// Wait until at most `n` async copy stages are still pending.

75 -template <int n>

76 -__device__ inline void cp_async_wait() {

77 - asm volatile("cp.async.wait_group %0;\n" ::"n"(n));

78 -}

79 -

80 -// m16n8k16 tensor core mma instruction with fp16 inputs and fp32

81 -// output/accumulation.

82 -__device__ inline void mma(const FragA& a_frag, const FragB& frag_b,

83 - FragC& frag_c) {

84 - const uint32_t* a = reinterpret_cast<const uint32_t*>(&a_frag);

85 - const uint32_t* b = reinterpret_cast<const uint32_t*>(&frag_b);

86 - float* c = reinterpret_cast<float*>(&frag_c);

87 - asm volatile(

88 - "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 "

89 - "{%0,%1,%2,%3}, {%4,%5,%6,%7}, {%8,%9}, {%10,%11,%12,%13};\n"

90 - : "=f"(c[0]), "=f"(c[1]), "=f"(c[2]), "=f"(c[3])

91 - : "r"(a[0]), "r"(a[1]), "r"(a[2]), "r"(a[3]), "r"(b[0]), "r"(b[1]),

92 - "f"(c[0]), "f"(c[1]), "f"(c[2]), "f"(c[3]));

93 -}

94 -

95 -// Instruction for loading a full 16x16 matrix fragment of operand A from shared

96 -// memory, directly in tensor core layout.

97 -__device__ inline void ldsm4(FragA& frag_a, const void* smem_ptr) {

98 - uint32_t* a = reinterpret_cast<uint32_t*>(&frag_a);

99 - uint32_t smem = static_cast<uint32_t>(__cvta_generic_to_shared(smem_ptr));

100 - asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0,%1,%2,%3}, [%4];\n"

101 - : "=r"(a[0]), "=r"(a[1]), "=r"(a[2]), "=r"(a[3])

102 - : "r"(smem));

103 -}

104 -

105 -// Lookup-table based 3-input logical operation; explicitly used for

106 -// dequantization as the compiler does not seem to automatically recognize it in

107 -// all cases.

108 -template <int lut>

109 -__device__ inline int lop3(int a, int b, int c) {

110 - int res;

111 - asm volatile("lop3.b32 %0, %1, %2, %3, %4;\n"

112 - : "=r"(res)

113 - : "r"(a), "r"(b), "r"(c), "n"(lut));

114 - return res;

115 -}

116 -

117 -// Constructs destination register by taking bytes from 2 sources (based on

118 -// mask)

119 -template <int start_byte, int mask>

120 -__device__ inline uint32_t prmt(uint32_t a) {

121 - uint32_t res;

122 - asm volatile("prmt.b32 %0, %1, %2, %3;\n"

123 - : "=r"(res)

124 - : "r"(a), "n"(start_byte), "n"(mask));

125 - return res;

126 -}

127 -

128 -template <vllm::ScalarTypeId w_type_id>

129 -__device__ inline FragB dequant(int q);

130 -

131 -// Efficiently dequantize 4bit values packed in an int32 value into a full

132 -// B-fragment of 4 fp16 values. We mostly follow the strategy in the link below,

133 -// with some small changes:

134 -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L215-L287

135 -template <>

136 -__device__ inline FragB dequant<vllm::kU4B8.id()>(int q) {

137 - const int LO = 0x000f000f;

138 - const int HI = 0x00f000f0;

139 - const int EX = 0x64006400;

140 - // Guarantee that the `(a & b) | c` operations are LOP3s.

141 - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);

142 - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);

143 - // We want signed int4 outputs, hence we fuse the `-8` symmetric zero point

144 - // directly into `SUB` and `ADD`.

145 - const int SUB = 0x64086408;

146 - const int MUL = 0x2c002c00;

147 - const int ADD = 0xd480d480;

148 - FragB frag_b;

149 - frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),

150 - *reinterpret_cast<const half2*>(&SUB));

151 - frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),

152 - *reinterpret_cast<const half2*>(&MUL),

153 - *reinterpret_cast<const half2*>(&ADD));

154 - return frag_b;

155 -}

156 -

157 -// Fast Int8ToFp16: Efficiently dequantize 8bit int values to fp16

158 -// Reference:

159 -// https://github.com/NVIDIA/FasterTransformer/blob/release/v5.3_tag/src/fastertransformer/cutlass_extensions/include/cutlass_extensions/interleaved_numeric_conversion.h#L53-L85

160 -template <>

161 -__device__ inline FragB dequant<vllm::kU8B128.id()>(int q) {

162 - static constexpr uint32_t mask_for_elt_01 = 0x5250;

163 - static constexpr uint32_t mask_for_elt_23 = 0x5351;

164 - static constexpr uint32_t start_byte_for_fp16 = 0x64646464;

165 -

166 - uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);

167 - uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);

168 -

169 - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64806480;

170 -

171 - FragB frag_b;

172 - frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),

173 - *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));

174 - frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),

175 - *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));

176 - return frag_b;

177 -}

178 -

179 -template <>

180 -__device__ inline FragB dequant<vllm::kU4.id()>(int q) {

181 - const int LO = 0x000f000f;

182 - const int HI = 0x00f000f0;

183 - const int EX = 0x64006400;

184 - // Guarantee that the `(a & b) | c` operations are LOP3s.

185 - int lo = lop3<(0xf0 & 0xcc) | 0xaa>(q, LO, EX);

186 - int hi = lop3<(0xf0 & 0xcc) | 0xaa>(q, HI, EX);

187 -

188 - const int SUB = 0x64006400;

189 - const int MUL = 0x2c002c00;

190 - const int ADD = 0xd400d400;

191 - FragB frag_b;

192 - frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),

193 - *reinterpret_cast<const half2*>(&SUB));

194 - frag_b[1] = __hfma2(*reinterpret_cast<half2*>(&hi),

195 - *reinterpret_cast<const half2*>(&MUL),

196 - *reinterpret_cast<const half2*>(&ADD));

197 - return frag_b;

198 -}

199 -

200 -template <>

201 -__device__ inline FragB dequant<vllm::kU8.id()>(int q) {

202 - static constexpr uint32_t mask_for_elt_01 = 0x5250;

203 - static constexpr uint32_t mask_for_elt_23 = 0x5351;

204 - static constexpr uint32_t start_byte_for_fp16 = 0x64646464;

205 -

206 - uint32_t lo = prmt<start_byte_for_fp16, mask_for_elt_01>(q);

207 - uint32_t hi = prmt<start_byte_for_fp16, mask_for_elt_23>(q);

208 -

209 - static constexpr uint32_t I8s_TO_F16s_MAGIC_NUM = 0x64006400;

210 -

211 - FragB frag_b;

212 - frag_b[0] = __hsub2(*reinterpret_cast<half2*>(&lo),

213 - *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));

214 - frag_b[1] = __hsub2(*reinterpret_cast<half2*>(&hi),

215 - *reinterpret_cast<const half2*>(&I8s_TO_F16s_MAGIC_NUM));

216 - return frag_b;

217 -}

218 -

219 -// Multiply dequantized values by the corresponding quantization scale; used

220 -// only for grouped quantization.

221 -__device__ inline void scale(FragB& frag_b, FragS& frag_s, int i) {

222 - half2 s = __half2half2(reinterpret_cast<__half*>(&frag_s)[i]);

223 - frag_b[0] = __hmul2(frag_b[0], s);

224 - frag_b[1] = __hmul2(frag_b[1], s);

225 -}

226 -

227 -__device__ inline void sub_zp(FragB& frag_b, half2& frag_zp, int i) {

228 - half2 zp = __half2half2(reinterpret_cast<__half*>(&frag_zp)[i]);

229 - frag_b[0] = __hsub2(frag_b[0], zp);

230 - frag_b[1] = __hsub2(frag_b[1], zp);

231 -}

232 -

233 -// Same as above, but for act_order (each K is multiplied individually)

234 -__device__ inline void scale4(FragB& frag_b, FragS& frag_s_1, FragS& frag_s_2,

235 - FragS& frag_s_3, FragS& frag_s_4, int i) {

236 - __half2 s_val_1_2;

237 - s_val_1_2.x = reinterpret_cast<__half*>(&frag_s_1)[i];

238 - s_val_1_2.y = reinterpret_cast<__half*>(&frag_s_2)[i];

239 -

240 - __half2 s_val_3_4;

241 - s_val_3_4.x = reinterpret_cast<__half*>(&frag_s_3)[i];

242 - s_val_3_4.y = reinterpret_cast<__half*>(&frag_s_4)[i];

243 -

244 - frag_b[0] = __hmul2(frag_b[0], s_val_1_2);

245 - frag_b[1] = __hmul2(frag_b[1], s_val_3_4);

246 -}

247 -

248 -// Given 2 floats multiply by 2 scales (halves)

249 -__device__ inline void scale_float(float* c, FragS& s) {

250 - __half* s_ptr = reinterpret_cast<__half*>(&s);

251 - c[0] = __fmul_rn(c[0], __half2float(s_ptr[0]));

252 - c[1] = __fmul_rn(c[1], __half2float(s_ptr[1]));

253 -}

254 -

255 -// Wait until barrier reaches `count`, then lock for current threadblock.

256 -__device__ inline void barrier_acquire(int* lock, int count) {

257 - if (threadIdx.x == 0) {

258 - int state = -1;

259 - do

260 - // Guarantee that subsequent writes by this threadblock will be visible

261 - // globally.

262 - asm volatile("ld.global.acquire.gpu.b32 %0, [%1];\n"

263 - : "=r"(state)

264 - : "l"(lock));

265 - while (state != count);

266 - }

267 - __syncthreads();

268 -}

269 -

270 -// Release barrier and increment visitation count.

271 -__device__ inline void barrier_release(int* lock, bool reset = false) {

272 - __syncthreads();

273 - if (threadIdx.x == 0) {

274 - if (reset) {

275 - lock[0] = 0;

276 - return;

277 - }

278 - int val = 1;

279 - // Make sure that all writes since acquiring this barrier are visible

280 - // globally, while releasing the barrier.

281 - asm volatile("fence.acq_rel.gpu;\n");

282 - asm volatile("red.relaxed.gpu.global.add.s32 [%0], %1;\n"

283 - :

284 - : "l"(lock), "r"(val));

285 - }

286 -}

287 -

288 -template <const vllm::ScalarTypeId w_type_id, // weight ScalarType id

289 - const int threads, // number of threads in a threadblock

290 - const int thread_m_blocks, // number of 16x16 blocks in the m

291 - // dimension (batchsize) of the

292 - // threadblock

293 - const int thread_n_blocks, // same for n dimension (output)

294 - const int thread_k_blocks, // same for k dimension (reduction)

295 - const int stages, // number of stages for the async global->shared

296 - // fetch pipeline

297 - const bool has_act_order, // whether act_order is enabled

298 - const bool has_zp, // whether zero-points are enabled

299 - const int group_blocks = -1 // number of consecutive 16x16 blocks

300 - // with a separate quantization scale

301 - >

302 -__device__ void MarlinMoESingle(

303 - const int4* __restrict__ A, // fp16 input matrix of shape mxk

304 - const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn

305 - int4* __restrict__ C, // fp16 output buffer of shape mxn

306 - const int* __restrict__ sorted_ids, // int32 sorted ids of experts

307 - const float* __restrict__ topk_weights, // float topk weights

308 - const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape

309 - // (k/groupsize)xn

310 - const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape

311 - // (k/groupsize)x(n/pack_factor)

312 - const int* __restrict__ g_idx, // int32 group indices of shape k

313 - const int* __restrict__ expert_offsets,

314 - int num_groups, // number of scale groups per output channel

315 - int expert_idx, // idx of current expert

316 - int num_experts, // number of experts

317 - int topk, // topk parameter of moe

318 - int prob_m, // batch dimension m

319 - int prob_n, // output dimension n

320 - int prob_k, // reduction dimension k

321 - int tot_m, // total number of rows in A and C

322 - int* locks, // extra global storage for barrier synchronization

323 - bool replicate_input, // do we use the same input for each expert?

324 - bool apply_weights, // apply weights to output

325 - int current_m_block // current m block to start kernel computation from

326 -) {

327 - static constexpr auto w_type = vllm::ScalarType::from_id(w_type_id);

328 - constexpr int pack_factor = 32 / w_type.size_bits();

329 -

330 - // For larger GEMMs we run multiple batchsize 64 versions in parallel for a

331 - // better partitioning with less reductions

332 - int parallel = 1;

333 - if (prob_m > 16 * thread_m_blocks) {

334 - parallel = prob_m / (16 * thread_m_blocks);

335 - prob_m = 16 * thread_m_blocks;

336 - }

337 -

338 - int k_tiles = prob_k / 16 / thread_k_blocks;

339 - int n_tiles = prob_n / 16 / thread_n_blocks;

340 - int iters = ceildiv(k_tiles * n_tiles * parallel, gridDim.x);

341 -

342 - if constexpr (!has_act_order && group_blocks != -1) {

343 - if (group_blocks >= thread_k_blocks) {

344 - // Ensure that the number of tiles in each stripe is a multiple of the

345 - // groupsize; this avoids an annoying special case where a stripe starts

346 - // in the middle of group.

347 - iters = (group_blocks / thread_k_blocks) *

348 - ceildiv(iters, (group_blocks / thread_k_blocks));

349 - }

350 - }

351 -

352 - int slice_row = (iters * blockIdx.x) % k_tiles;

353 - int slice_col_par = (iters * blockIdx.x) / k_tiles;

354 - int slice_col = slice_col_par;

355 - int slice_iters; // number of threadblock tiles in the current slice

356 - int slice_count =

357 - 0; // total number of active threadblocks in the current slice

358 - int slice_idx; // index of threadblock in current slice; numbered bottom to

359 - // top

360 -

361 - // We can easily implement parallel problem execution by just remapping

362 - // indices and advancing global pointers

363 - if (slice_col_par >= n_tiles) {

364 - locks += (slice_col_par / n_tiles) * n_tiles;

365 - slice_col = slice_col_par % n_tiles;

366 - sorted_ids += (slice_col_par / n_tiles) * 16 * thread_m_blocks;

367 - }

368 -

369 - // Compute all information about the current slice which is required for

370 - // synchronization.

371 - auto init_slice = [&]() {

372 - slice_iters =

373 - iters * (blockIdx.x + 1) - (k_tiles * slice_col_par + slice_row);

374 - if (slice_iters < 0 || slice_col_par >= n_tiles * parallel) slice_iters = 0;

375 - if (slice_iters == 0) return;

376 - if (slice_row + slice_iters > k_tiles) slice_iters = k_tiles - slice_row;

377 - slice_count = 1;

378 - slice_idx = 0;

379 - int col_first = iters * ceildiv(k_tiles * slice_col_par, iters);

380 - if (col_first <= k_tiles * (slice_col_par + 1)) {

381 - int col_off = col_first - k_tiles * slice_col_par;

382 - slice_count = ceildiv(k_tiles - col_off, iters);

383 - if (col_off > 0) slice_count++;

384 - int delta_first = iters * blockIdx.x - col_first;

385 - if (delta_first < 0 || (col_off == 0 && delta_first == 0))

386 - slice_idx = slice_count - 1;

387 - else {

388 - slice_idx = slice_count - 1 - delta_first / iters;

389 - if (col_off > 0) slice_idx--;

390 - }

391 - }

392 - if (slice_col == n_tiles) {

393 - sorted_ids += 16 * thread_m_blocks;

394 - locks += n_tiles;

395 - slice_col = 0;

396 - }

397 - };

398 - init_slice();

399 -

400 - // A sizes/strides

401 -

402 - // stride of the A matrix in global memory

403 - int a_gl_stride = prob_k / 8;

404 - // stride of an A matrix tile in shared memory

405 - constexpr int a_sh_stride = 16 * thread_k_blocks / 8;

406 - // delta between subsequent A tiles in global memory

407 - constexpr int a_gl_rd_delta_o = 16 * thread_k_blocks / 8;

408 - // between subsequent accesses within a tile

409 - int a_gl_rd_delta_i = a_gl_stride * (threads / a_gl_rd_delta_o);

410 - // between shared memory writes

411 - constexpr int a_sh_wr_delta = a_sh_stride * (threads / a_gl_rd_delta_o);

412 - // between shared memory tile reads

413 - constexpr int a_sh_rd_delta_o = 2 * ((threads / 32) / (thread_n_blocks / 4));

414 - // within a shared memory tile

415 - constexpr int a_sh_rd_delta_i = a_sh_stride * 16;

416 - // overall size of a tile

417 - constexpr int a_sh_stage = a_sh_stride * (16 * thread_m_blocks);

418 - // number of shared write iterations for a tile

419 - constexpr int a_sh_wr_iters = ceildiv(a_sh_stage, a_sh_wr_delta);

420 -

421 - // B sizes/strides

422 - int b_gl_stride = 16 * prob_n / (pack_factor * 4);

423 - constexpr int b_sh_stride = ((thread_n_blocks * 16) * 16 / pack_factor) / 4;

424 - constexpr int b_thread_vecs = w_type.size_bits() == 4 ? 1 : 2;

425 - constexpr int b_sh_stride_threads = b_sh_stride / b_thread_vecs;

426 -

427 - int b_gl_rd_delta_o = b_gl_stride * thread_k_blocks;

428 - int b_gl_rd_delta_i = b_gl_stride * (threads / b_sh_stride_threads);

429 - constexpr int b_sh_wr_delta = threads * b_thread_vecs;

430 - constexpr int b_sh_rd_delta = threads * b_thread_vecs;

431 - constexpr int b_sh_stage = b_sh_stride * thread_k_blocks;

432 - constexpr int b_sh_wr_iters = b_sh_stage / b_sh_wr_delta;

433 -

434 - // Scale sizes/strides without act_order

435 - int s_gl_stride = prob_n / 8;

436 - constexpr int s_sh_stride = 16 * thread_n_blocks / 8;

437 - constexpr int s_tb_groups =

438 - !has_act_order && group_blocks != -1 && group_blocks < thread_k_blocks

439 - ? thread_k_blocks / group_blocks

440 - : 1;

441 - constexpr int s_sh_stage = s_tb_groups * s_sh_stride;

442 - int s_gl_rd_delta = s_gl_stride;

443 - // Scale size/strides with act_order

444 - constexpr int tb_k = 16 * thread_k_blocks;

445 - constexpr int g_idx_stage = has_act_order ? (tb_k * sizeof(int)) / 16 : 0;

446 - // constexpr int act_s_row_stride = 1;

447 - // int act_s_col_stride = act_s_row_stride * num_groups;

448 - int act_s_col_stride = 1;

449 - int act_s_col_warp_stride = act_s_col_stride * 8;

450 - int tb_n_warps = thread_n_blocks / 4;

451 - int act_s_col_tb_stride = act_s_col_warp_stride * tb_n_warps;

452 -

453 - // Zero-points sizes/strides

454 - int zp_gl_stride = (prob_n / pack_factor) / 4;

455 - constexpr int zp_sh_stride = ((16 * thread_n_blocks) / pack_factor) / 4;

456 - constexpr int zp_tb_groups = s_tb_groups;

457 - constexpr int zp_sh_stage = has_zp ? zp_tb_groups * zp_sh_stride : 0;

458 - int zp_gl_rd_delta = zp_gl_stride;

459 -

460 - // Global A read index of current thread.

461 - int a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +

462 - (threadIdx.x % a_gl_rd_delta_o);

463 - a_gl_rd += a_gl_rd_delta_o * slice_row;

464 - // Shared write index of current thread.

465 - int a_sh_wr = a_sh_stride * (threadIdx.x / a_gl_rd_delta_o) +

466 - (threadIdx.x % a_gl_rd_delta_o);

467 - // Shared read index.

468 - int a_sh_rd =

469 - a_sh_stride * ((threadIdx.x % 32) % 16) + (threadIdx.x % 32) / 16;

470 - a_sh_rd += 2 * ((threadIdx.x / 32) / (thread_n_blocks / 4));

471 -

472 - int b_gl_rd = b_gl_stride * (threadIdx.x / b_sh_stride_threads) +

473 - (threadIdx.x % b_sh_stride_threads) * b_thread_vecs;

474 - b_gl_rd += b_sh_stride * slice_col;

475 - b_gl_rd += b_gl_rd_delta_o * slice_row;

476 - int b_sh_wr = threadIdx.x * b_thread_vecs;

477 - int b_sh_rd = threadIdx.x * b_thread_vecs;

478 -

479 - // For act_order

480 - constexpr int k_iter_size = tb_k / b_sh_wr_iters;

481 - int slice_k_start = tb_k * slice_row;

482 - int slice_k_finish = slice_k_start + tb_k * slice_iters;

483 - int slice_k_start_shared_fetch = slice_k_start;

484 - int slice_n_offset = act_s_col_tb_stride * slice_col;

485 -

486 - // No act_order

487 - int s_gl_rd;

488 - if constexpr (!has_act_order) {

489 - if constexpr (group_blocks == -1) {

490 - s_gl_rd = s_sh_stride * slice_col + threadIdx.x;

491 - } else {

492 - s_gl_rd = s_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +

493 - s_sh_stride * slice_col + threadIdx.x;

494 - }

495 - }

496 - int s_sh_wr = threadIdx.x;

497 - bool s_sh_wr_pred = threadIdx.x < s_sh_stride;

498 -

499 - // Zero-points

500 - int zp_gl_rd;

501 - if constexpr (has_zp) {

502 - if constexpr (group_blocks == -1) {

503 - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;

504 - } else {

505 - zp_gl_rd = zp_gl_stride * ((thread_k_blocks * slice_row) / group_blocks) +

506 - zp_sh_stride * slice_col + threadIdx.x;

507 - }

508 - }

509 - int zp_sh_wr = threadIdx.x;

510 - bool zp_sh_wr_pred = threadIdx.x < zp_sh_stride;

511 -

512 - // We use a different scale layout for grouped and column-wise quantization as

513 - // we scale a `half2` tile in column-major layout in the former and in

514 - // row-major in the latter case.

515 - int s_sh_rd;

516 - if constexpr (group_blocks != -1)

517 - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +

518 - (threadIdx.x % 32) / 4;

519 - else

520 - s_sh_rd = 8 * ((threadIdx.x / 32) % (thread_n_blocks / 4)) +

521 - (threadIdx.x % 32) % 4;

522 -

523 - // Zero-points have the same read layout as the scales

524 - // (without column-wise case)

525 - constexpr int num_col_threads = 8;

526 - constexpr int num_row_threads = 4;

527 - constexpr int num_ints_per_thread = 8 / pack_factor;

528 - int zp_sh_rd;

529 - if constexpr (has_zp) {

530 - zp_sh_rd = num_ints_per_thread * num_col_threads *

531 - ((threadIdx.x / 32) % (thread_n_blocks / 4)) +

532 - num_ints_per_thread * ((threadIdx.x % 32) / num_row_threads);

533 - }

534 -

535 - int sh_first_group_id = -1;

536 - int sh_num_groups = -1;

537 - constexpr int sh_max_num_groups = 32;

538 -

539 - extern __shared__ int4 sh[];

540 - // Shared memory storage for global fetch pipelines.

541 - int4* sh_a = sh;

542 - int4* sh_b = sh_a + (stages * a_sh_stage);

543 - int4* sh_g_idx = sh_b + (stages * b_sh_stage);

544 - int4* sh_zp = sh_g_idx + (stages * g_idx_stage);

545 - int4* sh_s = sh_zp + (stages * zp_sh_stage);

546 -

547 - // Precompute which thread should not read memory in which iterations; this is

548 - // needed if there are more threads than required for a certain tilesize or

549 - // when the batchsize is not a multiple of 16.

550 - bool a_sh_wr_pred[a_sh_wr_iters];

551 - #pragma unroll

552 - for (int i = 0; i < a_sh_wr_iters; i++) {

553 - int a_idx = a_sh_wr_delta * i + a_sh_wr;

554 - int row = a_idx / a_gl_rd_delta_o;

555 - if (row >= prob_m) {

556 - a_sh_wr_pred[i] = false;

557 - } else {

558 - a_sh_wr_pred[i] = a_sh_wr_delta * i + a_sh_wr < a_sh_stride * prob_m;

559 - }

560 - }

561 -

562 - // To ensure that writing and reading A tiles to/from shared memory, the

563 - // latter in fragment format, is fully bank conflict free, we need to use a

564 - // rather fancy XOR-based layout. The key here is that neither reads nor

565 - // writes of the 16-byte `int4` blocks of 8 consecutive threads involve the

566 - // same shared memory banks. Further, it seems (based on NSight-Compute) that

567 - // each warp must also write a consecutive memory segment?

568 - auto transform_a = [&](int i) {

569 - int row = i / a_gl_rd_delta_o;

570 - return a_gl_rd_delta_o * row + (i % a_gl_rd_delta_o) ^ row;

571 - };

572 - // Since the computation of this remapping is non-trivial and, due to our main

573 - // loop unrolls, all shared memory accesses are static, we simply precompute

574 - // both transformed reads and writes.

575 - int a_sh_wr_trans[a_sh_wr_iters];

576 - #pragma unroll

577 - for (int i = 0; i < a_sh_wr_iters; i++)

578 - a_sh_wr_trans[i] = transform_a(a_sh_wr_delta * i + a_sh_wr);

579 - int a_sh_rd_trans[b_sh_wr_iters][thread_m_blocks];

580 - #pragma unroll

581 - for (int i = 0; i < b_sh_wr_iters; i++) {

582 - #pragma unroll

583 - for (int j = 0; j < thread_m_blocks; j++)

584 - a_sh_rd_trans[i][j] =

585 - transform_a(a_sh_rd_delta_o * i + a_sh_rd_delta_i * j + a_sh_rd);

586 - }

587 -

588 - // Since B-accesses have non-constant stride they have to be computed at

589 - // runtime; we break dependencies between subsequent accesses with a tile by

590 - // maintining multiple pointers (we have enough registers), a tiny

591 - // optimization.

592 - const int4* B_ptr[b_sh_wr_iters];

593 - #pragma unroll

594 - for (int i = 0; i < b_sh_wr_iters; i++)

595 - B_ptr[i] = B + b_gl_rd_delta_i * i + b_gl_rd;

596 -

597 - // Register storage for double buffer of shared memory reads.

598 - FragA frag_a[2][thread_m_blocks];

599 - I4 frag_b_quant[2][b_thread_vecs];

600 - FragC frag_c[thread_m_blocks][4][2];

601 - FragS frag_s[2][4]; // No act-order

602 - FragS act_frag_s[2][4][4]; // For act-order

603 - int frag_qzp[2][num_ints_per_thread]; // Zero-points

604 - FragZP frag_zp; // Zero-points in fp16

605 -

606 - // Zero accumulators.

607 - auto zero_accums = [&]() {

608 - #pragma unroll

609 - for (int i = 0; i < thread_m_blocks * 4 * 2 * 4; i++)

610 - reinterpret_cast<float*>(frag_c)[i] = 0;

611 - };

612 -

613 - auto fetch_scales_to_shared = [&](bool is_async, int first_group_id,

614 - int last_group_id) {

615 - sh_first_group_id = first_group_id;

616 - sh_num_groups = last_group_id - first_group_id + 1;

617 -

618 - if (sh_num_groups < sh_max_num_groups) {

619 - sh_num_groups = sh_max_num_groups;

620 - }

621 -

622 - if (sh_first_group_id + sh_num_groups > num_groups) {

623 - sh_num_groups = num_groups - sh_first_group_id;

624 - }

625 -

626 - int row_offset = first_group_id * s_gl_stride;

627 -

628 - if (is_async) {

629 - for (int i = 0; i < sh_num_groups; i++) {

630 - if (threadIdx.x < s_sh_stride) {

631 - cp_async4_pred(&sh_s[(i * s_sh_stride) + threadIdx.x],

632 - &scales_ptr[row_offset + (i * s_gl_stride) +

633 - slice_n_offset + threadIdx.x]);

634 - }

635 - }

636 - } else {

637 - for (int i = 0; i < sh_num_groups; i++) {

638 - if (threadIdx.x < s_sh_stride) {

639 - sh_s[(i * s_sh_stride) + threadIdx.x] =

640 - scales_ptr[row_offset + (i * s_gl_stride) + slice_n_offset +

641 - threadIdx.x];

642 - }

643 - }

644 - }

645 - };

646 - // Asynchronously fetch the next A, B and s tile from global to the next

647 - // shared memory pipeline location.

648 - auto fetch_to_shared = [&](int pipe, int a_off, bool pred = true) {

649 - if (pred) {

650 - int4* sh_a_stage = sh_a + a_sh_stage * pipe;

651 - #pragma unroll

652 - for (int i = 0; i < a_sh_wr_iters; i++) {

653 - int a_idx = a_gl_rd_delta_i * i + a_gl_rd + a_gl_rd_delta_o * a_off;

654 - int row = a_idx / a_gl_stride;

655 - int sorted_row =

656 - replicate_input ? sorted_ids[row] / topk : sorted_ids[row];

657 - int new_idx = sorted_row * a_gl_stride + a_idx % a_gl_stride;

658 - if (sorted_row < tot_m * (replicate_input ? 1 : topk) &&

659 - new_idx < a_gl_stride * tot_m * (replicate_input ? 1 : topk)) {

660 - cp_async4_pred(&sh_a_stage[a_sh_wr_trans[i]], &A[new_idx],

661 - a_sh_wr_pred[i]);

662 - }

663 - }

664 - int4* sh_b_stage = sh_b + b_sh_stage * pipe;

665 - #pragma unroll

666 - for (int i = 0; i < b_sh_wr_iters; i++) {

667 - #pragma unroll

668 - for (int j = 0; j < b_thread_vecs; j++) {

669 - cp_async4(&sh_b_stage[b_sh_wr_delta * i + b_sh_wr + j], B_ptr[i] + j);

670 - }

671 - B_ptr[i] += b_gl_rd_delta_o;

672 - }

673 -

674 - if constexpr (has_act_order) {

675 - // Fetch g_idx thread-block portion

676 - int full_pipe = a_off;

677 - int cur_k = slice_k_start_shared_fetch + tb_k * full_pipe;

678 - if (cur_k < prob_k && cur_k < slice_k_finish) {

679 - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;

680 -

681 - int4 const* cur_g_idx_stage_ptr =

682 - reinterpret_cast<int4 const*>(&g_idx[cur_k]);

683 -

684 - if (threadIdx.x < g_idx_stage) {

685 - cp_async4_pred(&sh_g_idx_stage[threadIdx.x],

686 - &cur_g_idx_stage_ptr[threadIdx.x]);

687 - }

688 - }

689 - } else {

690 - if constexpr (group_blocks != -1) {

691 - int4* sh_s_stage = sh_s + s_sh_stage * pipe;

692 -

693 - if constexpr (group_blocks >= thread_k_blocks) {

694 - // Only fetch scales if this tile starts a new group

695 - if (pipe % (group_blocks / thread_k_blocks) == 0) {

696 - if (s_sh_wr_pred) {

697 - cp_async4(&sh_s_stage[s_sh_wr], &scales_ptr[s_gl_rd]);

698 - }

699 - s_gl_rd += s_gl_rd_delta;

700 - }

701 - } else {

702 - for (int i = 0; i < s_tb_groups; i++) {

703 - if (s_sh_wr_pred) {

704 - cp_async4(&sh_s_stage[i * s_sh_stride + s_sh_wr],

705 - &scales_ptr[s_gl_rd]);

706 - }

707 - s_gl_rd += s_gl_rd_delta;

708 - }

709 - }

710 - }

711 -

712 - if constexpr (has_zp && group_blocks != -1) {

713 - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;

714 -

715 - if constexpr (group_blocks >= thread_k_blocks) {

716 - // Only fetch zero-points if this tile starts a new group

717 - if (pipe % (group_blocks / thread_k_blocks) == 0) {

718 - if (zp_sh_wr_pred) {

719 - cp_async4(&sh_zp_stage[zp_sh_wr], &zp_ptr[zp_gl_rd]);

720 - }

721 - zp_gl_rd += zp_gl_rd_delta;

722 - }

723 - } else {

724 - for (int i = 0; i < zp_tb_groups; i++) {

725 - if (zp_sh_wr_pred) {

726 - cp_async4(&sh_zp_stage[i * zp_sh_stride + zp_sh_wr],

727 - &zp_ptr[zp_gl_rd]);

728 - }

729 - zp_gl_rd += zp_gl_rd_delta;

730 - }

731 - }

732 - }

733 - }

734 - }

735 - // Insert a fence even when we are winding down the pipeline to ensure that

736 - // waiting is also correct at this point.

737 - cp_async_fence();

738 - };

739 -

740 - auto fetch_zp_to_shared = [&]() {

741 - if (zp_sh_wr_pred) {

742 - cp_async4(&sh_zp[zp_sh_wr], &zp_ptr[zp_gl_rd]);

743 - }

744 - };

745 -

746 - // Wait until the next thread tile has been loaded to shared memory.

747 - auto wait_for_stage = [&]() {

748 - // We only have `stages - 2` active fetches since we are double buffering

749 - // and can only issue the next fetch when it is guaranteed that the previous

750 - // shared memory load is fully complete (as it may otherwise be

751 - // overwritten).

752 - cp_async_wait<stages - 2>();

753 - __syncthreads();

754 - };

755 -

756 - // Load the next sub-tile from the current location in the shared memory pipe

757 - // into the current register buffer.

758 - auto fetch_to_registers = [&](int k, int pipe) {

759 - int4* sh_a_stage = sh_a + a_sh_stage * pipe;

760 - #pragma unroll

761 - for (int i = 0; i < thread_m_blocks; i++)

762 - ldsm4(frag_a[k % 2][i], &sh_a_stage[a_sh_rd_trans[k % b_sh_wr_iters][i]]);

763 - int4* sh_b_stage = sh_b + b_sh_stage * pipe;

764 -

765 - #pragma unroll

766 - for (int i = 0; i < b_thread_vecs; i++) {

767 - frag_b_quant[k % 2][i] = *reinterpret_cast<I4*>(

768 - &sh_b_stage[b_sh_rd_delta * (k % b_sh_wr_iters) + b_sh_rd + i]);

769 - }

770 - };

771 -

772 - bool is_same_group[stages];

773 - int same_group_id[stages];

774 -

775 - auto init_same_group = [&](int pipe) {

776 - if constexpr (!has_act_order) {

777 - is_same_group[pipe] = false;

778 - same_group_id[pipe] = 0;

779 - return;

780 - }

781 -

782 - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;

783 - int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);

784 -

785 - int group_id_1 = sh_g_idx_int_ptr[0];

786 - int group_id_2 = sh_g_idx_int_ptr[tb_k - 1];

787 -

788 - is_same_group[pipe] = group_id_1 == group_id_2;

789 - same_group_id[pipe] = group_id_1;

790 - };

791 -

792 - auto fetch_scales_to_registers = [&](int k, int full_pipe) {

793 - int pipe = full_pipe % stages;

794 -

795 - if constexpr (!has_act_order) {

796 - // No act-order case

797 - if constexpr (group_blocks != -1) {

798 - if constexpr (group_blocks >= thread_k_blocks) {

799 - int4* sh_s_stage =

800 - sh_s + s_sh_stage * ((group_blocks / thread_k_blocks) *

801 - (pipe / (group_blocks / thread_k_blocks)));

802 - reinterpret_cast<int4*>(&frag_s[k % 2])[0] = sh_s_stage[s_sh_rd];

803 - } else {

804 - int warp_id = threadIdx.x / 32;

805 - int n_warps = thread_n_blocks / 4;

806 -

807 - int warp_row = warp_id / n_warps;

808 -

809 - int cur_k = warp_row * 16;

810 - cur_k += k_iter_size * (k % b_sh_wr_iters);

811 -

812 - int k_blocks = cur_k / 16;

813 - int cur_group_id = k_blocks / group_blocks;

814 -

815 - int4* sh_s_stage = sh_s + s_sh_stage * pipe;

816 -

817 - reinterpret_cast<int4*>(&frag_s[k % 2])[0] =

818 - sh_s_stage[s_sh_rd + cur_group_id * s_sh_stride];

819 - }

820 - }

821 -

822 - return;

823 - }

824 -

825 - // Act-order case

826 -

827 - // Determine K of the "current" thread-block

828 - int cur_k = slice_k_start + tb_k * full_pipe;

829 - if (cur_k >= prob_k || cur_k >= slice_k_finish) {

830 - return;

831 - }

832 -

833 - // Reset (to current thread-block) since we read g_idx portion from the

834 - // shared memory

835 - cur_k = 0;

836 -

837 - // Progress to current iteration

838 - cur_k += k_iter_size * (k % b_sh_wr_iters);

839 -

840 - // Determine "position" inside the thread-block (based on warp and

841 - // thread-id)

842 - int warp_id = threadIdx.x / 32;

843 - int n_warps =

844 - thread_n_blocks / 4; // Each warp processes 4 16-size tiles over N

845 -

846 - int warp_row = warp_id / n_warps;

847 - int warp_col = warp_id % n_warps;

848 -

849 - cur_k += warp_row * 16;

850 -

851 - int th_id = threadIdx.x % 32;

852 - cur_k += (th_id % 4) * 2; // Due to tensor-core layout for fp16 B matrix

853 -

854 - int s_col_shift =

855 - /*slice_n_offset +*/ (act_s_col_warp_stride * warp_col) +

856 - (th_id / 4) * act_s_col_stride;

857 -

858 - if (is_same_group[pipe]) {

859 - if (k % 2 == 0) {

860 - *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =

861 - sh_s[(same_group_id[pipe] - sh_first_group_id) * s_sh_stride +

862 - s_col_shift];

863 - } else {

864 - *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0]))) =

865 - *(reinterpret_cast<int4*>(&(act_frag_s[(k - 1) % 2][0][0])));

866 - }

867 -

868 - for (int i = 1; i < 4; i++) {

869 - *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =

870 - *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][0][0])));

871 - }

872 - return;

873 - }

874 -

875 - int4* sh_g_idx_stage = sh_g_idx + g_idx_stage * pipe;

876 - int* sh_g_idx_int_ptr = reinterpret_cast<int*>(sh_g_idx_stage);

877 -

878 - constexpr int k_frag_offsets[4] = {0, 1, 8,

879 - 9}; // Tensor core offsets per thread

880 -

881 - #pragma unroll

882 - for (int i = 0; i < 4; i++) {

883 - int actual_k = cur_k + k_frag_offsets[i];

884 -

885 - int group_id = sh_g_idx_int_ptr[actual_k];

886 - int rel_group_id = group_id - sh_first_group_id;

887 -

888 - *(reinterpret_cast<int4*>(&(act_frag_s[k % 2][i][0]))) =

889 - sh_s[rel_group_id * s_sh_stride + s_col_shift];

890 - }

891 - };

892 -

893 - auto fetch_zp_to_registers = [&](int k, int full_pipe) {

894 - // This code does not handle group_blocks == 0,

895 - // which signifies act_order.

896 - // has_zp implies AWQ, which doesn't have act_order,

897 - static_assert(!has_zp || group_blocks != 0);

898 -

899 - if constexpr (has_zp) {

900 - int pipe = full_pipe % stages;

901 -

902 - if constexpr (group_blocks == -1) {

903 - for (int i = 0; i < num_ints_per_thread; i++) {

904 - frag_qzp[k % 2][i] = (reinterpret_cast<int*>(sh_zp))[zp_sh_rd + i];

905 - }

906 -

907 - } else if constexpr (group_blocks >= thread_k_blocks) {

908 - int4* sh_zp_stage =

909 - sh_zp + zp_sh_stage * ((group_blocks / thread_k_blocks) *

910 - (pipe / (group_blocks / thread_k_blocks)));

911 - for (int i = 0; i < num_ints_per_thread; i++) {

912 - frag_qzp[k % 2][i] =

913 - (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];

914 - }

915 - } else {

916 - int warp_id = threadIdx.x / 32;

917 - int n_warps = thread_n_blocks / 4;

918 -

919 - int warp_row = warp_id / n_warps;

920 -

921 - int cur_k = warp_row * 16;

922 - cur_k += k_iter_size * (k % b_sh_wr_iters);

923 -

924 - int k_blocks = cur_k / 16;

925 - int cur_group_id = 0;

926 -

927 - // Suppress bogus and persistent divide-by-zero warning

928 - #pragma nv_diagnostic push

929 - #pragma nv_diag_suppress divide_by_zero

930 - cur_group_id = k_blocks / group_blocks;

931 - #pragma nv_diagnostic pop

932 -

933 - int4* sh_zp_stage = sh_zp + zp_sh_stage * pipe;

934 -

935 - sh_zp_stage += cur_group_id * zp_sh_stride;

936 -

937 - for (int i = 0; i < num_ints_per_thread; i++) {

938 - frag_qzp[k % 2][i] =

939 - (reinterpret_cast<int*>(sh_zp_stage))[zp_sh_rd + i];

940 - }

941 - }

942 - }

943 - };

944 -

945 - // Execute the actual tensor core matmul of a sub-tile.

946 - auto matmul = [&](int k) {

947 - if constexpr (has_zp) {

948 - FragB frag_zp_0;

949 - FragB frag_zp_1;

950 - int zp_quant_0, zp_quant_1;

951 -

952 - if constexpr (w_type.size_bits() == 4) {

953 - zp_quant_0 = frag_qzp[k % 2][0];

954 - zp_quant_1 = zp_quant_0 >> 8;

955 - } else {

956 - static_assert(w_type.size_bits() == 8);

957 - zp_quant_0 = frag_qzp[k % 2][0];

958 - zp_quant_1 = frag_qzp[k % 2][1];

959 - }

960 -

961 - frag_zp_0 = dequant<w_type_id>(zp_quant_0);

962 - frag_zp_1 = dequant<w_type_id>(zp_quant_1);

963 -

964 - frag_zp[0] = frag_zp_0[0];

965 - frag_zp[1] = frag_zp_0[1];

966 - frag_zp[2] = frag_zp_1[0];

967 - frag_zp[3] = frag_zp_1[1];

968 - }

969 -

970 - // We have the m dimension as the inner loop in order to encourage overlapping

971 - // dequantization and matmul operations.

972 - #pragma unroll

973 - for (int j = 0; j < 4; j++) {

974 - int b_quant_0, b_quant_1;

975 - if constexpr (w_type.size_bits() == 4) {

976 - b_quant_0 = frag_b_quant[k % 2][0][j];

977 - b_quant_1 = b_quant_0 >> 8;

978 - } else {

979 - static_assert(w_type.size_bits() == 8);

980 - int* frag_b_quant_ptr = reinterpret_cast<int*>(frag_b_quant[k % 2]);

981 - b_quant_0 = frag_b_quant_ptr[j * 2 + 0];

982 - b_quant_1 = frag_b_quant_ptr[j * 2 + 1];

983 - }

984 -

985 - FragB frag_b0 = dequant<w_type_id>(b_quant_0);

986 - FragB frag_b1 = dequant<w_type_id>(b_quant_1);

987 - // Apply zero-point to frag_b0

988 - if constexpr (has_zp) {

989 - sub_zp(frag_b0, frag_zp[j], 0);

990 - }

991 -

992 - // Apply scale to frag_b0

993 - if constexpr (has_act_order) {

994 - scale4(frag_b0, act_frag_s[k % 2][0][j], act_frag_s[k % 2][1][j],

995 - act_frag_s[k % 2][2][j], act_frag_s[k % 2][3][j], 0);

996 - } else {

997 - if constexpr (group_blocks != -1) {

998 - scale(frag_b0, frag_s[k % 2][j], 0);

999 - }

1000 - }

1001 -

1002 - // Apply zero-point to frag_b1

1003 - if constexpr (has_zp) {

1004 - sub_zp(frag_b1, frag_zp[j], 1);

1005 - }

1006 -

1007 - // Apply scale to frag_b1

1008 - if constexpr (has_act_order) {

1009 - scale4(frag_b1, act_frag_s[k % 2][0][j], act_frag_s[k % 2][1][j],

1010 - act_frag_s[k % 2][2][j], act_frag_s[k % 2][3][j], 1);

1011 -

1012 - } else {

1013 - if constexpr (group_blocks != -1) {

1014 - scale(frag_b1, frag_s[k % 2][j], 1);

1015 - }

1016 - }

1017 -

1018 - #pragma unroll

1019 - for (int i = 0; i < thread_m_blocks; i++) {

1020 - mma(frag_a[k % 2][i], frag_b0, frag_c[i][j][0]);

1021 - mma(frag_a[k % 2][i], frag_b1, frag_c[i][j][1]);

1022 - }

1023 - }

1024 - };

1025 -

1026 - // Since we slice across the k dimension of a tile in order to increase the

1027 - // number of warps while keeping the n dimension of a tile reasonable, we have

1028 - // multiple warps that accumulate their partial sums of the same output

1029 - // location; which we have to reduce over in the end. We do in shared memory.

1030 - auto thread_block_reduce = [&]() {

1031 - constexpr int red_off = threads / b_sh_stride_threads / 2;

1032 - if (red_off >= 1) {

1033 - int red_idx = threadIdx.x / b_sh_stride_threads;

1034 - constexpr int red_sh_stride = b_sh_stride_threads * 4 * 2;

1035 - constexpr int red_sh_delta = b_sh_stride_threads;

1036 - int red_sh_rd = red_sh_stride * (threadIdx.x / b_sh_stride_threads) +

1037 - (threadIdx.x % b_sh_stride_threads);

1038 -

1039 - // Parallel logarithmic shared memory reduction. We make sure to avoid any

1040 - // unnecessary read or write iterations, e.g., for two warps we write only

1041 - // once by warp 1 and read only once by warp 0.

1042 -

1043 - #pragma unroll

1044 - for (int m_block = 0; m_block < thread_m_blocks; m_block++) {

1045 - #pragma unroll

1046 - for (int i = red_off; i > 0; i /= 2) {

1047 - if (i <= red_idx && red_idx < 2 * i) {

1048 - #pragma unroll

1049 - for (int j = 0; j < 4 * 2; j++) {

1050 - int red_sh_wr =

1051 - red_sh_delta * j + (red_sh_rd - red_sh_stride * i);

1052 - if (i < red_off) {

1053 - float* c_rd =

1054 - reinterpret_cast<float*>(&sh[red_sh_delta * j + red_sh_rd]);

1055 - float* c_wr = reinterpret_cast<float*>(&sh[red_sh_wr]);

1056 - #pragma unroll

1057 - for (int k = 0; k < 4; k++)

1058 - reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + j][k] +=

1059 - c_rd[k] + c_wr[k];

1060 - }

1061 - sh[red_sh_wr] =

1062 - reinterpret_cast<int4*>(&frag_c)[4 * 2 * m_block + j];

1063 - }

1064 - }

1065 - __syncthreads();

1066 - }

1067 - if (red_idx == 0) {

1068 - #pragma unroll

1069 - for (int i = 0; i < 4 * 2; i++) {

1070 - float* c_rd =

1071 - reinterpret_cast<float*>(&sh[red_sh_delta * i + red_sh_rd]);

1072 - #pragma unroll

1073 - for (int j = 0; j < 4; j++)

1074 - reinterpret_cast<FragC*>(frag_c)[4 * 2 * m_block + i][j] +=

1075 - c_rd[j];

1076 - }

1077 - }

1078 - __syncthreads();

1079 - }

1080 - }

1081 - };

1082 -

1083 - // Since multiple threadblocks may process parts of the same column slice, we

1084 - // finally have to globally reduce over the results. As the striped

1085 - // partitioning minimizes the number of such reductions and our outputs are

1086 - // usually rather small, we perform this reduction serially in L2 cache.

1087 - auto global_reduce = [&](bool first = false, bool last = false) {

1088 - // We are very careful here to reduce directly in the output buffer to

1089 - // maximize L2 cache utilization in this step. To do this, we write out

1090 - // results in FP16 (but still reduce with FP32 compute).

1091 - constexpr int active_threads = 32 * thread_n_blocks / 4;

1092 - if (threadIdx.x < active_threads) {

1093 - int c_gl_stride = prob_n / 8;

1094 - int c_gl_wr_delta_o = 8 * c_gl_stride;

1095 - int c_gl_wr_delta_i = 4 * (active_threads / 32);

1096 - int c_gl_wr = c_gl_stride * ((threadIdx.x % 32) / 4) +

1097 - 4 * (threadIdx.x / 32) + threadIdx.x % 4;

1098 - c_gl_wr += (2 * thread_n_blocks) * slice_col;

1099 - constexpr int c_sh_wr_delta = active_threads;

1100 - int c_sh_wr = threadIdx.x;

1101 -

1102 - int row = (threadIdx.x % 32) / 4;

1103 -

1104 - if (!first) {

1105 - // Interestingly, doing direct global accesses here really seems to mess up

1106 - // the compiler and lead to slowdowns, hence we also use async-copies even

1107 - // though these fetches are not actually asynchronous.

1108 - #pragma unroll

1109 - for (int i = 0; i < thread_m_blocks * 4; i++) {

1110 - int c_idx =

1111 - c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);

1112 - int sorted_row = sorted_ids[c_idx / c_gl_stride];

1113 - int new_idx = sorted_row * c_gl_stride + c_idx % c_gl_stride;

1114 - cp_async4_pred(&sh[c_sh_wr + c_sh_wr_delta * i], &C[new_idx],

1115 - sorted_row < tot_m * topk &&

1116 - (8 * (i / 2) + row < prob_m &&

1117 - (i < (thread_m_blocks - 1) * 4 ||

1118 - sorted_ids[8 * (i / 2) + row] < tot_m * topk)));

1119 - }

1120 - cp_async_fence();

1121 - cp_async_wait<0>();

1122 - }

1123 -

1124 - #pragma unroll

1125 - for (int i = 0; i < thread_m_blocks * 4; i++) {

1126 - if (8 * (i / 2) + row < prob_m &&

1127 - (i < (thread_m_blocks - 1) * 4 ||

1128 - sorted_ids[8 * (i / 2) + row] < tot_m * topk)) {

1129 - if (!first) {

1130 - int4 c_red = sh[c_sh_wr + i * c_sh_wr_delta];

1131 - #pragma unroll

1132 - for (int j = 0; j < 2 * 4; j++) {

1133 - reinterpret_cast<float*>(

1134 - &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)] +=

1135 - __half2float(reinterpret_cast<__half*>(&c_red)[j]);

1136 - }

1137 - }

1138 - if (!last) {

1139 - int4 c;

1140 - #pragma unroll

1141 - for (int j = 0; j < 2 * 4; j++) {

1142 - reinterpret_cast<__half*>(&c)[j] =

1143 - __float2half(reinterpret_cast<float*>(

1144 - &frag_c)[4 * 2 * 4 * (i / 4) + 4 * j + (i % 4)]);

1145 - }

1146 - int c_idx =

1147 - c_gl_wr + c_gl_wr_delta_o * (i / 2) + c_gl_wr_delta_i * (i % 2);

1148 - int row = sorted_ids[c_idx / c_gl_stride];

1149 - if (row < tot_m * topk) {

1150 - int new_idx = row * c_gl_stride + c_idx % c_gl_stride;

1151 - C[new_idx] = c;

1152 - }

1153 - }

1154 - }

1155 - }

1156 - }

1157 - };

1158 -

1159 - // Write out the reduce final result in the correct layout. We only actually

1160 - // reshuffle matrix fragments in this step, the reduction above is performed

1161 - // in fragment layout.

1162 - auto write_result = [&]() {

1163 - int c_gl_stride = prob_n / 8;

1164 - constexpr int c_sh_stride = 2 * thread_n_blocks + 1;

1165 - int c_gl_wr_delta = c_gl_stride * (threads / (2 * thread_n_blocks));

1166 - constexpr int c_sh_rd_delta =

1167 - c_sh_stride * (threads / (2 * thread_n_blocks));

1168 -

1169 - int c_gl_wr = c_gl_stride * (threadIdx.x / (2 * thread_n_blocks)) +

1170 - (threadIdx.x % (2 * thread_n_blocks));

1171 - c_gl_wr += (2 * thread_n_blocks) * slice_col;

1172 - int c_sh_wr =

1173 - (4 * c_sh_stride) * ((threadIdx.x % 32) / 4) + (threadIdx.x % 32) % 4;

1174 - c_sh_wr += 32 * (threadIdx.x / 32);

1175 - int c_sh_rd = c_sh_stride * (threadIdx.x / (2 * thread_n_blocks)) +

1176 - (threadIdx.x % (2 * thread_n_blocks));

1177 -

1178 - int c_gl_wr_end = c_gl_stride * prob_m;

1179 -

1180 - // We first reorder in shared memory to guarantee the most efficient final

1181 - // global write patterns

1182 - auto write = [&](int idx, float c0, float c1, FragS& s) {

1183 - half2 res = __halves2half2(__float2half(c0), __float2half(c1));

1184 -

1185 - // For per-column quantization we finally apply the scale here (only for

1186 - // 4-bit)

1187 - if constexpr (!has_act_order && group_blocks == -1 &&

1188 - w_type.size_bits() == 4) {

1189 - res = __hmul2(res, s[0]);

1190 - }

1191 -

1192 - ((half2*)sh)[idx] = res;

1193 - };

1194 - if (threadIdx.x / 32 < thread_n_blocks / 4) {

1195 - #pragma unroll

1196 - for (int i = 0; i < thread_m_blocks; i++) {

1197 - #pragma unroll

1198 - for (int j = 0; j < 4; j++) {

1199 - int wr = c_sh_wr + 8 * j;

1200 - write(wr + (4 * c_sh_stride) * 0 + 0, frag_c[i][j][0][0],

1201 - frag_c[i][j][0][1], frag_s[j / 2][2 * (j % 2) + 0]);

1202 - write(wr + (4 * c_sh_stride) * 8 + 0, frag_c[i][j][0][2],

1203 - frag_c[i][j][0][3], frag_s[j / 2][2 * (j % 2) + 0]);

1204 - write(wr + (4 * c_sh_stride) * 0 + 4, frag_c[i][j][1][0],

1205 - frag_c[i][j][1][1], frag_s[j / 2][2 * (j % 2) + 1]);

1206 - write(wr + (4 * c_sh_stride) * 8 + 4, frag_c[i][j][1][2],

1207 - frag_c[i][j][1][3], frag_s[j / 2][2 * (j % 2) + 1]);

1208 - }

1209 - c_sh_wr += 16 * (4 * c_sh_stride);

1210 - }

1211 - }

1212 - __syncthreads();

1213 -

1214 - #pragma unroll

1215 - for (int i = 0;

1216 - i < ceildiv(16 * thread_m_blocks, threads / (2 * thread_n_blocks));

1217 - i++) {

1218 - if (c_gl_wr < c_gl_wr_end) {

1219 - int row = sorted_ids[c_gl_wr / c_gl_stride];

1220 - if (row < tot_m * topk) {

1221 - int off = row * c_gl_stride + c_gl_wr % c_gl_stride;

1222 - if (!apply_weights) {

1223 - C[off] = sh[c_sh_rd];

1224 - } else {

1225 - __half* ctrg = reinterpret_cast<__half*>(&C[off]);

1226 - __half* csrc = reinterpret_cast<__half*>(&sh[c_sh_rd]);

1227 - for (int j = 0; j < 8; ++j) {

1228 - ctrg[j] = __float2half(topk_weights[row] * __half2float(csrc[j]));

1229 - }

1230 - }

1231 - c_gl_wr += c_gl_wr_delta;

1232 - c_sh_rd += c_sh_rd_delta;

1233 - }

1234 - }

1235 - }

1236 - };

1237 -

1238 - // Start global fetch and register load pipelines.

1239 - auto start_pipes = [&]() {

1240 -

1241 - #pragma unroll

1242 - for (int i = 0; i < stages - 1; i++) {

1243 - if (has_act_order && i == 0) {

1244 - int last_g_idx = slice_k_start + stages * tb_k * 2;

1245 - if (last_g_idx >= prob_k) {

1246 - last_g_idx = prob_k - 1;

1247 - }

1248 - fetch_scales_to_shared(true, g_idx[slice_k_start], g_idx[last_g_idx]);

1249 - }

1250 -

1251 - if constexpr (has_zp && group_blocks == -1) {

1252 - if (i == 0) {

1253 - fetch_zp_to_shared();

1254 - }

1255 - }

1256 - fetch_to_shared(i, i, i < slice_iters);

1257 - }

1258 -

1259 - zero_accums();

1260 - wait_for_stage();

1261 - init_same_group(0);

1262 - fetch_to_registers(0, 0);

1263 - fetch_scales_to_registers(0, 0);

1264 - fetch_zp_to_registers(0, 0);

1265 - a_gl_rd += a_gl_rd_delta_o * (stages - 1);

1266 - slice_k_start_shared_fetch += tb_k * (stages - 1);

1267 - };

1268 - if (slice_iters) {

1269 - start_pipes();

1270 - }

1271 -

1272 - // Main loop.

1273 - while (slice_iters) {

1274 - // We unroll over both the global fetch and the register load pipeline to

1275 - // ensure all shared memory accesses are static. Note that both pipelines

1276 - // have even length meaning that the next iteration will always start at

1277 - // index 0.

1278 - #pragma unroll

1279 - for (int pipe = 0; pipe < stages;) {

1280 - #pragma unroll

1281 - for (int k = 0; k < b_sh_wr_iters; k++) {

1282 - fetch_to_registers(k + 1, pipe % stages);

1283 - fetch_scales_to_registers(k + 1, pipe);

1284 - fetch_zp_to_registers(k + 1, pipe);

1285 - if (k == b_sh_wr_iters - 2) {

1286 - fetch_to_shared((pipe + stages - 1) % stages, pipe,

1287 - slice_iters >= stages);

1288 - pipe++;

1289 - wait_for_stage();

1290 - init_same_group(pipe % stages);

1291 - }

1292 - matmul(k);

1293 - }

1294 - slice_iters--;

1295 - if (slice_iters == 0) {

1296 - break;

1297 - }

1298 - }

1299 -

1300 - a_gl_rd += a_gl_rd_delta_o * stages;

1301 - slice_k_start += tb_k * stages;

1302 - slice_k_start_shared_fetch += tb_k * stages;

1303 -

1304 - if constexpr (has_act_order) {

1305 - int first_group_id = g_idx[slice_k_start];

1306 - int last_g_idx = slice_k_start + stages * tb_k * 2;

1307 - if (last_g_idx >= prob_k) {

1308 - last_g_idx = prob_k - 1;

1309 - }

1310 - int last_group_id = g_idx[last_g_idx];

1311 - if (last_group_id >= sh_first_group_id + sh_num_groups) {

1312 - fetch_scales_to_shared(false, first_group_id, last_group_id);

1313 - __syncthreads();

1314 - }

1315 - }

1316 -

1317 - // Process results and, if necessary, proceed to the next column slice.

1318 - // While this pattern may not be the most readable, other ways of writing

1319 - // the loop seemed to noticeably worse performance after compilation.

1320 - if (slice_iters == 0) {

1321 - cp_async_wait<0>();

1322 - bool last = slice_idx == slice_count - 1;

1323 - if constexpr (!has_act_order && group_blocks == -1) {

1324 - if constexpr (w_type.size_bits() == 8) {

1325 - if (s_sh_wr_pred) {

1326 - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);

1327 - }

1328 - cp_async_fence();

1329 - } else {

1330 - // For 4-bit per-column scales, we only fetch them here in the

1331 - // final step before write-out

1332 - if (last) {

1333 - if (s_sh_wr_pred) {

1334 - cp_async4(&sh_s[s_sh_wr], &scales_ptr[s_gl_rd]);

1335 - }

1336 - cp_async_fence();

1337 - }

1338 - }

1339 - }

1340 -

1341 - thread_block_reduce();

1342 - if constexpr (!has_act_order && group_blocks == -1) {

1343 - if constexpr (w_type.size_bits() == 8) {

1344 - cp_async_wait<0>();

1345 - __syncthreads();

1346 - if (threadIdx.x / 32 < thread_n_blocks / 4) {

1347 - reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];

1348 - reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];

1349 - }

1350 -

1351 - } else {

1352 - if (last) {

1353 - cp_async_wait<0>();

1354 - __syncthreads();

1355 - if (threadIdx.x / 32 < thread_n_blocks / 4) {

1356 - reinterpret_cast<int4*>(&frag_s)[0] = sh_s[s_sh_rd + 0];

1357 - reinterpret_cast<int4*>(&frag_s)[1] = sh_s[s_sh_rd + 4];

1358 - }

1359 - }

1360 - }

1361 - }

1362 -

1363 - // For 8-bit channelwise, we apply the scale before the global reduction

1364 - // that converts the fp32 results to fp16 (so that we avoid possible

1365 - // overflow in fp16)

1366 - if constexpr (!has_act_order && group_blocks == -1 &&

1367 - w_type.size_bits() == 8) {

1368 - if (threadIdx.x / 32 < thread_n_blocks / 4) {

1369 - #pragma unroll

1370 - for (int i = 0; i < thread_m_blocks; i++) {

1371 - #pragma unroll

1372 - for (int j = 0; j < 4; j++) {

1373 - scale_float(reinterpret_cast<float*>(&frag_c[i][j][0][0]),

1374 - frag_s[j / 2][2 * (j % 2) + 0]);

1375 - scale_float(reinterpret_cast<float*>(&frag_c[i][j][0][2]),

1376 - frag_s[j / 2][2 * (j % 2) + 0]);

1377 -

1378 - scale_float(reinterpret_cast<float*>(&frag_c[i][j][1][0]),

1379 - frag_s[j / 2][2 * (j % 2) + 1]);

1380 - scale_float(reinterpret_cast<float*>(&frag_c[i][j][1][2]),

1381 - frag_s[j / 2][2 * (j % 2) + 1]);

1382 - }

1383 - }

1384 - }

1385 - }

1386 -

1387 - if (slice_count > 1) { // only globally reduce if there is more than one

1388 - // block in a slice

1389 - barrier_acquire(&locks[slice_col], slice_idx);

1390 - global_reduce(slice_idx == 0, last);

1391 - barrier_release(&locks[slice_col], last);

1392 - }

1393 - if (last) // only the last block in a slice actually writes the result

1394 - write_result();

1395 - slice_row = 0;

1396 - slice_col_par++;

1397 - slice_col++;

1398 - init_slice();

1399 - if (slice_iters) {

1400 - a_gl_rd = a_gl_stride * (threadIdx.x / a_gl_rd_delta_o) +

1401 - (threadIdx.x % a_gl_rd_delta_o);

1402 - #pragma unroll

1403 - for (int i = 0; i < b_sh_wr_iters; i++)

1404 - B_ptr[i] += b_sh_stride - b_gl_rd_delta_o * k_tiles;

1405 - if (slice_col == 0) {

1406 - #pragma unroll

1407 - for (int i = 0; i < b_sh_wr_iters; i++) B_ptr[i] -= b_gl_stride;

1408 - }

1409 -

1410 - // Update slice k/n for scales loading

1411 - if constexpr (has_act_order) {

1412 - slice_k_start = tb_k * slice_row;

1413 - slice_k_finish = slice_k_start + tb_k * slice_iters;

1414 - slice_k_start_shared_fetch = slice_k_start;

1415 - slice_n_offset = act_s_col_tb_stride * slice_col;

1416 -

1417 - } else {

1418 - s_gl_rd = s_sh_stride * slice_col + threadIdx.x;

1419 - zp_gl_rd = zp_sh_stride * slice_col + threadIdx.x;

1420 - }

1421 -

1422 - start_pipes();

1423 - }

1424 - }

1425 - }

1426 -}

1427 -

1428 -template <const vllm::ScalarTypeId w_type_id, // weight ScalarType id

1429 - const int threads, // number of threads in a threadblock

1430 - const int thread_n_blocks, // same for n dimension (output)

1431 - const int thread_k_blocks, // same for k dimension (reduction)

1432 - const int stages, // number of stages for the async global->shared

1433 - // fetch pipeline

1434 - const bool has_act_order, // whether act_order is enabled

1435 - const bool has_zp, // whether zero-points are enabled

1436 - const int group_blocks = -1 // number of consecutive 16x16 blocks

1437 - // with a separate quantization scale

1438 - >

1439 -__global__ void MarlinMoE(

1440 - const int4* __restrict__ A, // fp16 input matrix of shape mxk

1441 - const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn

1442 - int4* __restrict__ C, // fp16 output buffer of shape mxn

1443 - const int* __restrict__ sorted_ids_base, // int32 sorted ids of experts

1444 - const float* __restrict__ topk_weights, // float topk weights

1445 - const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape

1446 - // (k/groupsize)xn

1447 - const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape

1448 - // (k/groupsize)x(n/pack_factor)

1449 - const int* __restrict__ g_idx, // int32 group indices of shape k

1450 - const int* __restrict__ expert_offsets,

1451 - int num_groups, // number of scale groups per output channel

1452 - int expert_idx, // idx of current expert

1453 - int num_experts, // number of experts

1454 - int topk, // topk parameter of moe

1455 - int prob_m, // batch dimension m

1456 - int prob_n, // output dimension n

1457 - int prob_k, // reduction dimension k

1458 - int tot_m, // total number of rows in A and C

1459 - int* locks, // extra global storage for barrier synchronization

1460 - bool replicate_input, // do we use the same input for each expert?

1461 - bool apply_weights, // apply weights to output

1462 - int current_m_block, // current m block to start kernel computation from

1463 - int max_par, // maximum parallelism

1464 - int cfg_max_m_blocks // upper bound on m blocks

1465 -) {

1466 - int m_block_ctr = current_m_block;

1467 -

1468 - const int* sorted_ids_expert =

1469 - sorted_ids_base + expert_offsets[expert_idx] + m_block_ctr * 4 * max_par;

1470 - int tot_its = expert_offsets[expert_idx + 1] - expert_offsets[expert_idx];

1471 - if (tot_its == 0) {

1472 - return;

1473 - }

1474 - int tot_m_blocks = ceildiv(tot_its, 16);

1475 - int pad = 16 * tot_m_blocks - tot_its;

1476 -

1477 - if (m_block_ctr >= tot_m_blocks) {

1478 - return;

1479 - }

1480 -

1481 - int max_block = tot_m_blocks - m_block_ctr;

1482 - prob_m = tot_its - 16 * m_block_ctr;

1483 -

1484 - int par = 1;

1485 - if (max_block > cfg_max_m_blocks) {

1486 - // Note that parallel > 1 currently only works for inputs without any

1487 - // padding

1488 - par = (16 * max_block - pad) / (16 * cfg_max_m_blocks);

1489 - if (par > max_par) par = max_par;

1490 - prob_m = (16 * cfg_max_m_blocks) * par;

1491 - m_block_ctr += cfg_max_m_blocks * (par - 1);

1492 - max_block = cfg_max_m_blocks;

1493 - }

1494 -

1495 - if (max_block == 1) {

1496 - MarlinMoESingle<w_type_id, threads, 1, thread_n_blocks, thread_k_blocks,

1497 - stages, has_act_order, has_zp, group_blocks>(

1498 - A, B, C, sorted_ids_expert, topk_weights, scales_ptr, zp_ptr, g_idx,

1499 - expert_offsets, num_groups, expert_idx, num_experts, topk, prob_m,

1500 - prob_n, prob_k, tot_m, locks, replicate_input, apply_weights,

1501 - current_m_block);

1502 - } else if (max_block == 2) {

1503 - MarlinMoESingle<w_type_id, threads, 2, thread_n_blocks, thread_k_blocks,

1504 - stages, has_act_order, has_zp, group_blocks>(

1505 - A, B, C, sorted_ids_expert, topk_weights, scales_ptr, zp_ptr, g_idx,

1506 - expert_offsets, num_groups, expert_idx, num_experts, topk, prob_m,

1507 - prob_n, prob_k, tot_m, locks, replicate_input, apply_weights,

1508 - current_m_block);

1509 - } else if (max_block == 3) {

1510 - MarlinMoESingle<w_type_id, threads, 3, thread_n_blocks, thread_k_blocks,

1511 - stages, has_act_order, has_zp, group_blocks>(

1512 - A, B, C, sorted_ids_expert, topk_weights, scales_ptr, zp_ptr, g_idx,

1513 - expert_offsets, num_groups, expert_idx, num_experts, topk, prob_m,

1514 - prob_n, prob_k, tot_m, locks, replicate_input, apply_weights,

1515 - current_m_block);

1516 - } else {

1517 - MarlinMoESingle<w_type_id, threads, 4, thread_n_blocks, thread_k_blocks,

1518 - stages, has_act_order, has_zp, group_blocks>(

1519 - A, B, C, sorted_ids_expert, topk_weights, scales_ptr, zp_ptr, g_idx,

1520 - expert_offsets, num_groups, expert_idx, num_experts, topk, prob_m,

1521 - prob_n, prob_k, tot_m, locks, replicate_input, apply_weights,

1522 - current_m_block);

1523 - }

1524 -}

1525 -

1526 -#else

1527 -

1528 -template <const vllm::ScalarTypeId w_type_id, // weight ScalarType id

1529 - const int threads, // number of threads in a threadblock

1530 - const int thread_n_blocks, // same for n dimension (output)

1531 - const int thread_k_blocks, // same for k dimension (reduction)

1532 - const int stages, // number of stages for the async global->shared

1533 - // fetch pipeline

1534 - const bool has_act_order, // whether act_order is enabled

1535 - const bool has_zp, // whether zero-points are enabled

1536 - const int group_blocks = -1 // number of consecutive 16x16 blocks

1537 - // with a separate quantization scale

1538 - >

1539 -__global__ void MarlinMoE(

1540 - const int4* __restrict__ A, // fp16 input matrix of shape mxk

1541 - const int4* __restrict__ B, // 4bit quantized weight matrix of shape kxn

1542 - int4* __restrict__ C, // fp16 output buffer of shape mxn

1543 - const int* __restrict__ sorted_ids, // int32 sorted ids of experts

1544 - const float* __restrict__ topk_weights, // float topk weights

1545 - const int4* __restrict__ scales_ptr, // fp16 quantization scales of shape

1546 - // (k/groupsize)xn

1547 - const int4* __restrict__ zp_ptr, // 4bit packed zero-points of shape

1548 - // (k/groupsize)x(n/pack_factor)

1549 - const int* __restrict__ g_idx, // int32 group indices of shape k

1550 - const int* __restrict__ expert_offsets,

1551 - int num_groups, // number of scale groups per output channel

1552 - int expert_idx, // idx of current expert

1553 - int num_experts, // number of experts

1554 - int topk, // topk parameter of moe

1555 - int prob_m, // batch dimension m

1556 - int prob_n, // output dimension n

1557 - int prob_k, // reduction dimension k

1558 - int tot_m, // total number of rows in A and C

1559 - int* locks, // extra global storage for barrier synchronization

1560 - bool replicate_input, // do we use the same input for each expert?

1561 - bool apply_weights, // apply weights to output

1562 - int current_m_block, // current m block to start kernel computation from

1563 - int max_par, // maximum parallelism

1564 - int cfg_max_m_blocks // upper bound on m blocks

1565 -) {

1566 - // Marlin is not implemented yet for SM < 8.0

1567 - assert(false);

1568 - return;

1569 -}

1570 -

1571 -#endif

1572 -

1573 -// 8 warps are a good choice since every SM has 4 schedulers and having more

1574 -// than 1 warp per schedule allows some more latency hiding. At the same time,

1575 -// we want relatively few warps to have many registers per warp and small tiles.

1576 -const int USER_THREADS =

1577 - 256; // Note: This is only used with user-provided thread_k/n

1578 -const int STAGES = 4; // 4 pipeline stages fit into shared memory

1579 -

1580 -static constexpr int min_thread_n = 64;

1581 -static constexpr int min_thread_k = 64;

1582 -

1583 -#define __CALL_IF_MOE(W_TYPE, THREAD_N_BLOCKS, THREAD_K_BLOCKS, HAS_ACT_ORDER, \

1584 - HAS_ZP, GROUP_BLOCKS, NUM_THREADS) \

1585 - else if (q_type == W_TYPE && thread_n_blocks == THREAD_N_BLOCKS && \

1586 - thread_k_blocks == THREAD_K_BLOCKS && \

1587 - has_act_order == HAS_ACT_ORDER && has_zp == HAS_ZP && \

1588 - group_blocks == GROUP_BLOCKS && num_threads == NUM_THREADS) { \

1589 - cudaFuncSetAttribute( \

1590 - MarlinMoE<W_TYPE.id(), NUM_THREADS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \

1591 - STAGES, HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS>, \

1592 - cudaFuncAttributeMaxDynamicSharedMemorySize, max_shared_mem); \

1593 - MarlinMoE<W_TYPE.id(), NUM_THREADS, THREAD_N_BLOCKS, THREAD_K_BLOCKS, \

1594 - STAGES, HAS_ACT_ORDER, HAS_ZP, GROUP_BLOCKS> \

1595 - <<<blocks, NUM_THREADS, max_shared_mem, stream>>>( \

1596 - A_ptr, B_ptr, C_ptr, sorted_ids_ptr, topk_weights_ptr, s_ptr, \

1597 - zp_ptr, g_idx_ptr, expert_offsets_ptr, num_groups, expert_idx, \

1598 - num_experts, topk, prob_m, prob_n, prob_k, tot_m, locks, \

1599 - replicate_input, apply_weights, m_block, max_par, \

1600 - cfg_max_m_blocks); \

1601 - }

1602 -

1603 -#define GPTQ_CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \

1604 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, true, false, 0, NUM_THREADS) \

1605 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, false, -1, NUM_THREADS) \

1606 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, false, 2, NUM_THREADS) \

1607 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, false, 4, NUM_THREADS) \

1608 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, false, 8, NUM_THREADS)

1609 -

1610 -#define AWQ_CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, NUM_THREADS) \

1611 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, true, -1, NUM_THREADS) \

1612 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, true, 2, NUM_THREADS) \

1613 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, true, 4, NUM_THREADS) \

1614 - __CALL_IF_MOE(W_TYPE, N_BLOCKS, K_BLOCKS, false, true, 8, NUM_THREADS)

1615 -

1616 -} // namespace marlin_moe

csrc/moe/marlin_kernels/marlin_moe_kernel_ku4.cu deleted

-31

@@ -1,31 +0,0 @@

1 -#include "marlin_moe_kernel_ku4.h"

2 -

3 -namespace marlin_moe {

4 -

5 -// We return bool so we can create these different kernel calls as a sequence

6 -// of if-elseif's.

7 -bool call_marlin_moe_kernel_ku4(

8 - vllm::ScalarType const& q_type, int thread_n_blocks, int thread_k_blocks,

9 - bool has_act_order, int group_blocks, int num_threads, int blocks,

10 - int max_shared_mem, cudaStream_t stream, const int4* A_ptr,

11 - const int4* B_ptr, int4* C_ptr, const int* sorted_ids_ptr,

12 - const float* topk_weights_ptr, const int4* s_ptr, const int4* zp_ptr,

13 - const int* g_idx_ptr, int* expert_offsets_ptr, int num_groups,

14 - int expert_idx, int num_experts, int topk, int prob_m, int prob_n,

15 - int prob_k, int tot_m, int* locks, bool replicate_input, bool apply_weights,

16 - int m_block, int max_par, int cfg_max_m_blocks) {

17 - bool has_zp = true;

18 -

19 - if (false) {

20 - }

21 - AWQ_CALL_IF_MOE(vllm::kU4, 16, 4, 256)

22 - AWQ_CALL_IF_MOE(vllm::kU4, 8, 8, 256)

23 - AWQ_CALL_IF_MOE(vllm::kU4, 8, 4, 128)

24 - AWQ_CALL_IF_MOE(vllm::kU4, 4, 8, 128)

25 - else {

26 - return false;

27 - }

28 - return true;

29 -}

30 -

31 -} // namespace marlin_moe

csrc/moe/marlin_kernels/marlin_moe_kernel_ku4.h deleted

-20

@@ -1,20 +0,0 @@

1 -#pragma once

2 -

3 -#include "marlin_moe_kernel.h"

4 -

5 -namespace marlin_moe {

6 -

7 -// We return bool so we can create these different kernel calls as a sequence

8 -// of if-elseif's.

9 -bool call_marlin_moe_kernel_ku4(

10 - vllm::ScalarType const& q_type, int thread_n_blocks, int thread_k_blocks,

11 - bool has_act_order, int group_blocks, int num_threads, int blocks,

12 - int max_shared_mem, cudaStream_t stream, const int4* A_ptr,

13 - const int4* B_ptr, int4* C_ptr, const int* sorted_ids_ptr,

14 - const float* topk_weights_ptr, const int4* s_ptr, const int4* zp_ptr,

15 - const int* g_idx_ptr, int* expert_offsets_ptr, int num_groups,

16 - int expert_idx, int num_experts, int topk, int prob_m, int prob_n,

17 - int prob_k, int tot_m, int* locks, bool replicate_input, bool apply_weights,

18 - int m_block, int max_par, int cfg_max_m_blocks);

19 -

20 -} // namespace marlin_moe

csrc/moe/marlin_kernels/marlin_moe_kernel_ku4b8.cu deleted

-31

@@ -1,31 +0,0 @@

1 -#include "marlin_moe_kernel_ku4b8.h"

2 -

3 -namespace marlin_moe {

4 -

5 -// We return bool so we can create these different kernel calls as a sequence

6 -// of if-elseif's.

7 -bool call_marlin_moe_kernel_ku4b8(