From 69e6189973ec9f98f374d9b6cfd4db2e9ba89932 Mon Sep 17 00:00:00 2001 From: oyeong011 Date: Fri, 14 Aug 2026 20:55:37 +0900 Subject: [PATCH] Fix incorrect kv_head_idx computation in GQA/MHA head mapping kv_head_idx = head_idx / _nkvh is mathematically wrong. num_heads (the function parameter) already equals heads_per_kv (= num_attention_heads / num_kv_heads) at both call sites, and head_idx is always a multiple of heads_per_kv (the starting query-head index of that kv-heads group), so the correct divisor is num_heads, not _nkvh. Impact by head ratio (A = num_attention_heads, K_h = num_kv_heads): MHA (A == K_h): every head silently reads/writes kv_head_idx=0 instead of its own slice -- all heads collapse onto one KV head, no crash. GQA, A/K_h < K_h: some kv heads are never accessed (silent, no crash since indices stay in-bounds by coincidence). GQA, A/K_h > K_h: kv_head_idx exceeds nkvh, out-of-range index crash. Verified the fix restores the full expected kv_head_idx range [0, K_h) on four model configs (num_attention_heads/num_kv_heads): opt-125m 12/12 -> 0..11, Llama-2-7B 32/32 -> 0..31, Llama-3-8B 32/8 -> 0..7 (previously only 0..3 were reachable), Qwen2.5-7B 28/4 -> 0..3 (previously crashed with make_address: index is out of bound). Fixes #33 --- src/operations/Attention.cc | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) 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 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 dram_kv_addrs; // = _key[req_idx]->get_all_addrs();