Method _copy_embedding

src/infer.cpp:1217–1263

template void BlockMLA::_attention_impl<block_q3_K>(InferenceState&, int, int, int, int) const;

void Model::_copy_embedding(InferenceState& s, int token) {
  const Config& c = *config;
  switch (c.weight_quant) {
    case Quant::F32: {
      float* emb = static_cast<float*>(token_embedding_table->data);
      for (int i = 0; i < c.dim; ++i) {
        s.x()[i] = emb[token * c.dim + i];
      }
      break;
    }
    case Quant::F16: {
      f16_t* emb = static_cast<f16_t*>(token_embedding_table->data);
      for (int i = 0; i < c.dim; i+=1) {
        s.x()[i] = half_to_float(emb[token * c.dim + i]);
      }
      break;
    }
    case Quant::F8E5M2: {
      f8e5m2_t* emb = static_cast<f8e5m2_t*>(token_embedding_table->data);
      float* emb_scale = static_cast<float*>(token_embedding_scale->data);
      int* block_size = config->block_size.data();
      int scale_num_cols = (c.dim + block_size[1] - 1) / block_size[1];
      for (int i = 0; i < c.dim; i+=1) {
        int scale_i = token / block_size[0];
        int scale_j = i / block_size[1];
        float scale = emb_scale[scale_i * scale_num_cols + scale_j];
        s.x()[i] = float8e5m2_to_float(emb[token * c.dim + i]) * scale;
      }
      break;
    }
    case Quant::Q2_K: {
      block_q2_K* emb = static_cast<block_q2_K*>(token_embedding_table->data);
      int blocks_per_row = c.dim / QK_K;
      dequantize_row_q2_K(emb + token * blocks_per_row, s.x(), c.dim);
      break;
    }
    case Quant::Q3_K: {
      block_q3_K* emb = static_cast<block_q3_K*>(token_embedding_table->data);
      int blocks_per_row = c.dim / QK_K;
      dequantize_row_q3_K(emb + token * blocks_per_row, s.x(), c.dim);
      break;
    }
    default: {
      assert(false && "unsupported weight quantization");
    }
  }
}

void Model::_forward_cpu(InferenceState& s, int token, int pos, InferenceMode mode) {
  const Config& c = *config;