llmq/src/utilities/stack.h at dev · IST-DASLab/llmq · GitHub

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
// Copyright (c) 2025, IST Austria, developed by Erik Schultheis
// SPDX-License-Identifier: Apache-2.0
//

#ifndef LLMQ_SRC_UTILITIES_STACK_H
#define LLMQ_SRC_UTILITIES_STACK_H

#include <cstddef>
#include <vector>
#include "utilities/tensor.h"

class DeviceMemoryStack {
public:
    DeviceMemoryStack() = default;
    DeviceMemoryStack(std::byte* memory, std::size_t amount, int device_id);

    std::byte* allocate(std::size_t amount, const char* name="<unnamed>");
    Tensor allocate(ETensorDType dtype, const std::vector<long>& shape, const char* name="<unnamed>");

    void free(std::byte* ptr);
    void free(Tensor& tensor);

    std::size_t unused_capacity() const;
    std::size_t bytes_used() const;
    std::size_t max_utilization() const;
    int device_id() const;

    struct sAllocRecord {
        std::byte* Pointer;
        std::size_t Amount;
        const char* Name;
    };
    using AllocationList = std::vector<sAllocRecord>;

    const AllocationList& get_high_mark() const { return mHighMark; }
    void set_high_mark(const AllocationList& list) { mHighMark = list; }

    std::vector<std::pair<std::string, long>> get_allocation_stats() const;

private:
    int mDeviceID;
    std::byte* mBackingMemory;
    std::byte* mTop;
    std::size_t mCapacity;

    void _track_max();

    AllocationList mAlloc;

    std::size_t mMaxUtilization = 0;
    std::vector<sAllocRecord> mHighMark;
};

#endif //LLMQ_SRC_UTILITIES_STACK_H