Class Block

src/model.h:182–273

Transformer Block Base */

/* Transformer Block Base */
struct Block {
  Block(
    int layer_i,
    const std::shared_ptr<Config> config,
    const Tensor* rms_att_weight,
    const Tensor* rms_ffn_weight,
    const Tensor* w1,
    const Tensor* s1,
    const Tensor* w2,
    const Tensor* s2,
    const Tensor* w3,
    const Tensor* s3,
    const Tensor* shared_w1,
    const Tensor* shared_s1,
    const Tensor* shared_w2,
    const Tensor* shared_s2,
    const Tensor* shared_w3,
    const Tensor* shared_s3,
    const Tensor* moegate,
    const Tensor* moegate_bias
  );
  virtual ~Block();

  float* rms_att_weight() const { return _rms_att_weight ? static_cast<float*>(_rms_att_weight->data) : nullptr; }
  float* rms_ffn_weight() const { return _rms_ffn_weight ? static_cast<float*>(_rms_ffn_weight->data) : nullptr; }
  std::optional<QTensor> w1() const { return _w1; }
  std::optional<QTensor> w2() const { return _w2; }
  std::optional<QTensor> w3() const { return _w3; }
  std::optional<QTensor> moegate() const { return _moegate; }
  std::optional<QTensor> moegate_bias() const { return _moegate_bias; }
  std::optional<QTensor> shared_w1() const { return _shared_w1; }
  std::optional<QTensor> shared_w2() const { return _shared_w2; }
  std::optional<QTensor> shared_w3() const { return _shared_w3; }

  // Compute forward pass for this block and update the inference state accordingly.
  // PRECONDITIONS:
  // - `s.x()` contains the input to the block. Output will also go here.
  // - Block KV cache is hydrated.
  void block(
    InferenceState& s,  // inference state
    int pos,            // index of the current token in the sequence
    int kv_sink,        // number of sink tokens currently in the KV cache
    int kv_pos,         // index of the current token in the kv cache, must be in [0..kv_len) since kv cache is a ring buffer
    int kv_len          // number of tokens in the kv cache that we will attend over
  ) const;

  virtual double active_bytes(size_t pos) const;

protected:
  virtual void attention_impl(
    InferenceState& s,  // inference state
    int pos,            // index of the current token in the sequence
    int kv_sink,        // number of sink tokens currently in the KV cache
    int kv_pos,         // index of the current token in the kv cache, must be in [0..kv_len) since kv cache is a ring buffer
    int kv_len          // number of tokens in the kv cache that we will attend over
  ) const = 0;

  template <typename T>