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 →
1.2 KiB52 linesraw
1#pragma once
2#include "../Whisper/WhisperModel.h"
3#include "../CPU/MlContext.h"
4#include "../CPU/BufferAllocator.h"
5#include "KeyValueDownloader.h"
6#include "../CPU/KvTensors.h"
7
8// This version of the hybrid context uses the new, custom-built kernels
9class HybridContext
10{
11	CpuCompute::MlContext ml;
12	CpuCompute::VirtualAllocator allocCompute, allocComputeLayer;
13	
14	class AllocSingle : public CpuCompute::iArenaAllocator
15	{
16		CpuCompute::LargeBuffer buffer;
17		size_t capacity = 0;
18		bool allocated = false;
19		// Inherited via iArenaAllocator
20		virtual void* allocate( size_t cb, size_t align ) override final;
21
22	public:
23		virtual void resetArena() override final;
24	};
25	AllocSingle allocLayerOutput;
26
27	const CpuCompute::DecoderTensors& model;
28	const Whisper::WhisperModel& whisperModel;
29	KeyValueDownloader kvCross;
30	CpuCompute::KvTensors kv;
31
32	class SetAllocatorRaii;
33
34public:
35
36	HybridContext( const Whisper::WhisperModel& wm );
37
38	HRESULT create();
39
40	HRESULT downloadKeyValues( const DirectCompute::KeyValueBuffers& source )
41	{
42		return kvCross.download( source );
43	}
44
45	struct sDecParams
46	{
47		int n_threads;
48		int M;
49	};
50
51	HRESULT decode( const int* tokens, const int n_tokens, const int n_past, const sDecParams& dp, std::vector<float>& probs_out );
52};