rank_worker.hpp 4.69 KB
Newer Older
1
2
#pragma once

3
#include "../backends/attention_backends.hpp"
4
#include "../cache/cache.hpp"
5
#include "../config/model_config.hpp"
6
#include "../models/model_factory.hpp"
7
#include "compiler/general_compiler.hpp"
8
#include "distributed/distributed.hpp"
9
#include "rank_barrier.hpp"
10
11
12
13

#include <any>
#include <condition_variable>
#include <mutex>
14
#include <random>
15
16
17
18
19
20
21
22
23
24
25
#include <string>
#include <thread>
#include <vector>

namespace infinilm::engine {

class RankWorker {
    enum class Command {
        INIT,
        LOAD,
        RUN,
Ceng's avatar
Ceng committed
26
        RESET_CACHE,
27
        COMPILE,
28
29
30
31
        STOP
    };

public:
32
33
34
35
36
37
    struct Input {
        /// Token IDs tensor of shape `[batch, seq_len]`.
        std::optional<infinicore::Tensor> input_ids;
        /// Position IDs tensor of shape `[batch, seq_len]` or `[seq_len]`.
        std::optional<infinicore::Tensor> position_ids;
        /// Past Lengths of cached sequence for each request, of shape `[num_requests]`.
38
39
40
        std::optional<infinicore::Tensor> past_sequence_lengths;
        /// ToTal Lengths for each request sequence, of shape `[num_requests]`.
        std::optional<infinicore::Tensor> total_sequence_lengths;
41
        /// Offsets of each request in a continous-batched sequence, of shape `[num_requests + 1]`.
42
        std::optional<infinicore::Tensor> input_offsets;
43
44
        /// Cumulative total sequence lengths for each request, of shape `[num_requests + 1]`.
        std::optional<infinicore::Tensor> cu_seqlens;
45
46
47
48
49
50
51
52
53
54
55
        /// Block ids for each request `[batch, max_block_table_length]`. Used for paged cache.
        std::optional<infinicore::Tensor> block_tables;
        /// Slot ids for each token `[seq]`. Used for paged cache.
        std::optional<infinicore::Tensor> slot_mapping;

        float temperature{1};

        int top_k{50};

        float top_p{1};

56
        infinilm::InfinilmModel::Input to_model_input(infinicore::Device device) const;
57
58
59
60
61
62
    };

    struct Output {
        infinicore::Tensor output_ids;
    };

Jiacheng Huang's avatar
Jiacheng Huang committed
63
    RankWorker(const InfinilmModel::Config &model_config,
64
               const distributed::RankInfo &rank_info,
65
               const cache::CacheConfig *cache_config,
66
               RankBarrier *barrier,
67
68
               bool enable_graph_compiling,
               backends::AttentionBackend attention_backend);
69

70
71
72
73
    RankWorker(std::shared_ptr<infinilm::config::ModelConfig> model_config,
               const distributed::RankInfo &rank_info,
               const cache::CacheConfig *cache_config,
               RankBarrier *barrier,
74
75
               bool enable_graph_compiling,
               backends::AttentionBackend attention_backend);
76

77
78
79
80
    // Submit a parameter load job and wait until the load completes on the worker thread.
    void load_param(const std::string &name,
                    const infinicore::Tensor &param);

81
82
83
    // return the parameters (i.e. weights and biases).
    std::unordered_map<std::string, infinicore::nn::Parameter> state_dict();

PanZezhong's avatar
PanZezhong committed
84
    // Submit a run (forward) job.
85
    void run(const Input &args);
86

87
    // Reset the internal cache with a new configuration
PanZezhong's avatar
PanZezhong committed
88
    void reset_cache(const cache::CacheConfig *new_config);
Ceng's avatar
Ceng committed
89

90
91
92
    // Compile the model graph if enabled.
    void compile();

PanZezhong's avatar
PanZezhong committed
93
94
95
    // Wait until run job completes. The result can be retrieved with get_output().
    void wait();

96
97
98
99
    // Request worker shutdown and join the thread.
    void close();

    // Thread-safe accessor for last output produced by RUN.
100
    Output get_output();
101
102
103
104
105
106
107
108

    std::string info() const;

private:
    void thread_loop();

private:
    // Worker properties
109
110
    const InfinilmModel::Config &legacy_model_config_ = InfinilmModel::Config();
    std::shared_ptr<infinilm::config::ModelConfig> model_config_;
111
112
    distributed::RankInfo rank_info_;
    std::shared_ptr<InfinilmModel> model_;
PanZezhong's avatar
PanZezhong committed
113
    std::shared_ptr<cache::Cache> cache_;
114

115
116
117
    // Backends
    backends::AttentionBackend attention_backend_;

118
119
120
121
    // Graph Compiling
    bool enable_graph_compiling_;
    std::unique_ptr<GraphCompiler> compiler_;

122
123
124
125
126
127
128
129
130
131
132
133
    // Command for the pending job (protected by mutex_)
    Command job_cmd_;

    // State flags (protected by mutex_)
    bool has_job_ = false;     // a job is pending
    bool job_done_ = false;    // last job completed
    bool should_exit_ = false; // request to stop
    bool init_done_ = false;   // initialization finished

    // Task payloads (protected by mutex)
    std::string pending_param_name_;
    infinicore::Tensor pending_param_;
134
    Input pending_args_;
PanZezhong's avatar
PanZezhong committed
135
    std::unique_ptr<cache::CacheConfig> pending_cache_config_;
136
137

    // Output (protected by mutex)
138
    Output output_;
139
140
141
142
143

    // Thread sync
    std::thread thread_;
    std::mutex mutex_;
    std::condition_variable cv_;
144
145
146

    // Random
    std::mt19937 rng_;
147
148

    RankBarrier *barrier_;
149
150
151
};

} // namespace infinilm::engine