blob: a7897d3a4d423cf6fffcf9bcfd7f1e9aaa471f7e (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
|
#pragma once
#include "Tensor.h"
#include "LargeBuffer.h"
#include "../Whisper/sModelParams.h"
namespace CpuCompute
{
class KvTensors
{
uint16_t* keys = nullptr;
uint16_t* values = nullptr;
uint32_t size = 0;
CpuCompute::LargeBuffer memory;
public:
// Create these two large tensors, FP16 precision
HRESULT create( const Whisper::sModelParams& mp );
// A slice of model.memory_cross_k tensor
Tensor keysView( uint32_t len, uint32_t off ) const
{
if( len + off <= size )
return Tensor::fromData( keys + off, eDataType::FP16, len );
throw E_BOUNDS;
}
// A slice of model.memory_cross_v tensor
Tensor valuesView( uint32_t len, uint32_t off ) const
{
if( len + off <= size )
return Tensor::fromData( values + off, eDataType::FP16, len );
throw E_BOUNDS;
}
};
}
|
