File size: 4,898 Bytes
0431133 | 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 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 | import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))
import torch
import pytest
from src.resource_allocator import PIDDualSubgradientController
from src.hardware_cost_model import HardwareRooflineAnalyzer
from src.attention import MultiHeadAttention
from src.blocks import HybridBlock
from src.quantum_backend import compute_meyer_wallach_entanglement, compute_quantum_expressibility
from src.hf_model import QTensorFormerConfig, QTensorFormerForCausalLM
def test_pid_dual_subgradient_controller():
controller = PIDDualSubgradientController(
target_latency_ms=5.0,
target_energy_uj=10.0,
kp=0.1,
ki=0.01,
kd=0.005,
)
# Initial multipliers
b0 = controller.get_budget()
assert b0.lambda_latency == 1.0
# Simulate observed latency higher than target (e.g. 8.0 ms > 5.0 ms)
# The controller should penalize latency more heavily (increase lambda_latency)
for _ in range(5):
b = controller.update(measured_latency_ms=8.0, measured_energy_uj=12.0)
assert b.lambda_latency > 1.0, f"Expected lambda_latency to increase, got {b.lambda_latency}"
assert b.lambda_energy > 0.2, f"Expected lambda_energy to increase, got {b.lambda_energy}"
# Simulate observed latency lower than target (e.g. 3.0 ms < 5.0 ms)
for _ in range(15):
b = controller.update(measured_latency_ms=3.0, measured_energy_uj=7.0)
assert b.lambda_latency < 1.5, f"Expected lambda_latency to relax, got {b.lambda_latency}"
print("[SUCCESS] test_pid_dual_subgradient_controller passed")
def test_hardware_roofline_analyzer():
analyzer_a100 = HardwareRooflineAnalyzer("gpu_a100")
# Low arithmetic intensity (Memory-Bound)
res_mem = analyzer_a100.analyze("TTLinear_r2", flops=1000, memory_traffic_bytes=5000)
assert res_mem.arithmetic_intensity == 0.2
assert res_mem.regime == "Memory-Bound"
# High arithmetic intensity (Compute-Bound)
res_comp = analyzer_a100.analyze("DenseLinear", flops=500000, memory_traffic_bytes=2000)
assert res_comp.arithmetic_intensity == 250.0
assert res_comp.regime == "Compute-Bound"
# Test Apple M2
analyzer_m2 = HardwareRooflineAnalyzer("apple_m2")
res_m2 = analyzer_m2.analyze("Attention_SDPA", flops=20000, memory_traffic_bytes=4000)
assert res_m2.hardware_name == "Apple M2 Max"
print("[SUCCESS] test_hardware_roofline_analyzer passed")
def test_grouped_query_attention():
# 8 query heads, 2 KV heads (4:1 GQA)
d_model = 64
n_heads = 8
n_kv_heads = 2
gqa = MultiHeadAttention(d_model=d_model, n_heads=n_heads, n_kv_heads=n_kv_heads)
x = torch.randn(2, 16, d_model)
out, weights = gqa(x)
assert out.shape == (2, 16, d_model)
assert not torch.isnan(out).any()
# Verify MQA (1 KV head)
mqa = MultiHeadAttention(d_model=d_model, n_heads=n_heads, n_kv_heads=1)
out_mqa, _ = mqa(x)
assert out_mqa.shape == (2, 16, d_model)
print("[SUCCESS] test_grouped_query_attention passed")
def test_early_exit_hybrid_block():
d_model = 64
vocab_size = 500
block = HybridBlock(
d_model=d_model,
n_heads=4,
vocab_size=vocab_size,
enable_early_exit=True,
early_exit_threshold=0.99, # High threshold to trigger early exit
)
x = torch.randn(2, 8, d_model)
out, stats = block(x)
assert stats.get("early_exit_triggered", False) is True
assert "early_exit_logits" in stats
assert stats["early_exit_logits"].shape == (2, 8, vocab_size)
print("[SUCCESS] test_early_exit_hybrid_block passed")
def test_quantum_metrics():
# 1. Product state: |00> -> Meyer-Wallach should be 0.0
psi_product = torch.tensor([1.0, 0.0, 0.0, 0.0])
q_prod = compute_meyer_wallach_entanglement(psi_product)
assert q_prod == 0.0, f"Expected 0.0 for product state, got {q_prod}"
# 2. Bell state: (|00> + |11>) / sqrt(2) -> Meyer-Wallach should be 1.0
psi_bell = torch.tensor([1.0 / (2**0.5), 0.0, 0.0, 1.0 / (2**0.5)])
q_bell = compute_meyer_wallach_entanglement(psi_bell)
assert abs(q_bell - 1.0) < 1e-4, f"Expected 1.0 for Bell state, got {q_bell}"
# 3. Expressibility test
def circuit_dummy(theta):
# 2 qubits
ang = theta.squeeze(0)
return torch.tensor([torch.cos(ang[0]), torch.sin(ang[0]), torch.cos(ang[1]), torch.sin(ang[1])])
expr = compute_quantum_expressibility(circuit_dummy, n_qubits=2, n_samples=50)
assert "expressibility_kl" in expr
assert expr["expressibility_kl"] >= 0.0
print("[SUCCESS] test_quantum_metrics passed")
if __name__ == "__main__":
test_pid_dual_subgradient_controller()
test_hardware_roofline_analyzer()
test_grouped_query_attention()
test_early_exit_hybrid_block()
test_quantum_metrics()
print("All Advanced Feature Tests Passed Successfully!")
|