Skip to content

Fix incorrect kv_head_idx computation breaks GQA/MHA head-to-KV-head mapping - #34

Open
oyeong011 wants to merge 1 commit into
PSAL-POSTECH:masterfrom
oyeong011:fix/kv-head-idx-gqa-mapping
Open

Fix incorrect kv_head_idx computation breaks GQA/MHA head-to-KV-head mapping#34
oyeong011 wants to merge 1 commit into
PSAL-POSTECH:masterfrom
oyeong011:fix/kv-head-idx-gqa-mapping

Conversation

@oyeong011

Copy link
Copy Markdown

Summary

kv_head_idx = head_idx / _nkvh is mathematically wrong. num_heads (the function parameter at both call sites of Attention::initialize_instructions) already equals heads_per_kv (= num_attention_heads / num_kv_heads), and head_idx is always a multiple of heads_per_kv (the starting query-head index of that kv-head's group). The correct divisor is num_heads, not _nkvh.

Symptoms by head ratio (A = num_attention_heads, K_h = num_kv_heads)

Ratio Example Symptom
MHA (A == K_h) opt-125m (12/12), Llama-2-7B (32/32) 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 Llama-3-8B (32/8, ratio 4) Some kv heads (e.g. 4 of 8) are never accessed. Silent, no crash — indices stay in-bounds by coincidence.
GQA, A/K_h > K_h Qwen2.5-7B (28/4, ratio 7) kv_head_idx exceeds num_kv_heads, out-of-range index crash (make_address: index is out of bound).

Verification

After the fix, kv_head_idx was logged across four models and matched the expected full range [0, num_kv_heads) in every case:

  • 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 ever reached)
  • Qwen2.5-7B (28/4): 0..3 (previously crashed)

Fixes #33.

Scope

This PR contains only the two-line fix. It does not include any of the research instrumentation (dequant engine, KV-separated metrics, quantization metadata modeling) built on top of it for the study that surfaced this bug.

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 PSAL-POSTECH#33
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Incorrect kv_head_idx computation breaks GQA/MHA head-to-KV-head mapping

1 participant