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 →
4.0 KiB140 linesraw
1#include "stdafx.h"
2#include "HybridLoader.h"
3using namespace CpuCompute;
4using namespace ComLight;
5
6static void populateDecodeTensorsMap( CAtlMap<CStringA, Tensor*>& map, int layersDec, DecoderTensors& dec )
7{
8	dec.layers.resize( layersDec );
9
10	map[ "decoder.positional_embedding" ] = &dec.positionalEmbedding;
11	map[ "decoder.token_embedding.weight" ] = &dec.tokenEmbedding;
12	map[ "decoder.ln.weight" ] = &dec.ln.w;
13	map[ "decoder.ln.bias" ] = &dec.ln.b;
14
15	CStringA tempString;
16	auto add = [ & ]( const char* name, int i, Tensor& t )
17	{
18		tempString.Format( "decoder.blocks.%i.%s", i, name );
19		map[ tempString ] = &t;
20	};
21
22	auto add2 = [ & ]( const char* name, int i, TensorPair& tensors )
23	{
24		tempString.Format( "decoder.blocks.%i.%s.weight", i, name );
25		map[ tempString ] = &tensors.w;
26		tempString.Format( "decoder.blocks.%i.%s.bias", i, name );
27		map[ tempString ] = &tensors.b;
28	};
29
30	for( int i = 0; i < layersDec; i++ )
31	{
32		auto& gpu = dec.layers[ i ];
33		add2( "mlp_ln", i, gpu.mlpLn );
34		add2( "mlp.0", i, gpu.mlp0 );
35		add2( "mlp.2", i, gpu.mlp1 );
36		add2( "attn_ln", i, gpu.attnLn0 );
37		add2( "attn.query", i, gpu.attnQuery );
38		add( "attn.key.weight", i, gpu.attnKey );
39
40		add2( "attn.value", i, gpu.attnValue );
41		add2( "attn.out", i, gpu.attnLn1 );
42
43		add2( "cross_attn_ln", i, gpu.crossAttnLn0 );
44		add2( "cross_attn.query", i, gpu.crossAttnQuery );
45
46		// These 3 tensors are used by the encode() method, to compute cross-attention buffers
47		// Need them in VRAM even for the hybrid model
48		// add( "cross_attn.key.weight", i, gpu.cross_attn_k_w );
49		// add2( "cross_attn.value", i, gpu.cross_attn_v_w, gpu.cross_attn_v_b );
50		add2( "cross_attn.out", i, gpu.crossAttnLn1 );
51	}
52}
53
54HybridLoader::HybridLoader( DecoderTensors& m, int countLayers ) :
55	destination( m )
56{
57	populateDecodeTensorsMap( map, countLayers, destination );
58	pending.reserve( map.GetCount() );
59}
60
61HRESULT HybridLoader::setupTensor( const CStringA& name, int n_dims, int ftype, const std::array<int, 4>& ne, ComLight::iReadStream* stream, int64_t& postponedBytes )
62{
63	auto p = map.Lookup( name );
64	if( nullptr == p )
65		return S_FALSE;
66
67	Tensor& rdi = *p->m_value;
68	PendingTensor& pt = pending.emplace_back();
69
70	__m128i vec = load16( ne.data() );
71	vec = _mm_insert_epi32( vec, 1, 3 );
72	store16( &rdi.ne, vec );
73	rdi.setDenseStrides();
74
75	pt.destPointer = p->m_value;
76	CHECK( stream->getPosition( pt.streamOffset ) );
77	pt.bufferOffset = bufferBytes;
78
79	size_t cbElement;
80	if( ftype == 0 )
81	{
82		rdi.setType( eDataType::FP32 );
83		cbElement = 4;
84	}
85	else
86	{
87		rdi.setType( eDataType::FP16 );
88		cbElement = 2;
89	}
90
91	const size_t totalElts = (size_t)(uint32_t)ne[ 0 ] * (uint32_t)ne[ 1 ] * (uint32_t)ne[ 2 ];
92	if( totalElts * cbElement > UINT_MAX )
93		return DISP_E_OVERFLOW;
94
95	size_t payloadBytes = cbElement * totalElts;
96	pt.payloadBytes = payloadBytes;
97	CHECK( stream->seek( payloadBytes, eSeekOrigin::Current ) );
98	postponedBytes += (int64_t)payloadBytes;
99
100	payloadBytes = ( payloadBytes + 31 ) & ( ~( (size_t)31 ) );
101	bufferBytes += payloadBytes;
102	return S_OK;
103}
104
105HRESULT HybridLoader::completeLoad( ComLight::iReadStream* stream, iLoaderProgressSink& progressSink )
106{
107	if( pending.size() != map.GetCount() )
108	{
109		logError( u8"Not all tensors loaded from model file - expected %zu, got %zu", map.GetCount(), pending.size() );
110		return E_INVALIDARG;
111	}
112
113	LargeBuffer buffer;
114	CHECK( buffer.allocate( bufferBytes ) );
115
116	uint8_t* rdi = buffer.pointer();
117
118	for( const auto& pt : pending )
119	{
120		if( pt.payloadBytes > INT_MAX )
121			return DISP_E_OVERFLOW;
122		CHECK( stream->seek( pt.streamOffset, eSeekOrigin::Begin ) );
123
124		int written = 0;
125		CHECK( stream->read( rdi, (int)pt.payloadBytes, written ) );
126		CHECK( progressSink.gotBytes( (int64_t)pt.payloadBytes ) );
127
128		pt.destPointer->setDataPointer( rdi );
129
130		const size_t cb = ( pt.payloadBytes + 31 ) & ( ~( (size_t)31 ) );
131		rdi += cb;
132	}
133
134	CHECK( buffer.setReadOnly( bufferBytes ) );
135	destination.setMemoryBuffer( std::move( buffer ) );
136
137	constexpr double mulMb = 1.0 / ( 1 << 20 );
138	logDebug( u8"Loaded %zu decoder tensors, %g MB RAM", pending.size(), mulMb * (double)(int64_t)bufferBytes );
139	return S_OK;
140}