yum-archive/TaSTT-Whisper

High-performance GPGPU inference of OpenAI's Whisper automatic speech recognition (ASR) model

git clone https://git.yummers.dev/yum-archive/TaSTT-Whisper

KonstantinSource codes8c4603c

master
branchesmasterall refs →
2.8 KiB114 linesraw
1// Dispatch with [ ( neq1*neq2*neq3 + 31 ) / 32, 1, 1 ] thread groups
2#include "flashAttentionCommon.hlsli"
3Buffer<uint> lookupTable: register( t3 );
4
5void scaleTempVector( uint i, const uint length, const float multiplier )
6{
7	const uint end = i + length;
8	for( ; i < end; i++ )
9	{
10		float f = temp[ i ];
11		f *= multiplier;
12		// Rounding in this shader causes numerical errors on my GeForce 1080 Ti GPU, driver 527.56
13		// f = roundToFp16( f );
14		temp[ i ] = f;
15	}
16}
17
18inline float computeTempVectorMax( uint i, const uint length )
19{
20	// Compute per-thread maximum
21	const uint end = i + length;
22	float ax = negativeInfinity;
23	for( ; i < end; i++ )
24		ax = max( ax, temp[ i ] );
25	return ax;
26}
27
28#include "miscUtils.hlsli"
29#include "fp64Utils.hlsli"
30
31// Transform temp[ i ] = exp( temp[ i ] - tempMax ), and return the sum of these values
32inline double applySoftMax( uint i, const uint length, const float tempMax )
33{
34	// Transform the values, and compute per-thread sum
35	const uint end = i + length;
36	double sum = 0;
37	for( ; i < end; i++ )
38	{
39		float f = temp[ i ];
40		[branch]
41		if( f != negativeInfinity )
42		{
43			f -= tempMax;
44			const uint index = fp16Rounded( f );
45			const uint res16 = lookupTable[ index ];
46			f = f16tof32( res16 );
47			sum += f;
48		}
49		else
50			f = 0;
51
52		temp[ i ] = f;
53	}
54	return sum;
55}
56
57[ numthreads( 32, 1, 1 ) ]
58void main( uint3 dtid: SV_DispatchThreadID )
59{
60	const uint neq0 = q_elements[ 0 ];
61	const uint neq1 = q_elements[ 1 ];
62	const uint neq2 = q_elements[ 2 ];
63	const uint neq3 = q_elements[ 3 ];
64
65	const uint nek0 = k_elements[ 0 ];
66	const uint nek1 = k_elements[ 1 ];
67
68	const uint nev1 = v_elements[ 1 ];
69
70	const uint ne0 = res_elements[ 0 ];
71	const uint ne1 = res_elements[ 1 ];
72
73	const uint nbk0 = k_strides[ 0 ];
74	const uint nbk1 = k_strides[ 1 ];
75	const uint nbk2 = k_strides[ 2 ];
76	const uint nbk3 = k_strides[ 3 ];
77
78	const uint nbq0 = q_strides[ 0 ];
79	const uint nbq1 = q_strides[ 1 ];
80	const uint nbq2 = q_strides[ 2 ];
81	const uint nbq3 = q_strides[ 3 ];
82
83	const uint nbv0 = v_strides[ 0 ];
84	const uint nbv1 = v_strides[ 1 ];
85	const uint nbv2 = v_strides[ 2 ];
86	const uint nbv3 = v_strides[ 3 ];
87
88	const uint nb0 = res_strides[ 0 ];
89	const uint nb1 = res_strides[ 1 ];
90	const uint nb2 = res_strides[ 2 ];
91	const uint nb3 = res_strides[ 3 ];
92
93	const uint D = neq0;
94	const uint N = neq1;
95	const uint P = nek1 - N;
96	// const uint M = P + N;
97	const uint M = nek1;
98
99	const uint ir = dtid.x;
100	if( ir >= neq1 * neq2 * neq3 )
101		return;
102
103	const uint iq3 = ir / ( neq2 * neq1 );
104	const uint iq2 = ( ir - iq3 * neq2 * neq1 ) / neq1;
105	const uint iq1 = ( ir - iq3 * neq2 * neq1 - iq2 * neq1 );
106
107	const uint tempIndex = ir * tempBufferStride;
108
109	// Softmax
110	float tvm = computeTempVectorMax( tempIndex, M );
111	double sum = applySoftMax( tempIndex, M, tvm );
112
113	scaleTempVector( tempIndex, M, (float)( 1.0 / sum ) );
114}