[Kernel] optimize performance of gptq marlin kernel when n is small

[jinzhen-lin]

The performance of gptq/awq marlin kernel is low when n is small. The reason is that when n is small, the barrier_acquire and barrier_release cost too many time.

The case where n is small often occurs during inference using larger tensor parallelism or in some special models. For example, in deepseek-v3 model, the shape of kv_a_proj_with_mqa is k=7168 and n=576.

This PR add an argument use_atomic_add to marlin kernel, so that we don't require barrier_acquire and barrier_release anymore.

Some simple benchmarks:

import torch
import triton

import pandas as pd
import vllm._custom_ops as vllm_ops
from vllm.model_executor.layers.quantization.utils.marlin_utils_test import \
    awq_marlin_quantize
from vllm.model_executor.layers.quantization.utils.quant_utils import \
    quantize_weights
from vllm.scalar_type import scalar_types


def bench_gptq_gemm(inputs, weight):
    group_size = 128
    quant_type = quant_type = scalar_types.uint4
    _, qweight, scales, qzeros = quantize_weights(
        weight, quant_type, group_size, True, False)
    g_idx = torch.empty((0,), device=qweight.device)

    def run():
        return vllm_ops.gptq_gemm(inputs, qweight, qzeros, scales, g_idx, True, 4)

    with torch.cuda.stream(torch.cuda.Stream()):
        t = triton.testing.do_bench(run, rep=1000)
    return t


def bench_gptq_marlin_gemm(inputs, weight, use_atomic_add=False):
    group_size = 128

    _, qweight, scales, qzeros = awq_marlin_quantize(
        weight, scalar_types.uint4, group_size)

    g_idx = torch.empty(0, dtype=torch.int, device=qweight.device)
    sort_indices = torch.empty(0, dtype=torch.int, device=qweight.device)
    workspace = torch.zeros_like(qweight)

    def run():
        return vllm_ops.gptq_marlin_gemm(
            inputs,
            qweight,
            scales,
            qzeros,
            g_idx,
            sort_indices,
            workspace,
            scalar_types.uint4,
            inputs.shape[0],
            weight.shape[1],
            inputs.shape[1],
            is_k_full=True,
            has_zp=True,
            use_atomic_add=use_atomic_add,
            use_fp32_reduce=False,
            is_zp_float=False,
        )

    with torch.cuda.stream(torch.cuda.Stream()):
        t = triton.testing.do_bench(run, rep=1000)
    return t


def bench_torch_gemm(inputs, weight):
    def run():
        return inputs.matmul(weight)

    with torch.cuda.stream(torch.cuda.Stream()):
        t = triton.testing.do_bench_cudagraph(run, rep=1000)
    return t


shapes = [
    (16, 8192, 128),
    (16, 8192, 256),
    (16, 8192, 512),
    (16, 8192, 1024),
    (16, 8192, 2048),
    (16, 8192, 4096),
    (16, 8192, 8192),
    (16, 8192, 16384),
    (16, 8192, 32768)
]

res = []
for m, k, n in shapes:
    inputs = torch.randn((m, k), device="cuda:0", dtype=torch.half) / 10
    weight = torch.randn((k, n), device="cuda:0", dtype=torch.half) / 10
    res.append({
        "(m, k, n)": str((m, k, n)),
        "torch": bench_torch_gemm(inputs, weight),
        "gptq": bench_gptq_gemm(inputs, weight),
        "marlin-atomic-add-false": bench_gptq_marlin_gemm(inputs, weight, False),
        "marlin-atomic-add-true": bench_gptq_marlin_gemm(inputs, weight, True),
    })


print(pd.DataFrame(res).to_markdown(index=False))

(m, k, n)	torch	gptq	marlin-atomic-add-false	marlin-atomic-add-true
(16, 8192, 128)	0.00746296	0.0248551	0.095163	0.0159413
(16, 8192, 256)	0.00754125	0.0256122	0.0539876	0.0168648
(16, 8192, 512)	0.0091006	0.02665	0.0413818	0.0188273
(16, 8192, 1024)	0.0130155	0.0347707	0.0320496	0.0213404
(16, 8192, 2048)	0.0203446	0.0416565	0.0272889	0.024638
(16, 8192, 4096)	0.0497934	0.0682435	0.0300557	0.0323643
(16, 8192, 8192)	0.0888581	0.115765	0.0467808	0.0503323
(16, 8192, 16384)	0.17135	0.212554	0.0737511	0.0780345
(16, 8192, 32768)	0.326693	0.404916	0.12399	0.130165

The benchmark result shows that atomicAdd can improve that performance when n <=2048. However, this atomicAdd use fp16, I don't sure if the issue fixed by #6795 would be introduced again (need more test). So I disable the atomicAdd by default, and use the env variable VLLM_MARLIN_USE_ATOMIC_ADD to control the behavior.

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[jinzhen-lin]

Benchmark results of deepseek-v3-awq on 8*A800 (tokens/s):

bs	#13321	#13321 + this PR
1	46.2	50.2
2	79.1	84.1
4	124.8	130.7
8	183.6	190.2
16	261.5	267.4
32	356.5	368.8
64	480.1	503.5

Benchmark results of deepseek-r1-awq on 8*A800 (tokens/s):

bs	#13321	#13321 + this PR
1	44.8	49.4
2	76.5	83.4
4	118.7	126.5
8	171.3	180.3
16	240.8	249.5
32	332.5	361.2
64	454.6	478.6

Note:

1. The result is tested with VLLM_MLA_DISABLE=1
2. The baseline speed is slower than the result in [Kernel] moe wna16 cuda kernel #13321 since A800 is slightly slower than A100.

Accuracy test:

VLLM_MARLIN_USE_ATOMIC_ADD=1 lm_eval --model vllm     --model_args pretrained=/root/DeepSeek-R1-AWQ/,tensor_parallel_size=8,gpu_memory_utilization=0.95,trust_remote_code=True,max_model_len=16384,quantization=moe_wna16,dtype=half,max_num_batched_tokens=16384     --task gsm8k   --num_fewshot 5 --batch_size auto

|Tasks|Version|     Filter     |n-shot|  Metric   |   |Value |   |Stderr|
|-----|------:|----------------|-----:|-----------|---|-----:|---|-----:|
|gsm8k|      3|flexible-extract|     5|exact_match|↑  |0.9568|±  |0.0056|
|     |       |strict-match    |     5|exact_match|↑  |0.9553|±  |0.0057|

[mgoin]

This seems reasonable and the speedups are compelling, thank you @jinzhen-lin !
I think it would be best to enable this by default rather than hide it behind an env var for most users to observe. However, we can keep this for now and try to remove once we better understand heuristics on more hardware.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @jinzhen-lin.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[jinzhen-lin]

@mgoin All tests are passed, can you merge it ?

[mgoin]

Yes, thanks again @jinzhen-lin