Repository: ggml-org/llama.cpp
Pull request: ggml-zdnn: fix #15414, activate FP16 and BF16 acceleration and incorrect zTensor free

fixes #15414

Not sure if my `.supports_buft` is implemented inaccurately but the weights tensor are not going through the `.set_tensor` function, and thus, will have to re-initialise the weight zTensors on-the-fly during matmul. Not ideal though.

Activates the following data types:
1. FP16
2. BF16

Fixes:
1. `LLAMA_SET_ROWS=1` causing the inference to be incorrect (see: #15414)
2. zTensor not freeing correctly and would exhaust all available memory when `llama-bench` was used with more than 1 model
3. Moved bias zTensor to `.init_tensor` for performance improvements

### Performance

| model              | size     | params | threads | test  | t/s master | t/s PR | speedup |
| ------------------ | -------- | ------ | ------- | ----- | ---------- | ------ | ------- |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 1       | pp512 | 52.14      | 51.92  | 1.00    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 1       | tg128 | 3.92       | 3.86   | 0.98    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 2       | pp512 | 92.60      | 81.92  | 0.88    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 2       | tg128 | 4.44       | 4.48   | 1.01    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 4       | pp512 | 141.14     | 144.85 | 1.03    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 4       | tg128 | 4.83       | 4.86   | 1.01    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 8       | pp512 | 216.55     | 215.82 | 1.00    |
| granite 3B all F32 | 9.44 GiB | 2.53 B | 8       | tg128 | 4.97       | 4.95   | 1.00    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 1       | pp512 | 10.42      | 51.68  | 4.96    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 1       | tg128 | 0.45       | 3.43   | 7.62    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 2       | pp512 | 19.61      | 81.78  | 4.17    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 2       | tg128 | 0.89       | 4.17   | 4.69    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 4       | pp512 | 38.99      | 138.58 | 3.55    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 4       | tg128 | 1.73       | 4.67   | 2.70    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 8       | pp512 | 74.60      | 213.83 | 2.87    |
| granite 3B F16     | 4.72 GiB | 2.53 B | 8       | tg128 | 3.17       | 4.9    | 1.55    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 1       | pp512 | 11.30      | 51.6   | 4.57    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 1       | tg128 | 0.31       | 3.08   | 9.94    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 2       | pp512 | 21.40      | 82.45  | 3.85    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 2       | tg128 | 0.61       | 3.88   | 6.36    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 4       | pp512 | 42.28      | 142.97 | 3.38    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 4       | tg128 | 1.22       | 4.41   | 3.61    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 8       | pp512 | 80.90      | 213.85 | 2.64    |
| granite 3B BF16    | 4.72 GiB | 2.53 B | 8       | tg128 | 2.40       | 4.79   | 2.00    |

> [!NOTE]
> Tests were conducted on an IBM z17 Mainframe with 40 IFLs (cores) and 128 GB Memory on a shared R&D LPAR.

### `test-backend-ops`

```
build/bin/test-backend-ops -b zDNN | grep -v "not supported"
ggml_zdnn_init: allocating
ggml_zdnn_init: found 1 device
ggml_zdnn_init: picking default device: zDNN
ggml_zdnn_init: NNPA name: zDNN
ggml_zdnn_init: NNPA_PARMBLKFORMAT_0 = true
ggml_zdnn_init: NNPA_PARMBLKFORMAT_1 = true
Testing 3 devices

Backend 1/3: zDNN
  Device description: IBM Z Neural Network Processing Assist (NNPA)
  Device memory: 0 MB (0 MB free)

  MUL_MAT(type_a=f32,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=4,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=5,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=6,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=7,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=8,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=9,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=4,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=5,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=6,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=7,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=8,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=9,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=4,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=5,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=6,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=7,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=8,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=9,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=4,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=5,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=6,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=7,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=8,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=9,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=16,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=1,k=4,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f32,type_b=f32,m=16,n=16,k=4,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=16,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=1,k=4,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f32,m=16,n=16,k=4,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f16,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f16,m=16,n=16,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f16,m=16,n=1,k=4,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=f16,type_b=f16,m=16,n=16,k=4,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=1,k=1,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
  MUL_MAT(type_a=bf16,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK
ggml_zdnn_free: deallocating
  12353/12353 tests passed
  Backend zDNN: OK
Backend 2/3: BLAS
  Skipping
Backend 3/3: CPU
  Skipping
3/3 backends passed
OK
```

Changed files:
filename: docs/build-s390x.md
patch: @@ -241,8 +241,8 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 |            | VX/VXE/VXE2 | zDNN | Spyre |
 |------------|-------------|------|-------|
 | FP32       | ✅           | ✅    | ❓     |
-| FP16       | ✅           | ❓    | ❓     |
-| BF16       | 🚫           | ❓    | ❓     |
+| FP16       | ✅           | ✅    | ❓     |
+| BF16       | 🚫           | ✅    | ❓     |
 | Q4_0       | ✅           | ❓    | ❓     |
 | Q4_1       | ✅           | ❓    | ❓     |
 | MXFP4      | 🚫           | ❓    | ❓     |
@@ -272,4 +272,4 @@ IBM VXE/VXE2 SIMD acceleration depends on the BLAS implementation. It is strongl
 -   🚫 - acceleration unavailable, will still run using scalar implementation
 -   ❓ - acceleration unknown, please contribute if you can test it yourself
 
-Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Sep 6, 2025.
+Last Updated by **Aaron Teo (aaron.teo1@ibm.com)** on Sep 7, 2025.

filename: docs/ops.md
patch: @@ -18,6 +18,7 @@ Legend:
 |                              ACC | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
 |                              ADD | ❌ | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ✅ | ❌ |
 |                             ADD1 | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
+|                           ADD_ID | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                           ARANGE | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ |
 |                           ARGMAX | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
 |                          ARGSORT | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
@@ -26,6 +27,7 @@ Legend:
 |                             CONT | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | 🟡 | 🟡 | ❌ |
 |                          CONV_2D | ❌ | ❌ | ✅ | ❌ | ❌ | ✅ | ❌ | ✅ | ❌ |
 |                       CONV_2D_DW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+|                          CONV_3D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                CONV_TRANSPOSE_1D | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
 |                CONV_TRANSPOSE_2D | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                              COS | ❌ | ✅ | ✅ | ✅ | 🟡 | ❌ | ✅ | 🟡 | ❌ |
@@ -49,9 +51,11 @@ Legend:
 |                         GET_ROWS | ❌ | 🟡 | ✅ | 🟡 | ✅ | 🟡 | 🟡 | 🟡 | ❌ |
 |                    GET_ROWS_BACK | ❌ | ❌ | 🟡 | 🟡 | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                       GROUP_NORM | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
+|               GROUP_NORM_MUL_ADD | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                      HARDSIGMOID | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | 🟡 | ❌ | ❌ |
 |                        HARDSWISH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | 🟡 | ❌ | ❌ |
 |                           IM2COL | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | ✅ | ❌ |
+|                        IM2COL_3D | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                          L2_NORM | ❌ | ❌ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
 |                       LEAKY_RELU | ❌ | ✅ | ✅ | ✅ | ✅ | ❌ | ✅ | ✅ | ❌ |
 |                              LOG | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | ❌ |
@@ -61,7 +65,9 @@ Legend:
 |                       MUL_MAT_ID | ❌ | 🟡 | ✅ | ✅ | ✅ | 🟡 | 🟡 | ✅ | ❌ |
 |                              NEG | ❌ | ✅ | ✅ | 🟡 | 🟡 | ❌ | 🟡 | ❌ | ❌ |
 |                             NORM | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ |
+|                     NORM_MUL_ADD | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                   OPT_STEP_ADAMW | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ✅ | ❌ |
+|                     OPT_STEP_SGD | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                         OUT_PROD | 🟡 | ❌ | 🟡 | 🟡 | ❌ | ❌ | 🟡 | ❌ | ❌ |
 |                              PAD | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 |                   PAD_REFLECT_1D | ❌ | ✅ | ✅ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ |
@@ -98,6 +104,7 @@ Legend:
 |                              SUM | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ✅ | ❌ |
 |                         SUM_ROWS | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 |                           SWIGLU | ❌ | ✅ | ✅ | ✅ | 🟡 | ✅ | ✅ | 🟡 | ❌ |
+|                       SWIGLU_OAI | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ |
 |                             TANH | ❌ | ✅ | ✅ | 🟡 | 🟡 | ✅ | 🟡 | 🟡 | ❌ |
 |               TIMESTEP_EMBEDDING | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ |
 |                          UPSCALE | ❌ | 🟡 | ✅ | ✅ | 🟡 | ✅ | 🟡 | ✅ | ❌ |

filename: ggml/include/ggml-zdnn.h
patch: @@ -7,8 +7,6 @@
 extern "C" {
 #endif
 
-GGML_BACKEND_API ggml_backend_t ggml_backend_zdnn_init(void);
-
 GGML_BACKEND_API ggml_backend_reg_t ggml_backend_zdnn_reg(void);
 
 #ifdef __cplusplus

filename: ggml/src/ggml-zdnn/ggml-zdnn-impl.h
patch: @@ -76,6 +76,7 @@ struct ggml_backend_zdnn_context {
 
 struct ggml_backend_zdnn_buffer {
     void * data;
+    ggml_backend_zdnn_buffer * extra;  // for bias, etc.
     size_t size;
 
     zdnn_tensor_desc pre_tfm_desc;

filename: ggml/src/ggml-zdnn/ggml-zdnn.cpp
patch: @@ -115,9 +115,7 @@ static void ggml_zdnn_mul_mat_op(ggml_backend_zdnn_context * ctx, const ggml_ten
     ggml_backend_zdnn_buffer * weights_extra = (ggml_backend_zdnn_buffer *)weights->extra;
     ggml_backend_zdnn_buffer * inputs_extra  = (ggml_backend_zdnn_buffer *)inputs->extra;
     ggml_backend_zdnn_buffer * output_extra  = (ggml_backend_zdnn_buffer *)output->extra;
-
-    zdnn_tensor_desc ptd_bias, td_bias;
-    zdnn_ztensor zt_bias;
+    ggml_backend_zdnn_buffer * bias_extra    = (ggml_backend_zdnn_buffer *)output_extra->extra;
 
     const int64_t weights_rows = ne01;
     const int64_t weights_cols = ne00;
@@ -129,13 +127,10 @@ static void ggml_zdnn_mul_mat_op(ggml_backend_zdnn_context * ctx, const ggml_ten
     const int64_t output_rows = ne1;
     const int64_t output_cols = ne0;
 
-    const int64_t bias_dim  [GGML_MAX_DIMS]  = { 1, 1, 1, output_cols };
-    ggml_zdnn_create_tensor(ptd_bias, td_bias, zt_bias, output, bias_dim, ZDNN_1D);
-
-    void * bias_data = (void *)calloc(ne0, ggml_element_size(output));
+    // TODO: Weights are somehow not going through `ggml_backend_zdnn_buffer_set_tensor` during model loading.
+    //       So we need to load the weights here. Remove this when the issue is fixed.
+    //       Problem might be residing in `ggml_backend_zdnn_device_supports_buft`.
     if (weights_extra->ztensor.is_transformed == false) ggml_zdnn_load_tensor(weights_extra->ztensor, weights->data);
-    if (inputs_extra->ztensor.is_transformed == false) ggml_zdnn_load_tensor(inputs_extra->ztensor, inputs->data);
-    ggml_zdnn_load_tensor(zt_bias, bias_data);
 
     // GGML_LOG_INFO("%s: tensor '%s' tensor dimensions: [%ld, %ld, %ld, %ld] pre_tfm_desc dimensions: [%ld, %ld, %ld, %ld]\n",
     //               __func__, weights_extra->name,
@@ -158,29 +153,21 @@ static void ggml_zdnn_mul_mat_op(ggml_backend_zdnn_context * ctx, const ggml_ten
     GGML_ASSERT(inputs_extra->pre_tfm_desc.dim1  == inputs->ne[0]  && "inputs_extra->pre_tfm_desc.dim1 must match inputs->ne[0]");
     GGML_ASSERT(inputs_extra->pre_tfm_desc.dim2  == inputs->ne[1]  && "inputs_extra->pre_tfm_desc.dim2 must match inputs->ne[1]");
 
-    ZDNN_CHECK(zdnn_matmul_transpose_op(&inputs_extra->ztensor, &weights_extra->ztensor, &zt_bias,
+    ZDNN_CHECK(zdnn_matmul_transpose_op(&inputs_extra->ztensor, &weights_extra->ztensor, &bias_extra->ztensor,
                                         false, true, MATMUL_OP_ADDITION, &output_extra->ztensor));
     // TODO: Remove in the future as we are currently DLF16 -> FP32 then in the next op, FP32 -> DLF16 again. Inefficient.
     ZDNN_CHECK(zdnn_transform_origtensor(&output_extra->ztensor, output->data));
 
-    ZDNN_CHECK(zdnn_free_ztensor_buffer(&zt_bias));
-    free(bias_data);
+    GGML_UNUSED(ctx);
+    GGML_UNUSED(weights_rows);
+    GGML_UNUSED(weights_cols);
+    GGML_UNUSED(inputs_rows);
+    GGML_UNUSED(inputs_cols);
+    GGML_UNUSED(output_rows);
+    GGML_UNUSED(output_cols);
 }
 
 static void ggml_zdnn_mul_mat_dispatch(ggml_backend_zdnn_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
-    bool use_mul_mat_vec =
-        (src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_F16)
-        && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32
-        && src0->ne[0] % 2 == 0 && src1->ne[1] == 1;
-
-    bool use_mul_mat_vec_q =
-        ggml_is_quantized(src0->type)
-        && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32;
-
-    bool use_mul_mat_q =
-        ggml_is_quantized(src0->type)
-        && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32;
-
     // debug helpers
     // GGML_LOG_INFO("%s: use_mul_mat_vec   = %d\n", __func__, use_mul_mat_vec);
     // GGML_LOG_INFO("%s: use_mul_mat_vec_q = %d\n", __func__, use_mul_mat_vec_q);
@@ -192,25 +179,7 @@ static void ggml_zdnn_mul_mat_dispatch(ggml_backend_zdnn_context * ctx, const gg
     // GGML_LOG_INFO("%s: src0 is contiguous %d, transposed %d, type = %s, name = %s\n", __func__, ggml_is_contiguous(src0), ggml_is_transposed(src0), ggml_type_name(src0->type), src0->name);
     // GGML_LOG_INFO("%s: src1 is contiguous %d, transposed %d, type = %s, name = %s\n", __func__, ggml_is_contiguous(src1), ggml_is_transposed(src1), ggml_type_name(src1->type), src1->name);
 
-    if (src0->type == GGML_TYPE_F16 && src1->type == GGML_TYPE_F16
-        && !ggml_is_transposed(src0) && !ggml_is_transposed(src1)
-        && src1->ne[2] * src1->ne[3] > 1) {
-        // general KQ + KQV multi-batch
-        GGML_LOG_INFO("%s: using zdnn_mul_mat_batched for KQ + KQV multi-batch\n", __func__);
-        // ggml_zdnn_mul_mat_batched(ctx, src0, src1, dst);
-    } else if (use_mul_mat_vec) {
-        GGML_LOG_INFO("%s: using zdnn_op_mul_mat_vec for vector multiplication\n", __func__);
-        // ggml_zdnn_op_mul_mat(ctx, src0, src1, dst, ggml_zdnn_op_mul_mat_vec, nullptr);
-    } else if (use_mul_mat_vec_q) {
-        GGML_LOG_INFO("%s: using zdnn_op_mul_mat_vec_q for quantized vector multiplication\n", __func__);
-        // ggml_zdnn_op_mul_mat(ctx, src0, src1, dst, ggml_zdnn_op_mul_mat_vec_q, ggml_zdnn_quantize_row_q8_1);
-    } else if (use_mul_mat_q) {
-        GGML_LOG_INFO("%s: using zdnn_op_mul_mat_q for quantized matrix multiplication\n", __func__);
-        // ggml_zdnn_op_mul_mat(ctx, src0, src1, dst, ggml_zdnn_op_mul_mat_q, ggml_zdnn_quantize_mmq_q8_1);
-    } else {
-        // GGML_LOG_INFO("%s: using zdnn_op_mul_mat for general matrix multiplication\n", __func__);
-        ggml_zdnn_mul_mat_op(ctx, src0, src1, dst);
-    }
+    ggml_zdnn_mul_mat_op(ctx, src0, src1, dst);
 }
 
 static bool ggml_zdnn_compute_forward(ggml_backend_zdnn_context * ctx, ggml_tensor * dst) {
@@ -253,6 +222,8 @@ static enum ggml_status ggml_zdnn_graph_compute(ggml_backend_t backend, ggml_cgr
     }
 
     return GGML_STATUS_SUCCESS;
+
+    GGML_UNUSED(ctx_dev);
 }
 
 static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_dev, const ggml_tensor * op) {
@@ -266,22 +237,30 @@ static bool ggml_zdnn_supports_op(const ggml_backend_zdnn_device_context * ctx_d
 
         case GGML_OP_MUL_MAT:
             {
-                const ggml_tensor * src0 = op->src[0];
-                const ggml_tensor * src1 = op->src[1];
+                const ggml_tensor * weights = op->src[0];
+                const ggml_tensor * inputs  = op->src[1];
 
-                const int64_t ne10 = src1->ne[0];
-                const int64_t ne0 = op->ne[0];
-                const int64_t ne1 = op->ne[1];
+                const int64_t ne10 = inputs->ne[0];
+                const int64_t ne0  = op->ne[0];
+                const int64_t ne1  = op->ne[1];
 
                 const int64_t max_batch = ctx_dev->max_size;
 
-                return ggml_is_matrix(src0) &&
-                       ggml_is_matrix(src1) &&
-                       ggml_is_contiguous(src0) &&
-                       ggml_is_contiguous(src1) &&
-                       src0->view_src == nullptr && src1->view_src == nullptr &&
-                       src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F32 &&
-                       (ne0 <= max_batch && ne1 <= max_batch && ne10 <= max_batch);
+                if (!ggml_is_matrix(weights) || !ggml_is_matrix(inputs) ||
+                    !ggml_is_contiguous(weights) || !ggml_is_contiguous(inputs) ||
+                    weights->view_src != nullptr || inputs->view_src != nullptr ||
+                    ne0 > max_batch || ne1 > max_batch || ne10 > max_batch) {
+                        return false;
+                }
+
+                switch (weights->type) {
+                    case GGML_TYPE_F32:
+                    case GGML_TYPE_F16:
+                    case GGML_TYPE_BF16:
+                        return true;
+                    default:
+                        return false;
+                }
             } break;
 
         default:
@@ -374,10 +353,9 @@ static void ggml_zdnn_free(ggml_backend_zdnn_context * ctx) {
 static void ggml_backend_zdnn_buffer_free_buffer(ggml_backend_buffer_t buffer) {
     ggml_backend_zdnn_buffer_context * ctx = (ggml_backend_zdnn_buffer_context *)buffer->context;
 
-    for (int i = 0; i < ctx->n_buffers; i++) {
-        if (ctx->buffers[i]->ztensor.buffer != NULL && ctx->buffers[i]->ztensor.is_transformed) {
-            ZDNN_CHECK(zdnn_free_ztensor_buffer(&ctx->buffers[i]->ztensor));
-        }
+    for (const auto & buf_ptr : ctx->buffers) {
+        ggml_backend_zdnn_buffer * buf = buf_ptr.get();
+        if (buf->ztensor.buffer_size > 0) ZDNN_CHECK(zdnn_free_ztensor_buffer(&buf->ztensor));
     }
 
     delete ctx;
@@ -402,18 +380,46 @@ static enum ggml_status ggml_backend_zdnn_buffer_init_tensor(ggml_backend_buffer
     std::unique_ptr<ggml_backend_zdnn_buffer> zdnn_buffer = std::make_unique<ggml_backend_zdnn_buffer>();
     zdnn_buffer->data = tensor->data;
     zdnn_buffer->size = tsize;
-    strncpy(zdnn_buffer->name, tensor->name, GGML_MAX_NAME - 1);
+    zdnn_buffer->extra = nullptr;
+    snprintf(zdnn_buffer->name, GGML_MAX_NAME, "%s", tensor->name);
 
     ggml_zdnn_init_tensor(zdnn_buffer.get(), tensor);
     tensor->extra = zdnn_buffer.get();
 
+    switch (tensor->op) {
+        case GGML_OP_MUL_MAT:
+            {
+                std::unique_ptr<ggml_backend_zdnn_buffer> zdnn_bias_buffer = std::make_unique<ggml_backend_zdnn_buffer>();
+                zdnn_bias_buffer->data = (void *)calloc(tensor->ne[0], ggml_element_size(tensor));
+                zdnn_bias_buffer->size = ggml_element_size(tensor) * tensor->ne[0];
+                snprintf(zdnn_bias_buffer->name, GGML_MAX_NAME, "%.*s (bias)",
+                         GGML_MAX_NAME - (int)sizeof(" (bias)"), tensor->name);
+
+                const int64_t bias_dim[GGML_MAX_DIMS] = { 1, 1, 1, tensor->ne[0] };
+                ggml_zdnn_create_tensor(zdnn_bias_buffer->pre_tfm_desc,
+                                        zdnn_bias_buffer->tfm_desc,
+                                        zdnn_bias_buffer->ztensor,
+                                        tensor, bias_dim, ZDNN_1D);
+
+                ggml_zdnn_load_tensor(zdnn_bias_buffer->ztensor, zdnn_bias_buffer->data);
+                zdnn_buffer->extra = zdnn_bias_buffer.get();
+
+                ctx->buffers.push_back(std::move(zdnn_bias_buffer));
+                ctx->n_buffers++;
+            } break;
+        default:
+            break;
+    }
+
     ctx->buffers.push_back(std::move(zdnn_buffer));
     ctx->n_buffers++;
 
     // GGML_LOG_INFO("%s: initialised tensor '%s' in buffer %d, size = %8.2f MiB\n",
     //               __func__, tensor->name, buffer_idx, tsize);
 
     return GGML_STATUS_SUCCESS;
+
+    GGML_UNUSED(buffer_idx);
 }
 
 static void ggml_backend_zdnn_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) {
@@ -425,6 +431,10 @@ static void ggml_backend_zdnn_buffer_memset_tensor(ggml_backend_buffer_t buffer,
 static void ggml_backend_zdnn_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) {
     memcpy((char *)tensor->data + offset, data, size);
 
+    ggml_backend_zdnn_buffer * extra = (ggml_backend_zdnn_buffer *)tensor->extra;
+    if (extra->ztensor.is_transformed) zdnn_reset_ztensor(&extra->ztensor);
+    ggml_zdnn_load_tensor(extra->ztensor, tensor->data);
+
     GGML_UNUSED(buffer);
 }
 
@@ -594,27 +604,6 @@ static ggml_guid_t ggml_backend_zdnn_guid(void) {
     return reinterpret_cast<ggml_guid_t>((void *)guid_str);
 }
 
-// TODO: remove in the future
-ggml_backend_t ggml_backend_zdnn_init(void) {
-    ggml_backend_dev_t dev = ggml_backend_reg_dev_get(ggml_backend_zdnn_reg(), 0);
-
-    ggml_backend_zdnn_context * ctx = ggml_zdnn_init(dev);
-    if (ctx == NULL) {
-        GGML_LOG_ERROR("%s: error: failed to allocate context\n", __func__);
-        return NULL;
-    }
-
-    ggml_backend_t backend = (ggml_backend_t)malloc(sizeof(ggml_backend));
-    *backend = (ggml_backend) {
-        /* .guid       = */ ggml_backend_zdnn_guid(),
-        /* .iface      = */ ggml_backend_zdnn_i,
-        /* .device     = */ dev,
-        /* .context    = */ ctx,
-    };
-
-    return backend;
-}
-
 bool ggml_backend_is_zdnn(ggml_backend_t backend) {
     return backend != NULL &&
            ggml_guid_matches(backend->guid, ggml_backend_zdnn_guid());
@@ -634,11 +623,15 @@ static const char * ggml_backend_zdnn_device_get_name(ggml_backend_dev_t dev) {
 
 static const char * ggml_backend_zdnn_device_get_description(ggml_backend_dev_t dev) {
     return "IBM Z Neural Network Processing Assist (NNPA)";
+
+    GGML_UNUSED(dev);
 }
 
 static void ggml_backend_zdnn_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) {
     *free  = 0;
     *total = 0;
+
+    GGML_UNUSED(dev);
 }
 
 static enum ggml_backend_dev_type ggml_backend_zdnn_device_get_type(ggml_backend_dev_t dev) {
@@ -656,7 +649,7 @@ static void ggml_backend_zdnn_device_get_props(ggml_backend_dev_t dev, ggml_back
         /* .async                = */ false,
         /* .host_buffer          = */ false,
         /* .buffer_from_host_ptr = */ true,
-        /* .events               = */ false,
+        /* .events               = */ false
     };
 }
 
@@ -672,7 +665,7 @@ static ggml_backend_t ggml_backend_zdnn_device_init(ggml_backend_dev_t dev, cons
         /* .guid       = */ ggml_backend_zdnn_guid(),
         /* .iface      = */ ggml_backend_zdnn_i,
         /* .device     = */ dev,
-        /* .context    = */ ctx,
+        /* .context    = */ ctx
     };
 
     return backend;
@@ -724,6 +717,8 @@ static ggml_backend_buffer_t ggml_backend_zdnn_device_buffer_from_ptr(ggml_backe
     ++ctx->n_buffers;
 
     return ggml_backend_buffer_init(ggml_backend_zdnn_buffer_from_ptr_type(), ggml_backend_zdnn_buffer_i, ctx, size);
+
+    GGML_UNUSED(max_tensor_size);
 }
 
 static bool ggml_backend_zdnn_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) {
@@ -813,7 +808,7 @@ static ggml_backend_reg_i ggml_backend_zdnn_reg_i = {
     /* .get_name         = */ ggml_backend_zdnn_reg_get_name,
     /* .get_device_count = */ ggml_backend_zdnn_reg_device_count,
     /* .get_device       = */ ggml_backend_zdnn_reg_device_get,
-    /* .get_proc_address = */ ggml_backend_zdnn_get_proc_address,
+    /* .get_proc_address = */ ggml_backend_zdnn_get_proc_address
 };
 
 static void ggml_zdnn_cleanup(void) {
@@ -831,13 +826,13 @@ ggml_backend_reg_t ggml_backend_zdnn_reg(void) {
         g_ggml_backend_zdnn_reg = (ggml_backend_reg) {
             /* .api_version = */ GGML_ZDNN_VERSION,
             /* .iface       = */ ggml_backend_zdnn_reg_i,
-            /* .context     = */ NULL,
+            /* .context     = */ NULL
         };
 
         g_ggml_backend_zdnn_device = (ggml_backend_device) {
             /* .iface       = */ ggml_backend_zdnn_device_i,
             /* .reg         = */ &g_ggml_backend_zdnn_reg,
-            /* .context     = */ &g_ggml_ctx_dev_main,
+            /* .context     = */ &g_ggml_ctx_dev_main
         };
 
         return &g_ggml_backend_zdnn_reg;