Function test_attn

src/test.cpp:70–126

}

void test_attn() {
  constexpr int TEST_SEQ_LEN = 4;
  constexpr int TEST_DIM = 6;
  constexpr int TEST_HEAD_DIM = 3;
  constexpr int TEST_N_HEADS = 2;
  std::shared_ptr<Config> config = std::make_shared<Config>();
  config->dim = TEST_DIM;
  config->hidden_dim = TEST_DIM;
  config->head_dim = TEST_HEAD_DIM;
  config->n_heads = TEST_N_HEADS;
  config->vocab_size = 1;
  config->max_seq_len = TEST_SEQ_LEN;
  InferenceState s(config);
  // (n_heads, head_dim) - query vectors
  std::vector<float> q{
    0., 1e4, 0., // h=0
    0., 0., 1e4 // h=1
  };
  for (size_t i = 0; i < q.size(); i++) {
    s.q()[i] = q[i];
  }
  std::vector<f16_t> kb = float_array_to_half({
    1., 0., 0., // t=0
    0., 1., 0., // t=1
    0., 0., 1., // t=2
    -1., 0., 0. // t=3
  }); // (kv_len, n_heads, head_dim) - buffer containing key vectors of the sequence for all KV heads
  std::vector<f16_t> vb = float_array_to_half({
    1., 0., 0., // t=0
    0., 1., 0., // t=1
    0., 0., 1., // t=2
    -1., 0., 0. // t=3
  }); // (kv_len, n_heads, head_dim) - buffer containing value vectors of the sequence for all KV heads

  // Multihead attention. Iterate over all heads.
  int h;
#pragma omp parallel for private(h)
  for (h = 0; h < TEST_N_HEADS; h++) {
    int kv_head_offset = h * TEST_HEAD_DIM;
    f16_t* kh = kb.data() + kv_head_offset;
    f16_t* vh = vb.data() + kv_head_offset;
    attn(s.xb(h), s.att(h), s.q(h), kh, vh, TEST_HEAD_DIM, TEST_HEAD_DIM, TEST_N_HEADS, TEST_SEQ_LEN);
  }
  // attention scores
  // h=0
  assertArrayEquals(s.att(0), {
    0., 1., 0., 0.
  }, "att(h=0)");
  // h=1
  assertArrayEquals(s.att(1), {
    0., 0., 1., 0.
  }, "att(h=1)");
  assertArrayEquals(s.xb(), {
    0., 1., 0., // h=0
    0., 0., 1. // h=1
  }, "xout");
}