beehive-lab / beehive-lab/GPULlama3.java

Support direct execution of K-quant GGUF weights on GPU (Q4_K, Q5_K, Q6_K)

Open
#118 0 comments 0 reactions 0 assignees View on GitHub
enhancement quantization
Dominant language
Java
Stars
274
Forks
37
Avg merge
5d 1h
Merged PRs (30d)
2

Description

### Support direct execution of K-quant GGUF weights on GPU (Q4_K, Q5_K, Q6_K)

#### Motivation

GPU support for K-quant GGUF models (Q4_K_M, Q5_K_M, Q6_K) is currently implemented by #108 through load-time dequantization to Q8_0. This enables functional GPU execution but increases memory consumption and introduces an additional model-loading step.

To preserve the memory-efficiency benefits of K-quant models, we should add support for executing directly on K-quant weights without converting them to Q8_0.

#### Current behavior
Q4_K_M, Q5_K_M, and Q6_K tensors are dequantized to Q8_0 during model loading. In ModelLoader.java we see:

```java
/**
* Dispatcher method for loading a TornadoVM-compatible tensor based on GGML type.
* Used in GPU-path.
*/
public static TornadoTensor loadTornadoTensor(GGMLTensorEntry entry) {
GGMLType ggmlType = entry.ggmlType();
int size = FloatTensor.numberOfElements(entry.shape());
return switch (ggmlType) {
case F32 -> FP32TornadoTensor.fromTornadoMemorySegment(entry.memorySegment());
case F16 -> FP16TornadoTensor.fromTornadoMemorySegment(entry.memorySegment());
case Q8_0 -> Q8_0TornadoTensor.fromTornadoMemorySegment(entry.memorySegment());
case Q4_0 -> throw new UnsupportedOperationException("Q4 format not supported yet");
case Q4_K, Q5_K, Q6_K -> dequantizeToQ8_0TornadoTensor(entry);
case Q4_0 -> throw new UnsupportedOperationException("Q4_0 format not supported for TornadoVM yet");
default -> throw new UnsupportedOperationException("Quantization format " + ggmlType);
};
}
```

GPU kernels operate on the resulting Q8_0 representation.
Model execution is correct, but memory usage is significantly higher than the original GGUF quantization format.

#### Proposed work

Implement direct TornadoVM GPU support for K-quant tensor formats:

- Q4_K_M
- Q5_K_M
- Q6_K

by implementing:

1) the tornado-based quantized tensor class in `tensor.tornado` similarly to `Q8_0TornadoTensor`.
2) in `AbstractModelLoader.effectiveGpuWeightType()` point to the corresponding quantized type (i.e. Q4_K_M) instead of the effective one
3) implement corresponding embeddings copy in `InferenceCore.forwardTornadoVM`
3) implement corresponding helper dispatch methods in `QuantizationPlannerFactory.create`
4) implement corresponding TornadoVM kernels in `tornadovm.kernels` (i.e., see `matrixVectorRowMajorOptimizedQ8_0Byte`)

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.