Fix incorrect kv_head_idx computation breaks GQA/MHA head-to-KV-head mapping - #34
Open
oyeong011 wants to merge 1 commit into
Open
Fix incorrect kv_head_idx computation breaks GQA/MHA head-to-KV-head mapping#34oyeong011 wants to merge 1 commit into
oyeong011 wants to merge 1 commit into
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
kv_head_idx = head_idx / _nkvhis mathematically wrong.num_heads(the function parameter at both call sites ofAttention::initialize_instructions) already equalsheads_per_kv(=num_attention_heads / num_kv_heads), andhead_idxis always a multiple ofheads_per_kv(the starting query-head index of that kv-head's group). The correct divisor isnum_heads, not_nkvh.Symptoms by head ratio (A = num_attention_heads, K_h = num_kv_heads)
kv_head_idx=0instead of its own slice — all heads collapse onto one KV head. No crash.kv_head_idxexceedsnum_kv_heads, out-of-range index crash (make_address: index is out of bound).Verification
After the fix,
kv_head_idxwas logged across four models and matched the expected full range[0, num_kv_heads)in every case: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.