[Feature]: [vllm_gguf_plugin] Add --gguf-dequant-on-load option to unlock cuBLAS tensor cores for prefill-heavy workloads`
The feature, motivation and pitch
Motivation
vllm_gguf_plugin is popular for running quantized GGUF weights. During autoregressive decoding ($N=1$), memory bandwidth is the bottleneck, and on-the-fly dequantization (ops.ggml_mul_mat_a8) is fast.
However, for prefill-heavy workloads (such as long-context prompts, batch embeddings, or dataset scoring on $T=8,192$), linear projections are compute-bound.
On an RTX 4090 running Llama-3.1-8B-Q8_0 on an 8,192-token prompt:
ggml_mul_mat_a8takes ~18.2 ms per GEMM ➔ 15.03 seconds total model forward pass.- Native
bfloat16cuBLAS tensor cores take ~1.7 ms per GEMM ➔ 1.18 seconds total model forward pass (12.7x speedup).
Proposed Enhancement
Add an optional flag VLLM_GGUF_DEQUANT_ON_LOAD=1 or --gguf-dequant-on-load. When set, the GGUF loader converts Q8_0 (or other compatible formats) to torch.bfloat16 during model initialization:
# In vllm_gguf_plugin/quantization/linear.py:
if os.environ.get("VLLM_GGUF_DEQUANT_ON_LOAD") == "1":
weight = weight.to(torch.bfloat16)
bias = bias.to(torch.bfloat16) if bias is not None else None
return F.linear(x, weight, bias)
Users with enough VRAM to hold the unquantized weights (e.g. 16 GB for 8B on a 24 GB card) can immediately benefit from native cuBLAS tensor cores without having to convert GGUF files back to HuggingFace safetensors on disk.
I searched out a bug and a couple performance optimizations while doing some research. I had the option of not doing anything, or having the LLM write the text. I hope that doesn't bother anyone. The bug was confirmed in another bug post, I added context / appended to the other bug report in terms of fixing it / root cause I determined.
Alternatives
No response
Additional context
Test environment as per collect_env
==============================
System Info
==============================
OS : Ubuntu 22.04.5 LTS (x86_64) on WSL2
Kernel : Linux 6.6.87.1-microsoft-standard-WSL2
GCC version : (Ubuntu 11.4.0-1ubuntu1~22.04.3) 11.4.0
Libc version : glibc-2.35
==============================
PyTorch Info
==============================
PyTorch version : 2.13.0+cu130
Is debug build : False
CUDA used to build PyTorch : 13.0
==============================
Python Environment
==============================
Python version : 3.10.12 (64-bit runtime)
Python platform : Linux-6.6.87.1-microsoft-standard-WSL2-x86_64-with-glibc2.35
==============================
CUDA / GPU Info
==============================
Is CUDA available : True
GPU models and configuration : GPU 0: NVIDIA GeForce RTX 4090 (24 GB)
CUDA runtime / nvcc version : 13.3, V13.3.73
Nvidia driver version : 610.88 (WDDM / WSL2)
==============================
CPU Info
==============================
Architecture : x86_64
Model name : 12th Gen Intel(R) Core(TM) i9-12900KF (24 vCPUs)
==============================
Versions of relevant libraries
==============================
[pip3] vllm==0.28.0
[pip3] flashinfer-python==0.6.16.post3
[pip3] torch==2.13.0
[pip3] triton==3.7.1
[pip3] transformers==5.16.1
[pip3] numpy==2.2.6
[pip3] ninja==1.13.2
Before submitting a new issue...
- Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.
Source: vllm-project/vllm