几个注意力头共用 K/V，为什么它们仍给出不同答案？

"""One-query grouped attention: explicit KV grouping and logical cache payload."""
import json
import math
from pathlib import Path


def attention(query, keys, values):
    if not keys or len(keys) != len(values):
        raise ValueError('K/V lengths must match and be nonempty.')
    if any(len(key) != len(query) for key in keys):
        raise ValueError('Q/K dimensions must match.')
    scores = [sum(q*k for q,k in zip(query,key))/math.sqrt(len(query)) for key in keys]
    exp_scores = [math.exp(score-max(scores)) for score in scores]
    weights = [score/sum(exp_scores) for score in exp_scores]
    output = [sum(w*value[d] for w,value in zip(weights,values)) for d in range(len(values[0]))]
    return {'weights':weights,'output':output}


def groups(query_heads, kv_heads):
    if query_heads <= 0 or kv_heads <= 0 or query_heads % kv_heads:
        raise ValueError('Positive query head count must be divisible by KV head count.')
    return [head//(query_heads//kv_heads) for head in range(query_heads)]


def cache_payload(batch, kv_heads, length, head_dimension, element_bytes):
    return 2 * batch * kv_heads * length * head_dimension * element_bytes


def main():
    queries = [[1,0],[0,1],[1,1],[-1,1]]
    keys = [[[1,0],[0,1],[-1,0]], [[0,1],[1,0],[0,-1]]]
    values = [[[1,0],[0,2],[3,1]], [[10,0],[0,20],[30,10]]]
    mapping = groups(4,2)
    wrong_mapping = [head%2 for head in range(4)]
    outputs = [attention(q,keys[g],values[g]) for q,g in zip(queries,mapping)]
    wrong = [attention(q,keys[g],values[g]) for q,g in zip(queries,wrong_mapping)]
    expanded_keys = [key for key in keys for _ in range(2)]
    expanded_values = [value for value in values for _ in range(2)]
    repeated = [attention(q,k,v) for q,k,v in zip(queries,expanded_keys,expanded_values)]
    assert outputs == repeated
    assert outputs[0]['output'] != outputs[1]['output']
    print(f'contiguous KV mapping={mapping}; cyclic mapping={wrong_mapping}')
    changed = []
    for head,(correct,incorrect) in enumerate(zip(outputs,wrong)):
        delta = math.dist(correct['output'],incorrect['output'])
        if delta > 1e-12:
            changed.append(head)
        print(f'head {head}: weights=' + ','.join(f'{v:.6f}' for v in correct['weights'])
              + '; output=' + ','.join(f'{v:.6f}' for v in correct['output']) + f'; cyclic error={delta:.6f}')
        assert math.isclose(sum(correct['weights']),1.0,abs_tol=1e-12)
    assert changed == [1,2]
    payload = {name:cache_payload(2,heads,128,64,2) for name,heads in [('MHA',4),('GQA',2),('MQA',1)]}
    print('logical K+V payload, B=2 L=128 D=64 FP16: '+', '.join(f'{k}={v} bytes' for k,v in payload.items()))
    print(f"GQA/MHA payload={payload['GQA']/payload['MHA']:.6f}; no speed measurement")
    assert groups(4,4) == [0,1,2,3]
    assert groups(4,1) == [0,0,0,0]
    try:
        groups(5,2)
    except ValueError:
        pass
    else:
        raise AssertionError('Non-divisible groups must fail.')
    report = {'mapping':mapping,'cyclic_mapping':wrong_mapping,'heads':outputs,'cyclic_heads':wrong,
        'heads_changed_by_wrong_mapping':changed,'logical_payload_bytes':payload,
        'scope':'Fixed Q/K/V and one query per head; no trained model, GPU kernel, full cache allocator or speed benchmark.'}
    Path('report.json').write_text(json.dumps(report,indent=2,allow_nan=False)+'\n',encoding='utf-8')
    print('checks passed: expanded-KV equivalence, distinct Q outputs, grouping counterexample, boundary counts')


if __name__ == '__main__':
    main()

"""Small actual CPU SDPA comparison for a fixed contiguous GQA head layout."""
import json
import math
import platform
import sys
from pathlib import Path
import torch
from torch.nn import functional as F
from torch.nn.attention import SDPBackend, sdpa_kernel

torch.set_num_threads(1)
DTYPE = torch.float64


def main():
    q = torch.tensor([[[[1.,0.]],[[0.,1.]],[[1.,1.]],[[-1.,1.]]]],dtype=DTYPE,device='cpu')
    k = torch.tensor([[[[1.,0.],[0.,1.],[-1.,0.]],[[0.,1.],[1.,0.],[0.,-1.]]]],dtype=DTYPE,device='cpu')
    v = torch.tensor([[[[1.,0.],[0.,2.],[3.,1.]],[[10.,0.],[0.,20.],[30.,10.]]]],dtype=DTYPE,device='cpu')
    expanded_k = k.repeat_interleave(2,dim=1)
    expanded_v = v.repeat_interleave(2,dim=1)
    manual = torch.softmax(q @ expanded_k.transpose(-2,-1)/math.sqrt(2),dim=-1) @ expanded_v
    with sdpa_kernel(backends=[SDPBackend.MATH]):
        grouped = F.scaled_dot_product_attention(q,k,v,dropout_p=0.0,is_causal=False,enable_gqa=True)
        expanded = F.scaled_dot_product_attention(q,expanded_k,expanded_v,dropout_p=0.0,is_causal=False)
        cyclic_index = torch.tensor([0,1,0,1],device='cpu')
        cyclic = F.scaled_dot_product_attention(q,k[:,cyclic_index],v[:,cyclic_index],dropout_p=0.0,is_causal=False)
    manual_error = (grouped-manual).abs().max().item()
    expanded_error = (grouped-expanded).abs().max().item()
    assert manual_error < 1e-12 and expanded_error < 1e-12
    per_head = torch.linalg.vector_norm((grouped-cyclic).squeeze(0).squeeze(1),dim=-1)
    changed = [i for i,value in enumerate(per_head.tolist()) if value > 1e-12]
    assert changed == [1,2]
    assert not torch.allclose(grouped[:,0],grouped[:,1],rtol=1e-12,atol=1e-12)
    print(f'Python {platform.python_version()}; torch {torch.__version__}; device=cpu; dtype=float64; backend=MATH; threads=1')
    print(f'Q shape={list(q.shape)}; K/V shape={list(k.shape)}')
    for head,values in enumerate(grouped.squeeze(0).squeeze(1).tolist()):
        print(f'head {head}: output='+','.join(f'{value:.6f}' for value in values)
              + f'; cyclic error={per_head[head].item():.6f}')
    print(f'GQA vs explicit repeat max error={expanded_error:.3e}; vs manual formula={manual_error:.3e}')
    print('checks passed: actual enable_gqa CPU call, explicit-repeat equivalence, distinct shared-KV heads, cyclic counterexample')
    report = {'python':sys.version,'executable':sys.executable,'torch':torch.__version__,
        'torch_path':torch.__file__,'device':'cpu','dtype':'float64','backend_requested':'SDPBackend.MATH',
        'threads':torch.get_num_threads(),'q_shape':list(q.shape),'kv_shape':list(k.shape),
        'grouped':grouped.tolist(),'expanded':expanded.tolist(),'manual':manual.tolist(),'cyclic':cyclic.tolist(),
        'manual_max_error':manual_error,'expanded_max_error':expanded_error,'cyclic_head_errors':per_head.tolist(),
        'changed_heads':changed,
        'scope':'Actual CPU PyTorch SDPA with enable_gqa on fixed Q/K/V, one query, no mask and zero dropout; no weights, generation, GPU kernel, cache allocation measurement or speed benchmark.'}
    Path('gqa-torch-report.json').write_text(json.dumps(report,indent=2,allow_nan=False)+'\n',encoding='utf-8')


if __name__ == '__main__':
    main()