The calculation results of RoPE implemented by executorch and C++ are different.

Hello. I'm trying to reuse the KV cache and currently need to implement the decoupling of the KV cache and RoPE.

• model: llama2
• executorch version: 0.2.0

At present, I have implemented that the key before RoPE in the forward of executorch is output in the runtime phase every inference, so that the key without position encoding of each token can be obtained in runner.cpp.
Next, I implemented RoPE (C++) in runner.cpp, and then, the key with RoPE is rewrited to the key cache of the next inference.

// RoPE
std::vector<float> Runner::apply_rope(
  float k1,
  float k2,
  int32_t d,
  int32_t p,
  int32_t dim){
  std::vector<float> k_out(2);
  float theta = 1.0f / pow(10000.0f, (2.0f * d / dim));
  k_out[0] = k1 * cos(p*theta) - k2 * sin(p*theta);
  k_out[1] = k1 * sin(p*theta) + k2 * cos(p*theta);

  return k_out;
}

However, after the modification, during the inference, the generated text appears to be circular.

• Prompt: Once upon a time
• seq_len: 120
• temperature: 0

Once upon a time, there was a little girl who loved to read. nobody could tell her enough stories. She loved to read about the adventures of the little red hen. She loved to read about the adventures of the little red hen. She loved to read about the adventures of the little red hen. She loved to read about the adventures of the little red hen. She loved to read about the adventures of the little red hen. She loved to read about the adventures of the little red hen. She loved to read about the adventures of the little red hen

I'm surmising that there's something wrong with the RoPE calculations. Therefore, I compared the values of the key for the first layer and the first head at token 1, in the following three scenarios:

1. key cache before RoPE, which is without position encoding.
2. key cache with position encoding in forward() in executorch (I get these values in Tensor& key_cache in sdpa_with_kv_cache_out function in executorch/examples/models/llama2/custom_ops/op_sdpa.cpp)
3. key cache with position encoding obtained by calculating the C++ implemented RoPE for key cache without position encoding in the runner.cpp (My modified codes)

The valuse of 2. and 3. are different.

But, I tried to calculate the position encoding for key cache without position encoding with python according to the algorithm of RoPE in the llama2 model codes (apply_rotary_emb function in /executorch/examples/models/llama2/llama_transformer.py), and the result is the same as 3.

Thus, I'm very confused by the inconsistencies between the results of 2. and 3. And I think it's this inconsistency that leads to the failure of RoPE in my modfied codes, which in turn leads to the phenomenon of circulation in inference.

Is this because the operators or kernels invoked by executorch in the runtime phase are different from those invoked by C++/Python, resulting in different accuracy of the calculation results?

If so, how do I find out what operators or kernels invoked by executorch to calculate RoPE at runtime?

Thanks for the help!

cc @mergennachin @cccclai @helunwencser @dvorjackz

Start by comparing apply_rotary_emb in executorch/examples/models/llama2/llama_transformer.py with the RoPE path in executorch/examples/models/llama2/custom_ops/op_sdpa.cpp and the modified runner.cpp. Reproduce the token-1 key-cache values, trace the runtime operators or kernels used by sdpa_with_kv_cache_out, and identify the cause of the discrepancy so the externally rewritten cache matches the runtime result.