File size: 1,182 Bytes
2c20074
 
 
 
0becb0a
 
 
 
 
 
 
 
2c20074
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
#pragma once

#include <torch/torch.h>

#include <vector>

// The ggml type ids this build's `mul_mat_vec` implements. Backend-specific: the CUDA port reaches
// everything upstream's `mul_mat_vec_q_switch_type` covers, while Metal's list is the type table in
// gguf_metal/ggml_dispatch.mm. A caller must ask rather than assume, or it will route a type into a
// gemv that has no kernel for it.
std::vector<int64_t> gemv_types();

// The two entry points every backend implements. Both take a GGUF weight exactly as it is stored
// in the file — `(rows, bytes_per_row)` uint8 blocks — so nothing has to be unpacked to use them.

// Blocks -> values: `(rows, bytes_per_row)` uint8 -> `(rows, cols)` of `dtype`.
at::Tensor dequantize(const at::Tensor &blocks, int64_t ggml_type, int64_t rows, int64_t cols,
                      at::ScalarType dtype);

// Fused dequantize-gemv: `blocks` is `(out_features, bytes_per_row)`, `x` is `(rows, in_features)`
// with `rows <= MAX_GEMV_ROWS`. Returns `(rows, out_features)` f32, whatever `x`'s dtype.
at::Tensor mul_mat_vec(const at::Tensor &blocks, const at::Tensor &x, int64_t ggml_type,
                       int64_t out_features);