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

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, 19:45 UTC

[Kernel] optimize performance of gptq marlin kernel when n is small (#14138)

Signed-off-by: Jinzhen Lin <linjinzhen@hotmail.com>

Jinzhen Lin committed Mar 8, 2025 at 00:53 UTC d0feea31c70f9540a8993c4e96103a03cd935416

6 files changed +99 -24

csrc/quantization/gptq_marlin/gptq_marlin.cu

+46 -16

@@ -538,6 +538,7 @@ __global__ void Marlin(

538 int prob_n, // output dimension n

539 int prob_k, // reduction dimension k

540 int* locks, // extra global storage for barrier synchronization

541 + bool use_atomic_add, // whether to use atomic add to reduce

542 bool use_fp32_reduce // whether to use fp32 global reduce

543 ) {

544 // Each threadblock processes one "stripe" of the B matrix with (roughly) the

@@ -1542,7 +1543,17 @@ __global__ void Marlin(

1543 i < div_ceil(16 * thread_m_blocks, threads / (2 * thread_n_blocks));

1544 i++) {

1545 if (c_gl_wr < c_gl_wr_end) {

1545 - C[c_gl_wr] = sh_red[c_sh_rd];

1546 + if (use_atomic_add && slice_count > 1) {

1547 + scalar_t2* C_half2 = reinterpret_cast<scalar_t2*>(&C[c_gl_wr]);

1548 + scalar_t2* sh_red_half2 =

1549 + reinterpret_cast<scalar_t2*>(&sh_red[c_sh_rd]);

1550 + #pragma unroll

1551 + for (int a = 0; a < 4; a++) {

1552 + atomicAdd(&C_half2[a], sh_red_half2[a]);

1553 + }

1554 + } else {

1555 + C[c_gl_wr] = sh_red[c_sh_rd];

1556 + }

1557 c_gl_wr += c_gl_wr_delta;

1558 c_sh_rd += c_sh_rd_delta;

1559 }

@@ -1644,7 +1655,7 @@ __global__ void Marlin(

1655 }

1656 cp_async_fence();

1657 } else {

1647 - if (last) {

1658 + if (last || use_atomic_add) {

1659 if (s_sh_wr_pred) {

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

1661 }

@@ -1664,7 +1675,7 @@ __global__ void Marlin(

1675 }

1676

1677 } else {

1667 - if (last) {

1678 + if (last || use_atomic_add) {

1679 cp_async_wait<0>();

1680 __syncthreads();

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

@@ -1703,8 +1714,8 @@ __global__ void Marlin(

1714 }

1715 }

1716

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

1707 - // block in a slice

1717 + if (slice_count > 1 && !use_atomic_add) {

1718 + // only globally reduce if there is more than one block in a slice

1719 barrier_acquire(&locks[slice_col], slice_idx);

1720 if (use_fp32_reduce) {

1721 global_reduce_fp32(slice_idx == 0, last);

@@ -1713,7 +1724,8 @@ __global__ void Marlin(

1724 }

1725 barrier_release(&locks[slice_col], last);

1726 }

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

1727 + if (last || use_atomic_add)

1728 + // only the last block in a slice actuallywrites the result

1729 write_result();

1730 slice_row = 0;

1731 slice_col_par++;

@@ -1768,7 +1780,8 @@ __global__ void Marlin(

1780 HAS_ZP, GROUP_BLOCKS, IS_ZP_FLOAT> \

1781 <<<blocks, NUM_THREADS, max_shared_mem, stream>>>( \

1782 A_ptr, B_ptr, C_ptr, C_tmp_ptr, s_ptr, zp_ptr, g_idx_ptr, \

1771 - num_groups, prob_m, prob_n, prob_k, locks, use_fp32_reduce); \

1783 + num_groups, prob_m, prob_n, prob_k, locks, use_atomic_add, \

1784 + use_fp32_reduce); \

1785 } \

1786 }

1787

@@ -2062,7 +2075,8 @@ void marlin_mm(const void* A, const void* B, void* C, void* C_tmp, void* s,

2075 vllm::ScalarType const& q_type, bool has_act_order,

2076 bool is_k_full, bool has_zp, int num_groups, int group_size,

2077 int dev, cudaStream_t stream, int thread_k, int thread_n,

2065 - int sms, int max_par, bool use_fp32_reduce, bool is_zp_float) {

2078 + int sms, int max_par, bool use_atomic_add, bool use_fp32_reduce,

2079 + bool is_zp_float) {

2080 if (has_zp) {

2081 TORCH_CHECK(

2082 q_type == vllm::kU4 || q_type == vllm::kU8,

@@ -2243,7 +2257,7 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,

2257 torch::Tensor& workspace,

2258 vllm::ScalarTypeId const& b_q_type_id,

2259 int64_t size_m, int64_t size_n, int64_t size_k,

2246 - bool is_k_full, bool has_zp,

2260 + bool is_k_full, bool has_zp, bool use_atomic_add,

2261 bool use_fp32_reduce, bool is_zp_float) {

2262 vllm::ScalarType const b_q_type = vllm::ScalarType::from_id(b_q_type_id);

2263 if (has_zp) {

@@ -2306,19 +2320,34 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,

2320 // Alloc buffers

2321 const at::cuda::OptionalCUDAGuard device_guard(device_of(a));

2322 auto options = torch::TensorOptions().dtype(a.dtype()).device(a.device());

2309 - torch::Tensor c = torch::empty({size_m, size_n}, options);

2310 - torch::Tensor a_tmp = torch::empty({size_m, size_k}, options);

2323 + torch::Tensor c;

2324 + if (use_atomic_add) {

2325 + c = torch::zeros({size_m, size_n}, options);

2326 + } else {

2327 + c = torch::empty({size_m, size_n}, options);

2328 + }

2329 +

2330 + torch::Tensor a_tmp;

2331 + bool has_act_order = g_idx.size(0) != 0;

2332 + if (has_act_order) {

2333 + a_tmp = torch::empty({size_m, size_k}, options);

2334 + } else {

2335 + a_tmp = torch::empty({0}, options);

2336 + }

2337

2338 // Alloc C tmp buffer that is going to be used for the global reduce

2339 + torch::Tensor c_tmp;

2340 int reduce_max_m = marlin::determine_reduce_max_m(size_m, marlin::max_par);

2341 int reduce_n = size_n;

2342 auto options_fp32 =

2343 torch::TensorOptions().dtype(at::kFloat).device(a.device());

2317 - if (!use_fp32_reduce) {

2344 + if (use_fp32_reduce) {

2345 + c_tmp = torch::empty({reduce_max_m, reduce_n}, options_fp32);

2346 + } else {

2347 reduce_max_m = 0;

2348 reduce_n = 0;

2349 + c_tmp = torch::empty({0}, options_fp32);

2350 }

2321 - torch::Tensor c_tmp = torch::empty({reduce_max_m, reduce_n}, options_fp32);

2351

2352 // thread_k: `k` size of a thread_tile in `weights` (can usually be left as

2353 // auto -1)

@@ -2339,7 +2368,6 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,

2368 // Detect groupsize and act_order

2369 int num_groups = -1;

2370 int group_size = -1;

2342 - bool has_act_order = g_idx.size(0) != 0;

2371

2372 int rank = b_scales.sizes().size();

2373 TORCH_CHECK(rank == 2, "b_scales rank = ", rank, " is not 2");

@@ -2407,7 +2435,8 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,

2435 a_tmp.data_ptr<at::Half>(), size_m, size_n, size_k,

2436 workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,

2437 num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),

2410 - thread_k, thread_n, sms, marlin::max_par, use_fp32_reduce, is_zp_float);

2438 + thread_k, thread_n, sms, marlin::max_par, use_atomic_add,

2439 + use_fp32_reduce, is_zp_float);

2440 } else if (a.scalar_type() == at::ScalarType::BFloat16) {

2441 marlin::marlin_mm<nv_bfloat16>(

2442 a.data_ptr<at::BFloat16>(), b_q_weight.data_ptr(),

@@ -2416,7 +2445,8 @@ torch::Tensor gptq_marlin_gemm(torch::Tensor& a, torch::Tensor& b_q_weight,

2445 perm.data_ptr(), a_tmp.data_ptr<at::BFloat16>(), size_m, size_n, size_k,

2446 workspace.data_ptr(), b_q_type, has_act_order, is_k_full, has_zp,

2447 num_groups, group_size, dev, at::cuda::getCurrentCUDAStream(dev),

2419 - thread_k, thread_n, sms, marlin::max_par, use_fp32_reduce, is_zp_float);

2448 + thread_k, thread_n, sms, marlin::max_par, use_atomic_add,

2449 + use_fp32_reduce, is_zp_float);

2450 } else {

2451 TORCH_CHECK(false, "gpt_marlin_gemm only supports bfloat16 and float16");

2452 }

csrc/torch_bindings.cpp

+2 -1

@@ -272,7 +272,8 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {

272 "Tensor b_zeros, Tensor g_idx, Tensor perm, Tensor workspace, "

273 "int b_q_type, "

274 "SymInt size_m, SymInt size_n, SymInt size_k, bool is_k_full, "

275 - "bool has_zp, bool use_fp32_reduce, bool is_zp_float) -> Tensor",

275 + "bool has_zp, bool use_atomic_add, bool use_fp32_reduce, "

276 + "bool is_zp_float) -> Tensor",

277 {stride_tag});

278 // conditionally compiled so impl registration is in source file

279

tests/kernels/test_marlin_gemm.py

+10 -6

@@ -34,6 +34,7 @@ from vllm.scalar_type import scalar_types

35 ACT_ORDER_OPTS = [False, True]

36 K_FULL_OPTS = [False, True]

37 +USE_ATOMIC_ADD_OPTS = [False, True]

38 USE_FP32_REDUCE_OPTS = [False, True]

40 MARLIN_K_CHUNKS = [128]

@@ -194,6 +195,7 @@ def test_awq_marlin_repack(k_chunk, n_chunk, quant_type, group_size,

195 @pytest.mark.parametrize("mnk_factors", MNK_FACTORS)

196 @pytest.mark.parametrize("act_order", ACT_ORDER_OPTS)

197 @pytest.mark.parametrize("is_k_full", K_FULL_OPTS)

198 +@pytest.mark.parametrize("use_atomic_add", USE_ATOMIC_ADD_OPTS)

199 @pytest.mark.parametrize("use_fp32_reduce", USE_FP32_REDUCE_OPTS)

200 def test_gptq_marlin_gemm(

201 k_chunk,

@@ -203,6 +205,7 @@ def test_gptq_marlin_gemm(

205 mnk_factors,

206 act_order,

207 is_k_full,

208 + use_atomic_add,

209 use_fp32_reduce,

210 ):

211 m_factor, n_factor, k_factor = mnk_factors

@@ -228,12 +231,12 @@ def test_gptq_marlin_gemm(

231 workspace = MarlinWorkspace(size_n, GPTQ_MARLIN_MIN_THREAD_N,

232 GPTQ_MARLIN_MAX_PARALLEL)

233

231 - opcheck(

232 - torch.ops._C.gptq_marlin_gemm,

233 - (a_input, marlin_q_w, marlin_s, marlin_zp, g_idx, sort_indices,

234 - workspace.scratch, quant_type.id, a_input.shape[0], b_weight.shape[1],

235 - a_input.shape[1], is_k_full, False, use_fp32_reduce, False),

236 - test_utils=DEFAULT_OPCHECK_TEST_UTILS)

234 + opcheck(torch.ops._C.gptq_marlin_gemm,

235 + (a_input, marlin_q_w, marlin_s, marlin_zp, g_idx, sort_indices,

236 + workspace.scratch, quant_type.id, a_input.shape[0],

237 + b_weight.shape[1], a_input.shape[1], is_k_full, False,

238 + use_atomic_add, use_fp32_reduce, False),

239 + test_utils=DEFAULT_OPCHECK_TEST_UTILS)

240

241 output = ops.gptq_marlin_gemm(

242 a_input,

@@ -249,6 +252,7 @@ def test_gptq_marlin_gemm(

252 a_input.shape[1],

253 is_k_full=is_k_full,

254 has_zp=False,

255 + use_atomic_add=use_atomic_add,

256 use_fp32_reduce=use_fp32_reduce,

257 is_zp_float=False,

258 )

vllm/_custom_ops.py

+4 -1

@@ -301,6 +301,7 @@ if hasattr(torch.ops._C, "gptq_marlin_24_gemm"):

301 size_k: torch.SymInt,

302 is_k_full: bool,

303 has_zp: bool = False,

304 + use_atomic_add: bool = False,

305 use_fp32_reduce: bool = False,

306 is_zp_float: bool = False) -> torch.Tensor:

307 return torch.empty((size_m, size_n), device=a.device, dtype=a.dtype)

@@ -713,12 +714,14 @@ def gptq_marlin_gemm(a: torch.Tensor,

714 size_k: int,

715 is_k_full: bool,

716 has_zp: bool = False,

717 + use_atomic_add: bool = False,

718 use_fp32_reduce: bool = False,

719 is_zp_float: bool = False) -> torch.Tensor:

720 return torch.ops._C.gptq_marlin_gemm(a, b_q_weight, b_scales, b_zeros,

721 g_idx, perm, workspace, b_q_type.id,

722 size_m, size_n, size_k, is_k_full,

721 - has_zp, use_fp32_reduce, is_zp_float)

723 + has_zp, use_atomic_add,

724 + use_fp32_reduce, is_zp_float)

725

726

727 # fp8 marlin

vllm/envs.py

@@ -95,6 +95,7 @@ if TYPE_CHECKING:

95 VLLM_DP_SIZE: int = 1

96 VLLM_DP_MASTER_IP: str = ""

97 VLLM_DP_MASTER_PORT: int = 0

98 + VLLM_MARLIN_USE_ATOMIC_ADD: bool = False

100

101 def get_default_cache_root():

@@ -630,6 +631,10 @@ environment_variables: dict[str, Callable[[], Any]] = {

631 # Whether to use S3 path for model loading in CI via RunAI Streamer

632 "VLLM_CI_USE_S3":

633 lambda: os.environ.get("VLLM_CI_USE_S3", "0") == "1",

634 +

635 + # Whether to use atomicAdd reduce in gptq/awq marlin kernel.

636 + "VLLM_MARLIN_USE_ATOMIC_ADD":

637 + lambda: os.environ.get("VLLM_MARLIN_USE_ATOMIC_ADD", "0") == "1",

638 }

639

640 # end-env-vars-definition

vllm/model_executor/layers/quantization/utils/marlin_utils.py

+32

@@ -5,6 +5,7 @@ from typing import List, Optional, Tuple

5 import numpy

6 import torch

8 +import vllm.envs as envs

9 from vllm import _custom_ops as ops

10 from vllm.model_executor.layers.linear import LinearBase

11 from vllm.platforms import current_platform

@@ -290,6 +291,23 @@ def moe_awq_to_marlin_zero_points(q_zp_packed: torch.Tensor, size_k: int,

291 return output

292

293

294 +def should_use_atomic_add_reduce(m: int, n: int, k: int, device: torch.device,

295 + dtype: torch.dtype) -> bool:

296 + # disable atomicAdd reduce by default,

297 + # one can enable it with VLLM_MARLIN_USE_ATOMIC_ADD=1

298 + if not envs.VLLM_MARLIN_USE_ATOMIC_ADD or device.type != "cuda":

299 + return False

300 +

301 + # sm8x doesn't support atomicAdd + bfloat16 natively

302 + device_capability = torch.cuda.get_device_capability(device)

303 + if device_capability[0] < 9 and dtype == torch.bfloat16:

304 + return False

305 +

306 + # the performance of atomicAdd is better than global reduce

307 + # only when m*n is small and k is large

308 + return max(m, 64) * n < 64 * 2048 and k >= 2048

309 +

310 +

311 def apply_gptq_marlin_linear(

312 input: torch.Tensor,

313 weight: torch.Tensor,

@@ -307,6 +325,12 @@ def apply_gptq_marlin_linear(

325 reshaped_x = input.reshape(-1, input.shape[-1])

326 out_shape = input.shape[:-1] + (output_size_per_partition, )

327

328 + use_atomic_add = should_use_atomic_add_reduce(m=reshaped_x.size(0),

329 + n=output_size_per_partition,

330 + k=reshaped_x.size(1),

331 + device=input.device,

332 + dtype=input.dtype)

333 +

334 output = ops.gptq_marlin_gemm(reshaped_x,

335 weight,

336 weight_scale,

@@ -320,6 +344,7 @@ def apply_gptq_marlin_linear(

344 size_k=input_size_per_partition,

345 is_k_full=is_k_full,

346 has_zp=False,

347 + use_atomic_add=use_atomic_add,

348 use_fp32_reduce=use_fp32_reduce,

349 is_zp_float=False)

350

@@ -345,6 +370,12 @@ def apply_awq_marlin_linear(

370 reshaped_x = input.reshape(-1, input.shape[-1])

371 out_shape = input.shape[:-1] + (output_size_per_partition, )

372

373 + use_atomic_add = should_use_atomic_add_reduce(m=reshaped_x.size(0),

374 + n=output_size_per_partition,

375 + k=reshaped_x.size(1),

376 + device=input.device,

377 + dtype=input.dtype)

378 +

379 output = ops.gptq_marlin_gemm(reshaped_x,

380 weight,

381 weight_scale,

@@ -358,6 +389,7 @@ def apply_awq_marlin_linear(

389 size_k=input_size_per_partition,

390 is_k_full=True,

391 has_zp=True,

392 + use_atomic_add=use_atomic_add,

393 use_fp32_reduce=use_fp32_reduce,

394 is_zp_float=False)

395