242 lines
7.0 KiB
C
242 lines
7.0 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
/**
|
|
* @file qweight.h
|
|
* @brief Quantized weight handle and its matvec / matmul / dot kernels.
|
|
*
|
|
* Written from a behaviour specification (plans/cleanroom_r2_qweight_spec.md); no ggml or
|
|
* llama.cpp source was consulted.
|
|
*
|
|
* @status None
|
|
* @issues - Block kernels ignore a partial trailing block; Q2_0_g128 processes it.
|
|
* - Ternary is unsupported through pn_qdot (stderr + 0) but works via pn_qmv.
|
|
* - pn_qmm / pn_qmm_prefill output is undefined for types the row dequant does
|
|
* not cover and for non-block-multiple cols.
|
|
* - pn_qweight_from_gguf ignores middle dims, never writes scales/group_size.
|
|
* @todo None
|
|
*/
|
|
#ifndef PN_QWEIGHT_H
|
|
#define PN_QWEIGHT_H
|
|
|
|
#include <stddef.h>
|
|
|
|
#include "../formats/gguf.h"
|
|
#include "../formats/gguf_dequant.h"
|
|
|
|
#ifdef __cplusplus
|
|
extern "C" {
|
|
#endif
|
|
|
|
/** Largest batch pn_qmm runs through the shared-dequant path; larger B loops pn_qmv. */
|
|
#define PN_QMM_MAX_B 64
|
|
|
|
/**
|
|
* A quantized weight matrix [rows, cols]. Plain public struct: callers build it by hand.
|
|
*
|
|
* data row 0 of the quantized bytes; row r starts at data + r*row_stride
|
|
* type quant type of every row
|
|
* rows output dimension
|
|
* cols input dimension
|
|
* row_stride bytes between consecutive rows
|
|
* scales PN_QUANT_TERNARY only: per-group scales, scales[r*n_groups + g]
|
|
* group_size PN_QUANT_TERNARY only: elements per scale group
|
|
* gpu fm_buffer * holding the weight bytes on a GPU device, or NULL
|
|
* (unconditional field: the struct layout never depends on build flags)
|
|
*/
|
|
typedef struct {
|
|
const void *data;
|
|
pn_gguf_quant type;
|
|
int rows;
|
|
int cols;
|
|
size_t row_stride;
|
|
const float *scales;
|
|
int group_size;
|
|
void *gpu;
|
|
} pn_qweight;
|
|
|
|
/** Kind of dispatch reported to the dispatch observer. */
|
|
typedef enum {
|
|
PN_QWEIGHT_OP_MATVEC,
|
|
PN_QWEIGHT_OP_MATMUL,
|
|
PN_QWEIGHT_OP_DOT
|
|
} pn_qweight_op;
|
|
|
|
/** Observer called once per reported dispatch, on the dispatching thread. */
|
|
typedef void (*pn_qweight_dispatch_observer_fn)(const pn_qweight *w,
|
|
pn_qweight_op op, void *user);
|
|
|
|
struct fw_pool;
|
|
|
|
/**
|
|
* Fill a pn_qweight from a named GGUF tensor (data points into the mmap).
|
|
*
|
|
* @param qw Output handle; scales and group_size are never written
|
|
* @param gguf Open GGUF file
|
|
* @param name Tensor name (exact match)
|
|
* @return 0 on success; -1 on NULL args, a missing tensor (stderr line), or a type
|
|
* that cannot be sized (silent; data/type/gpu/rows/cols already written)
|
|
*
|
|
* Thread-safe: yes (writes only *qw)
|
|
* Reentrant: yes
|
|
* Signal-safe: no
|
|
*/
|
|
int pn_qweight_from_gguf(pn_qweight *qw, pn_gguf_file *gguf, const char *name);
|
|
|
|
/**
|
|
* out[r] = row r of w dotted with vec, for r in [0, rows).
|
|
*
|
|
* @param w Weight (NULL: no-op)
|
|
* @param vec Input [cols] (NULL: no-op)
|
|
* @param out Output [rows] (NULL: no-op)
|
|
*
|
|
* Thread-safe: no (uses the registered pool and GPU staging buffers)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmv(const pn_qweight *w, const float *vec, float *out);
|
|
|
|
/**
|
|
* out[i] = w[i] @ vec[i] for i < min(n, 8), in one parallel dispatch when possible.
|
|
* Results equal per-weight pn_qmv exactly.
|
|
*
|
|
* @param w Array of n weight pointers (NULL array: no-op)
|
|
* @param vec Array of n input vectors
|
|
* @param out Array of n output vectors
|
|
* @param n Number of weights (<= 0: no-op; > 8 is clamped to 8)
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmv_group(const pn_qweight *const *w, const float *const *vec,
|
|
float *const *out, int n);
|
|
|
|
/**
|
|
* Batched matvec: out[b*rows + r] = row r of w dotted with acts[b*cols ...].
|
|
*
|
|
* @param w Weight
|
|
* @param acts Activations [B, cols] row-major
|
|
* @param B Batch size (<= 0: no-op; > PN_QMM_MAX_B loops pn_qmv)
|
|
* @param out Output [B, rows] row-major
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmm(const pn_qweight *w, const float *acts, int B, float *out);
|
|
|
|
/**
|
|
* Prefill matmul: as pn_qmm with no B cap and panels of at least 64 rows.
|
|
*
|
|
* @param w Weight
|
|
* @param acts Activations [B, cols] row-major
|
|
* @param B Batch size (<= 0: no-op)
|
|
* @param out Output [B, rows] row-major
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmm_prefill(const pn_qweight *w, const float *acts, int B, float *out);
|
|
|
|
/**
|
|
* Install (or clear, fn = NULL) the process-wide dispatch observer.
|
|
*
|
|
* @param fn Observer, or NULL to disable
|
|
* @param user Passed back to fn on every call (stored regardless of fn)
|
|
*
|
|
* Thread-safe: no (call before the forward)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qweight_set_dispatch_observer(pn_qweight_dispatch_observer_fn fn, void *user);
|
|
|
|
/**
|
|
* Report a dispatch to the observer (if one is installed and w is non-NULL).
|
|
*
|
|
* @param w Weight being dispatched
|
|
* @param op Kind of dispatch
|
|
*
|
|
* Thread-safe: yes, if the observer is
|
|
* Reentrant: yes, if the observer is
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qweight_note_dispatch(const pn_qweight *w, pn_qweight_op op);
|
|
|
|
/**
|
|
* Register the thread pool the matvec / matmul entry points split rows over.
|
|
* Ignored in builds without USE_THREADS.
|
|
*
|
|
* @param pool Pool, or NULL to run serially
|
|
*
|
|
* Thread-safe: no (call before the forward)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmv_set_pool(struct fw_pool *pool);
|
|
|
|
/**
|
|
* Dot one quantized row with a float vector.
|
|
*
|
|
* @param row_data Row bytes (any alignment)
|
|
* @param vec Input [n]
|
|
* @param n Elements; block types use whole blocks only (Q2_0_g128 rounds up)
|
|
* @param type Quant type; unsupported types print one stderr line and return 0
|
|
* @return The dot product
|
|
*
|
|
* Thread-safe: yes
|
|
* Reentrant: yes
|
|
* Signal-safe: no
|
|
*/
|
|
float pn_qdot(const void *row_data, const float *vec, int n, pn_gguf_quant type);
|
|
|
|
/**
|
|
* @return 1 if the Metal GPU matvec is available (USE_METAL builds), else 0
|
|
*
|
|
* Thread-safe: TBD (follows fm_available)
|
|
* Reentrant: TBD
|
|
* Signal-safe: no
|
|
*/
|
|
int pn_qmv_gpu_available(void);
|
|
|
|
/**
|
|
* Upload a Q4_K weight's bytes to the GPU and attach the buffer to w->gpu.
|
|
*
|
|
* @param w Weight (NULL: returns 0)
|
|
* @return 1 if w->gpu is (already or now) set; 0 otherwise. Always 0 without USE_METAL.
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
int pn_qmv_gpu_upload(pn_qweight *w);
|
|
|
|
/**
|
|
* Free any buffer attached to w->gpu and clear it (USE_METAL builds; else no-op).
|
|
*
|
|
* @param w Weight (NULL: no-op)
|
|
*
|
|
* Thread-safe: no
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmv_gpu_free(pn_qweight *w);
|
|
|
|
/**
|
|
* Enable (non-zero) or disable GPU decode for pn_qmm on GPU-resident Q4_K weights.
|
|
* Off by default: GPU decode is slower than the CPU path; weights upload for prefill.
|
|
*
|
|
* @param on Flag, stored as-is
|
|
*
|
|
* Thread-safe: no (call before the forward)
|
|
* Reentrant: no
|
|
* Signal-safe: no
|
|
*/
|
|
void pn_qmm_set_gpu_decode(int on);
|
|
|
|
#ifdef __cplusplus
|
|
}
|
|
#endif
|
|
|
|
#endif /* PN_QWEIGHT_H */
|