diff --git a/src/operations/Attention.cc b/src/operations/Attention.cc
index 3688a44a..20f479ec 100644
--- a/src/operations/Attention.cc
+++ b/src/operations/Attention.cc
@@ -232,7 +232,7 @@ void Attention::initialize_instructions(Tile* tile, int head_idx, int num_heads)
     addr_type value_addr = get_operand_addr(_INPUT_OPERAND + 2);
     addr_type ouput_addr = get_operand_addr(_OUTPUT_OPERAND);
     assert(num_heads <= _nkvh);
-    int kv_head_idx = head_idx / _nkvh;
+    int kv_head_idx = head_idx / num_heads;
     addr_type sram_k_ofs = sram_key_base + kv_head_idx * (_dk * seq_len) * _config.precision;
     addr_type sram_v_ofs = sram_value_base + kv_head_idx * (_dk * seq_len) * _config.precision;
     std::set<addr_type> dram_kv_addrs;    // = _key[req_idx]->get_all_addrs();
@@ -356,7 +356,7 @@ void Attention::initialize_instructions(Tile* tile, Mapping mapping, int head_id
     addr_type ouput_addr = get_operand_addr(_OUTPUT_OPERAND);
     addr_type logits_addr = get_operand_addr(_OUTPUT_OPERAND) + _dmodel * _q_len * _config.precision; // logits addr for scale
     assert(num_heads <= _nkvh);
-    int kv_head_idx = head_idx / _nkvh;
+    int kv_head_idx = head_idx / num_heads;
     addr_type sram_k_ofs = sram_key_base + kv_head_idx * (_dk * seq_len) * _config.precision;
     addr_type sram_v_ofs = sram_value_base + kv_head_idx * (_dk * seq_len) * _config.precision;
     std::set<addr_type> dram_kv_addrs;    // = _key[req_idx]->get_all_addrs();