gary-boon Claude Opus 4.5 commited on
Commit
a79cb83
·
1 Parent(s): decb5ab

Add safety checks for missing QKV keys

Browse files

If a hook fails silently, the layer's qkv_captures dict might exist but
be missing the 'q', 'k', or 'v' keys. This would cause a KeyError.
Now we check for each key's existence before accessing it.

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Files changed (1) hide show
  1. backend/model_service.py +16 -6
backend/model_service.py CHANGED
@@ -1859,11 +1859,16 @@ async def analyze_research_attention(request: Dict[str, Any], authenticated: boo
1859
  k_matrix = None
1860
  v_matrix = None
1861
  if layer_idx in qkv_captures:
 
1862
  # Q/K/V shape: [seq_len, n_heads, head_dim]
1863
  # Convert to float32 for numpy (bfloat16 not supported)
1864
- q_matrix = qkv_captures[layer_idx]['q'][:, head_idx, :].float().numpy().tolist()
1865
- k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].float().numpy().tolist()
1866
- v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].float().numpy().tolist()
 
 
 
 
1867
 
1868
  critical_heads.append({
1869
  "head_idx": head_idx,
@@ -2414,9 +2419,14 @@ async def analyze_research_attention_stream(request: Dict[str, Any], authenticat
2414
  k_matrix = None
2415
  v_matrix = None
2416
  if layer_idx in qkv_captures:
2417
- q_matrix = qkv_captures[layer_idx]['q'][:, head_idx, :].float().numpy().tolist()
2418
- k_matrix = qkv_captures[layer_idx]['k'][:, head_idx, :].float().numpy().tolist()
2419
- v_matrix = qkv_captures[layer_idx]['v'][:, head_idx, :].float().numpy().tolist()
 
 
 
 
 
2420
 
2421
  critical_heads.append({
2422
  "head_idx": head_idx,
 
1859
  k_matrix = None
1860
  v_matrix = None
1861
  if layer_idx in qkv_captures:
1862
+ layer_qkv = qkv_captures[layer_idx]
1863
  # Q/K/V shape: [seq_len, n_heads, head_dim]
1864
  # Convert to float32 for numpy (bfloat16 not supported)
1865
+ # Check each key exists (hooks may have failed for some)
1866
+ if 'q' in layer_qkv:
1867
+ q_matrix = layer_qkv['q'][:, head_idx, :].float().numpy().tolist()
1868
+ if 'k' in layer_qkv:
1869
+ k_matrix = layer_qkv['k'][:, head_idx, :].float().numpy().tolist()
1870
+ if 'v' in layer_qkv:
1871
+ v_matrix = layer_qkv['v'][:, head_idx, :].float().numpy().tolist()
1872
 
1873
  critical_heads.append({
1874
  "head_idx": head_idx,
 
2419
  k_matrix = None
2420
  v_matrix = None
2421
  if layer_idx in qkv_captures:
2422
+ layer_qkv = qkv_captures[layer_idx]
2423
+ # Check each key exists (hooks may have failed for some)
2424
+ if 'q' in layer_qkv:
2425
+ q_matrix = layer_qkv['q'][:, head_idx, :].float().numpy().tolist()
2426
+ if 'k' in layer_qkv:
2427
+ k_matrix = layer_qkv['k'][:, head_idx, :].float().numpy().tolist()
2428
+ if 'v' in layer_qkv:
2429
+ v_matrix = layer_qkv['v'][:, head_idx, :].float().numpy().tolist()
2430
 
2431
  critical_heads.append({
2432
  "head_idx": head_idx,