ct2rs 0.9.21

Rust bindings for OpenNMT/CTranslate2
Documentation 
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
37
38
39
40
41
42
43
44
45

#pragma once

#include "op.h"

namespace ctranslate2 {
  namespace ops {
    class FlashAttention : public Op {
    public:
      FlashAttention(float queries_scale, dim_t sliding_window, bool is_causal = true);

      void operator()(StorageView& queries,
                      StorageView& keys,
                      StorageView& values,
                      StorageView& output,
                      StorageView* cached_keys,
                      StorageView* cached_values,
                      StorageView* attention,
                      bool return_normalized_attention,
                      StorageView* rotary_cos,
                      StorageView* rotary_sin,
                      const bool rotary_interleave,
                      StorageView* alibi,
                      dim_t offset) const;

    private:
      const float _queries_scale;
      const dim_t _sliding_window;
      const bool _is_causal;
      template <Device D>
      void compute(StorageView& queries,
                   StorageView& keys,
                   StorageView& values,
                   StorageView& output,
                   StorageView* cached_keys,
                   StorageView* cached_values,
                   StorageView* attention,
                   bool return_normalized_attention,
                   StorageView* rotary_cos,
                   StorageView* rotary_sin,
                   const bool rotary_interleave,
                   StorageView* alibi,
                   dim_t offset) const;
    };
  }
}