File size: 1,708 Bytes
3fd1a35 | 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 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | #pragma once
#include <cstddef>
#include <cstdint>
#include <memory>
#include <span>
#include <vector>
namespace ling3 {
struct GdnState {
std::vector<std::uint16_t> fp16;
std::vector<float> fp32;
bool fp32_valid = false;
};
struct GdnRunTimings {
double stage_ms = 0.0;
double npu_ms = 0.0;
double collect_ms = 0.0;
double total_ms = 0.0;
};
class GdnStep {
public:
GdnStep(
std::span<const std::byte> heads6_model,
std::span<const std::byte> heads5_model);
~GdnStep();
GdnStep(const GdnStep &) = delete;
GdnStep & operator=(const GdnStep &) = delete;
GdnStep(GdnStep &&) noexcept;
GdnStep & operator=(GdnStep &&) noexcept;
void Reset();
GdnState SaveState();
void RestoreState(const GdnState & state);
GdnRunTimings Run(
std::span<const float> query,
std::span<const float> key,
std::span<const float> value,
std::span<const float> decay,
std::span<const float> beta,
std::span<float> output);
GdnRunTimings RunBatch16(
std::span<const float> query,
std::span<const float> key,
std::span<const float> value,
std::span<const float> decay,
std::span<const float> beta,
std::span<float> output);
GdnRunTimings RunBatchCpu(
std::span<const float> query,
std::span<const float> key,
std::span<const float> value,
std::span<const float> decay,
std::span<const float> beta,
std::span<float> output);
bool has_batch16() const noexcept;
std::size_t state_bytes() const noexcept;
private:
struct Impl;
std::unique_ptr<Impl> impl_;
};
} // namespace ling3
|