Spaces:
Running
Running
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)
- models/sbla_attention.py +8 -0
- 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 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| 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!")
|