zhan1206 commited on
Commit
08f2cd2
·
1 Parent(s): 86985bf

fix: SBLAttention forward() guard for removed projections

Browse files

- Add runtime check in SBLAttention.forward() to prevent AttributeError
- Update test_sbla.py to use forward_with_qkv() API
- All unit tests passing (test_simple_import.py, test_sbla.py)

Files changed (2) hide show
  1. models/sbla_attention.py +8 -0
  2. tests/test_sbla.py +10 -1
models/sbla_attention.py CHANGED
@@ -288,6 +288,14 @@ class SBLAttention(nn.Module):
288
  device = hidden_states.device
289
 
290
  # ========== 1. Q/K/V 投影 ==========
 
 
 
 
 
 
 
 
291
  Q = self.q_proj(hidden_states) # (batch, seq_len, num_heads * head_dim)
292
  K = self.k_proj(hidden_states) # (batch, seq_len, num_kv_heads * head_dim)
293
  V = self.v_proj(hidden_states)
 
288
  device = hidden_states.device
289
 
290
  # ========== 1. Q/K/V 投影 ==========
291
+ # S-NEW-8: Q/K/V projections moved to FusionAttention (for RoPE)
292
+ # If you see this error, use FusionAttention or forward_with_qkv() instead
293
+ if not hasattr(self, 'q_proj') or self.q_proj is None:
294
+ raise RuntimeError(
295
+ "SBLAttention.forward() requires Q/K/V projections, but they were "
296
+ "removed to avoid parameter waste. Use FusionAttention wrapper or "
297
+ "call sbla.forward_with_qkv(Q, K, V, ...) instead."
298
+ )
299
  Q = self.q_proj(hidden_states) # (batch, seq_len, num_heads * head_dim)
300
  K = self.k_proj(hidden_states) # (batch, seq_len, num_kv_heads * head_dim)
301
  V = self.v_proj(hidden_states)
tests/test_sbla.py CHANGED
@@ -19,6 +19,15 @@ batch_size, seq_len = 2, 16
19
  hidden_states = torch.randn(batch_size, seq_len, 64)
20
  attention_mask = torch.ones(batch_size, seq_len)
21
 
22
- output, cache = sbla(hidden_states=hidden_states, attention_mask=attention_mask)
 
 
 
 
 
 
 
 
 
23
  print(f"OK: shape={output.shape}, no NaN={not torch.isnan(output).any()}, cache={cache}")
24
  print("[PASS] SBLA Attention working!")
 
19
  hidden_states = torch.randn(batch_size, seq_len, 64)
20
  attention_mask = torch.ones(batch_size, seq_len)
21
 
22
+ # Create Q/K/V manually (simulating FusionAttention's role)
23
+ import torch.nn.functional as F
24
+ q_proj = torch.nn.Linear(64, 64)
25
+ k_proj = torch.nn.Linear(64, 64)
26
+ v_proj = torch.nn.Linear(64, 64)
27
+ Q = q_proj(hidden_states).view(batch_size, seq_len, 4, 16).transpose(1, 2)
28
+ K = k_proj(hidden_states).view(batch_size, seq_len, 4, 16).transpose(1, 2)
29
+ V = v_proj(hidden_states).view(batch_size, seq_len, 4, 16).transpose(1, 2)
30
+
31
+ output, cache = sbla.forward_with_qkv(Q, K, V, attention_mask)
32
  print(f"OK: shape={output.shape}, no NaN={not torch.isnan(output).any()}, cache={cache}")
33
  print("[PASS] SBLA Attention working!")