torch-ext/torch_binding.h

#pragma once

#include <torch/torch.h>

// 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);