Spaces:
Running
Running
File size: 1,060 Bytes
03059b7 36e31c4 03059b7 36e31c4 03059b7 36e31c4 03059b7 08f2cd2 36e31c4 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 | """Quick unit test for SBLA Attention"""
import sys
sys.path.insert(0, ".")
import torch
print("[TEST] Testing SBLA Attention...")
from models.sbla_attention import SBLAttention
sbla = SBLAttention(
hidden_size=64,
num_heads=4,
block_size=8,
latent_dim=8,
window_size=16,
mode="pure_sbla",
)
batch_size, seq_len = 2, 16
hidden_states = torch.randn(batch_size, seq_len, 64)
attention_mask = torch.ones(batch_size, seq_len)
# Create Q/K/V manually (simulating FusionAttention's role)
import torch.nn.functional as F
q_proj = torch.nn.Linear(64, 64)
k_proj = torch.nn.Linear(64, 64)
v_proj = torch.nn.Linear(64, 64)
Q = q_proj(hidden_states).view(batch_size, seq_len, 4, 16).transpose(1, 2)
K = k_proj(hidden_states).view(batch_size, seq_len, 4, 16).transpose(1, 2)
V = v_proj(hidden_states).view(batch_size, seq_len, 4, 16).transpose(1, 2)
output, cache = sbla.forward_with_qkv(Q, K, V, attention_mask)
print(f"OK: shape={output.shape}, no NaN={not torch.isnan(output).any()}, cache={cache}")
print("[PASS] SBLA Attention working!")
|