steel_attention.h File Reference

Go to the source code of this file.

Classes
struct	TransformScale< T >
struct	MaxOp
struct	SumOp
struct	MulOp
struct	SubOp
struct	ExpSubOp
struct	DivOp

Functions
template<typename T, int BQ, int BK, int BD, int WM, int WN, typename AccumType = float>
void	attention (const device T *Q, const device T *K, const device T *V, device T *O, const constant AttnParams *params, uint simd_lane_id, uint simd_group_id, uint3 tid, uint3 lid)

Variables
constant bool	align_Q
constant bool	align_K

Function Documentation

attention()

template<typename T, int BQ, int BK, int BD, int WM, int WN, typename AccumType = float>

void attention	(	const device T *	Q,
		const device T *	K,
		const device T *	V,
		device T *	O,
		const constant AttnParams *	params,
		uint	simd_lane_id,
		uint	simd_group_id,
		uint3	tid,
		uint3	lid )

Variable Documentation

align_K

constant bool align_K

align_Q

constant bool align_Q