Buffer for all state used during a forward pass. Members are reused across subsequent blocks and passes. This lets us avoid allocations during inference.
| 99 | // Members are reused across subsequent blocks and passes. |
| 100 | // This lets us avoid allocations during inference. |
| 101 | struct InferenceState { |
| 102 | InferenceState(const std::shared_ptr<Config> config); |
| 103 | ~InferenceState(); |
| 104 | |
| 105 | // current activations |
| 106 | float* x() const { return _x; } |
| 107 | float* xb() const { return _xb; } |
| 108 | float* xb(int head) const { return _xb + _config->head_dim * head; } |
| 109 | // TODO: do we need xb2? |
| 110 | float* xb2() const { return _xb2; } |
| 111 | float* xb2(int head, int head_size) const { return _xb2 + head_size * head; } |
| 112 | float* hb() const { return _hb; } |
| 113 | float* hb2() const { return _hb2; } |
| 114 | float* q_a() const { return _q_a; } |
| 115 | float* q() const { return _q; } |
| 116 | float* q(int head) const { return _q + _config->head_dim * head; } |
| 117 | float* kv_a() const { return _kv_a; } |
| 118 | float* kv_b() const { return _kv_b; } |
| 119 | float* kv_b(int head) const { return _kv_b + (_config->head_dim - _config->qk_rope_head_dim + _config->v_head_dim) * head; } |
| 120 | float* ropebuf() const { return _ropebuf; } |
| 121 | float* k() const { return _k; } |
| 122 | float* k(int head) const { return _k + _config->head_dim * head; } |
| 123 | float* v() const { return _v; } |
| 124 | float* v(int head) const { return _v + _config->v_head_dim * head; } |
| 125 | float* att() const { return _att; } |
| 126 | float* att(int head) const { return _att + _config->max_seq_len * head; } |
| 127 | // MLA only |
| 128 | float* q_c() const { return _q_c; } |
| 129 | float* q_c(int head) const { return _q_c + _config->kv_lora_rank * head; } |
| 130 | float* q_rope() const { return _q_rope; } |
| 131 | float* q_rope(int head) const { return _q_rope + _config->qk_rope_head_dim * head; } |
| 132 | // mixture of experts |
| 133 | float* moe_weights() const { return _moe_weights; } |
| 134 | float* active_experts_weights() const { return _active_experts_weights; } |
| 135 | int* active_experts() const { return _active_experts; } |
| 136 | // LM head |
| 137 | float* logits() const { return _logits; } |
| 138 | int* logit_indices() const { return _logit_indices; } |
| 139 | // activation quantization buffer |
| 140 | void* aqb() const { return _aqb; } |
| 141 | |
| 142 | Device device() const { return _device; } |
| 143 | InferenceMode mode() const { return _mode; } |
| 144 | void set_mode(InferenceMode mode) { _mode = mode; } |
| 145 | |
| 146 | private: |
| 147 | std::shared_ptr<Config> _config; |
| 148 | Device _device = Device::CPU; |
| 149 | InferenceMode _mode = InferenceMode::OUTPUT_LOGITS; |
| 150 | |
| 151 | // current activations |
| 152 | float* _x = nullptr; // (dim,) - latest activation |
| 153 | float* _xb = nullptr; // (dim,) - activation inside a residual branch |
| 154 | float* _xb2 = nullptr; // (max{dim, n_heads * v_head_dim, n_heads * kv_lora_rank},) - activation inside a residual branch (second slot) |
| 155 | float* _hb = nullptr; // (max{dim, hidden_dim},) - buffer for hidden dimension in feedforward network |
| 156 | float* _hb2 = nullptr; // (hidden_dim,) - buffer for hidden dimension in feedforward network (second slot) |
| 157 | float* _q_a = nullptr; // (q_lora_rank,) - compressed (latent) query vector for latest timestamp |
| 158 | float* _q = nullptr; // (n_heads * head_dim,) - query vectors for latest timestamp |
nothing calls this directly
no outgoing calls
no test coverage detected