#include <metal_simdgroup>
#include <metal_stdlib>

Classes
struct	QuantizedBlockLoader< T, BROWS, BCOLS, dst_ld, reduction_dim, tgp_size, group_size, bits >

Macros
#define	MLX_MTL_CONST static constant constexpr const

Functions
template<typename T , typename U , int values_per_thread, int bits>
U	load_vector (const device T x, thread U x_thread)

template<typename T , typename U , int values_per_thread, int bits>
U	load_vector_safe (const device T x, thread U x_thread, int N)

template<typename U , int values_per_thread, int bits>
U	qdot (const device uint8_t w, const thread U x_thread, U scale, U bias, U sum)

template<typename U , int values_per_thread, int bits>
U	qdot_safe (const device uint8_t w, const thread U x_thread, U scale, U bias, U sum, int N)

template<typename U , int values_per_thread, int bits>
void	qouter (const thread uint8_t w, U x, U scale, U bias, thread U result)

template<typename U , int N, int bits>
void	dequantize (const device uint8_t w, U scale, U bias, threadgroup U w_local)

template<typename T , int group_size, int bits, int D>
METAL_FUNC void	qmv_quad_impl (const device uint32_t w, const device T scales, const device T biases, const device T x, device T *y, constant int &in_vec_size, const constant int &out_vec_size, uint3 tid, uint quad_gid, uint quad_lid)

template<typename T , int group_size, int bits>
METAL_FUNC void	qmv_fast_impl (const device uint32_t w, const device T scales, const device T biases, const device T x, device T *y, const constant int &in_vec_size, const constant int &out_vec_size, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , int group_size, int bits>
METAL_FUNC void	qmv_impl (const device uint32_t w, const device T scales, const device T biases, const device T x, device T *y, const constant int &in_vec_size, const constant int &out_vec_size, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits>
METAL_FUNC void	qvm_impl (const device uint32_t w, const device T scales, const device T biases, const device T x, device T *y, const int in_vec_size, const int out_vec_size, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32>
METAL_FUNC void	qmm_t_impl (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, threadgroup T Xs, threadgroup T *Ws, const constant int &K, const constant int &N, const constant int &M, uint3 tid, uint lid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, const int BM = 32, const int BK = 32, const int BN = 32>
METAL_FUNC void	qmm_n_impl (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, threadgroup T Xs, threadgroup T *Ws, const constant int &K, const constant int &N, const constant int &M, uint3 tid, uint lid, uint simd_gid, uint simd_lid)

template<typename T >
METAL_FUNC void	adjust_matrix_offsets (const device T &x, const device uint32_t &w, const device T &scales, const device T &biases, device T &y, int output_stride, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid)

template<typename T >
METAL_FUNC void	adjust_matrix_offsets (const device T &x, const device uint32_t &w, const device T &scales, const device T &biases, const device uint32_t lhs_indices, const device uint32_t rhs_indices, device T &y, int output_stride, const constant int &batch_ndims, const constant int batch_shape, const constant size_t lhs_strides, const constant size_t rhs_strides, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t b_strides, uint3 tid)

template<typename T , int group_size, int bits, int D, bool batched>
void	qmv_quad (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid, uint quad_gid, uint quad_lid)

template<typename T , int group_size, int bits, bool batched>
void	qmv_fast (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, bool batched>
void	qmv (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, bool batched>
void	qvm (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, int split_k = 32>
void	qvm_split_k (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, const constant int &final_block_size, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, const bool aligned_N, const bool batched, const int BM = 32, const int BK = 32, const int BN = 32>
void	qmm_t (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &K, const constant int &N, const constant int &M, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid, uint lid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, const bool batched, const int BM = 32, const int BK = 32, const int BN = 32>
void	qmm_n (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &K, const constant int &N, const constant int &M, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t *b_strides, uint3 tid, uint lid, uint simd_gid, uint simd_lid)

template<typename T , int group_size, int bits>
void	bs_qmv_fast (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t b_strides, const constant int &batch_ndims, const constant int batch_shape, const device uint32_t lhs_indices, const device uint32_t rhs_indices, const constant size_t lhs_strides, const constant size_t rhs_strides, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , int group_size, int bits>
void	bs_qmv (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t b_strides, const constant int &batch_ndims, const constant int batch_shape, const device uint32_t lhs_indices, const device uint32_t rhs_indices, const constant size_t lhs_strides, const constant size_t rhs_strides, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , int group_size, int bits>
void	bs_qvm (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &in_vec_size, const constant int &out_vec_size, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t b_strides, const constant int &batch_ndims, const constant int batch_shape, const device uint32_t lhs_indices, const device uint32_t rhs_indices, const constant size_t lhs_strides, const constant size_t rhs_strides, uint3 tid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32>
void	bs_qmm_t (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &K, const constant int &N, const constant int &M, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t b_strides, const constant int &batch_ndims, const constant int batch_shape, const device uint32_t lhs_indices, const device uint32_t rhs_indices, const constant size_t lhs_strides, const constant size_t rhs_strides, uint3 tid, uint lid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits, const int BM = 32, const int BK = 32, const int BN = 32>
void	bs_qmm_n (const device uint32_t w, const device T scales, const device T biases, const device T x, device T y, const constant int &K, const constant int &N, const constant int &M, const constant int &x_batch_ndims, const constant int x_shape, const constant size_t x_strides, const constant int &w_batch_ndims, const constant int w_shape, const constant size_t w_strides, const constant size_t s_strides, const constant size_t b_strides, const constant int &batch_ndims, const constant int batch_shape, const device uint32_t lhs_indices, const device uint32_t rhs_indices, const constant size_t lhs_strides, const constant size_t rhs_strides, uint3 tid, uint lid, uint simd_gid, uint simd_lid)

template<typename T , const int group_size, const int bits>
void	affine_quantize (const device T w, device uint8_t out, device T scales, device T biases, uint2 index, uint2 grid_dim)

template<typename T , const int group_size, const int bits>
void	affine_quantize_scales_biases (const device T w, const device T scales, const device T biases, device uint8_t out, uint2 index, uint2 grid_dim)

template<typename T , const int group_size, const int bits>
void	affine_dequantize (const device uint8_t w, const device T scales, const device T biases, device T out, uint2 index, uint2 grid_dim)

Variables
static constant constexpr const int	SIMD_SIZE = 32

static constant constexpr const int	QUAD_SIZE = 4

Macro Definition Documentation

◆ MLX_MTL_CONST

#define MLX_MTL_CONST static constant constexpr const

Function Documentation

◆ adjust_matrix_offsets() [1/2]

template<typename T >

METAL_FUNC void adjust_matrix_offsets	(	const device T *&	x,
		const device uint32_t *&	w,
		const device T *&	scales,
		const device T *&	biases,
		const device uint32_t *	lhs_indices,
		const device uint32_t *	rhs_indices,
		device T *&	y,
		int	output_stride,
		const constant int &	batch_ndims,
		const constant int *	batch_shape,
		const constant size_t *	lhs_strides,
		const constant size_t *	rhs_strides,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid )

◆ adjust_matrix_offsets() [2/2]

template<typename T >

METAL_FUNC void adjust_matrix_offsets	(	const device T *&	x,
		const device uint32_t *&	w,
		const device T *&	scales,
		const device T *&	biases,
		device T *&	y,
		int	output_stride,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid )

◆ affine_dequantize()

template<typename T , const int group_size, const int bits>

void affine_dequantize	(	const device uint8_t *	w,
		const device T *	scales,
		const device T *	biases,
		device T *	out,
		uint2	index,
		uint2	grid_dim )

◆ affine_quantize()

template<typename T , const int group_size, const int bits>

void affine_quantize	(	const device T *	w,
		device uint8_t *	out,
		device T *	scales,
		device T *	biases,
		uint2	index,
		uint2	grid_dim )

◆ affine_quantize_scales_biases()

template<typename T , const int group_size, const int bits>

void affine_quantize_scales_biases	(	const device T *	w,
		const device T *	scales,
		const device T *	biases,
		device uint8_t *	out,
		uint2	index,
		uint2	grid_dim )

◆ bs_qmm_n()

template<typename T , const int group_size, const int bits, const int BM = 32, const int BK = 32, const int BN = 32>

void bs_qmm_n	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	K,
		const constant int &	N,
		const constant int &	M,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		const constant int &	batch_ndims,
		const constant int *	batch_shape,
		const device uint32_t *	lhs_indices,
		const device uint32_t *	rhs_indices,
		const constant size_t *	lhs_strides,
		const constant size_t *	rhs_strides,
		uint3	tid,
		uint	lid,
		uint	simd_gid,
		uint	simd_lid )

◆ bs_qmm_t()

template<typename T , const int group_size, const int bits, const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32>

void bs_qmm_t	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	K,
		const constant int &	N,
		const constant int &	M,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		const constant int &	batch_ndims,
		const constant int *	batch_shape,
		const device uint32_t *	lhs_indices,
		const device uint32_t *	rhs_indices,
		const constant size_t *	lhs_strides,
		const constant size_t *	rhs_strides,
		uint3	tid,
		uint	lid,
		uint	simd_gid,
		uint	simd_lid )

◆ bs_qmv()

template<typename T , int group_size, int bits>

void bs_qmv	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		const constant int &	batch_ndims,
		const constant int *	batch_shape,
		const device uint32_t *	lhs_indices,
		const device uint32_t *	rhs_indices,
		const constant size_t *	lhs_strides,
		const constant size_t *	rhs_strides,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ bs_qmv_fast()

template<typename T , int group_size, int bits>

void bs_qmv_fast	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		const constant int &	batch_ndims,
		const constant int *	batch_shape,
		const device uint32_t *	lhs_indices,
		const device uint32_t *	rhs_indices,
		const constant size_t *	lhs_strides,
		const constant size_t *	rhs_strides,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ bs_qvm()

template<typename T , int group_size, int bits>

void bs_qvm	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		const constant int &	batch_ndims,
		const constant int *	batch_shape,
		const device uint32_t *	lhs_indices,
		const device uint32_t *	rhs_indices,
		const constant size_t *	lhs_strides,
		const constant size_t *	rhs_strides,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ dequantize()

template<typename U , int N, int bits>

void dequantize	(	const device uint8_t *	w,
		U	scale,
		U	bias,
		threadgroup U *	w_local )

inline

◆ load_vector()

template<typename T , typename U , int values_per_thread, int bits>

U load_vector	(	const device T *	x,
		thread U *	x_thread )

inline

◆ load_vector_safe()

template<typename T , typename U , int values_per_thread, int bits>

U load_vector_safe	(	const device T *	x,
		thread U *	x_thread,
		int	N )

inline

◆ qdot()

template<typename U , int values_per_thread, int bits>

U qdot	(	const device uint8_t *	w,
		const thread U *	x_thread,
		U	scale,
		U	bias,
		U	sum )

inline

◆ qdot_safe()

template<typename U , int values_per_thread, int bits>

U qdot_safe	(	const device uint8_t *	w,
		const thread U *	x_thread,
		U	scale,
		U	bias,
		U	sum,
		int	N )

inline

◆ qmm_n()

template<typename T , const int group_size, const int bits, const bool batched, const int BM = 32, const int BK = 32, const int BN = 32>

void qmm_n	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	K,
		const constant int &	N,
		const constant int &	M,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid,
		uint	lid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmm_n_impl()

template<typename T , const int group_size, const int bits, const int BM = 32, const int BK = 32, const int BN = 32>

METAL_FUNC void qmm_n_impl	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		threadgroup T *	Xs,
		threadgroup T *	Ws,
		const constant int &	K,
		const constant int &	N,
		const constant int &	M,
		uint3	tid,
		uint	lid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmm_t()

template<typename T , const int group_size, const int bits, const bool aligned_N, const bool batched, const int BM = 32, const int BK = 32, const int BN = 32>

void qmm_t	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	K,
		const constant int &	N,
		const constant int &	M,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid,
		uint	lid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmm_t_impl()

template<typename T , const int group_size, const int bits, const bool aligned_N, const int BM = 32, const int BK = 32, const int BN = 32>

METAL_FUNC void qmm_t_impl	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		threadgroup T *	Xs,
		threadgroup T *	Ws,
		const constant int &	K,
		const constant int &	N,
		const constant int &	M,
		uint3	tid,
		uint	lid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmv()

template<typename T , const int group_size, const int bits, bool batched>

void qmv	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmv_fast()

template<typename T , int group_size, int bits, bool batched>

void qmv_fast	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmv_fast_impl()

template<typename T , int group_size, int bits>

METAL_FUNC void qmv_fast_impl	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmv_impl()

template<typename T , int group_size, int bits>

METAL_FUNC void qmv_impl	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ qmv_quad()

template<typename T , int group_size, int bits, int D, bool batched>

void qmv_quad	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid,
		uint	quad_gid,
		uint	quad_lid )

◆ qmv_quad_impl()

template<typename T , int group_size, int bits, int D>

METAL_FUNC void qmv_quad_impl	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		constant int &	in_vec_size,
		const constant int &	out_vec_size,
		uint3	tid,
		uint	quad_gid,
		uint	quad_lid )

◆ qouter()

template<typename U , int values_per_thread, int bits>

void qouter	(	const thread uint8_t *	w,
		U	x,
		U	scale,
		U	bias,
		thread U *	result )

inline

◆ qvm()

template<typename T , const int group_size, const int bits, bool batched>

void qvm	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ qvm_impl()

template<typename T , const int group_size, const int bits>

METAL_FUNC void qvm_impl	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const int	in_vec_size,
		const int	out_vec_size,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

◆ qvm_split_k()

template<typename T , const int group_size, const int bits, int split_k = 32>

void qvm_split_k	(	const device uint32_t *	w,
		const device T *	scales,
		const device T *	biases,
		const device T *	x,
		device T *	y,
		const constant int &	in_vec_size,
		const constant int &	out_vec_size,
		const constant int &	x_batch_ndims,
		const constant int *	x_shape,
		const constant size_t *	x_strides,
		const constant int &	w_batch_ndims,
		const constant int *	w_shape,
		const constant size_t *	w_strides,
		const constant size_t *	s_strides,
		const constant size_t *	b_strides,
		const constant int &	final_block_size,
		uint3	tid,
		uint	simd_gid,
		uint	simd_lid )

Variable Documentation

◆ QUAD_SIZE

constant constexpr const int QUAD_SIZE = 4

staticconstexpr

◆ SIMD_SIZE

constant constexpr const int SIMD_SIZE = 32

staticconstexpr

Classes

Macros

Functions

Variables

Macro Definition Documentation

◆ MLX_MTL_CONST

Function Documentation

◆ adjust_matrix_offsets() [1/2]

◆ adjust_matrix_offsets() [2/2]

◆ affine_dequantize()

◆ affine_quantize()

◆ affine_quantize_scales_biases()

◆ bs_qmm_n()

◆ bs_qmm_t()

◆ bs_qmv()

◆ bs_qmv_fast()

◆ bs_qvm()

◆ dequantize()

◆ load_vector()

◆ load_vector_safe()

◆ qdot()

◆ qdot_safe()

◆ qmm_n()

◆ qmm_n_impl()

◆ qmm_t()

◆ qmm_t_impl()

◆ qmv()

◆ qmv_fast()

◆ qmv_fast_impl()

◆ qmv_impl()

◆ qmv_quad()

◆ qmv_quad_impl()

◆ qouter()

◆ qvm()

◆ qvm_impl()

◆ qvm_split_k()

Variable Documentation

◆ QUAD_SIZE

◆ SIMD_SIZE