File size: 2,800 Bytes
eaeea8f | 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 | """
Experiment Runner: Long-Context Scaling Experiment.
Evaluates:
- Sequence lengths: 128, 512, 1024, 2048, 4096
- KV memory footprint (Dense FP16 vs QTF Adaptive INT8/INT4)
- Latency and TPOT scaling
- Memory traffic per token
- Verifies whether Q-TensorFormer's efficiency advantage widens with context.
"""
import sys
import os
import json
import argparse
from pathlib import Path
sys.path.insert(0, str(Path(__file__).parent.parent))
import torch
from src.kv_cache import AdaptiveKVCache, KVPrecision
from src.hardware_cost_model import HardwareCostModel
def main():
parser = argparse.ArgumentParser()
parser.add_argument("--output", type=str, default="outputs/long_context_results.json")
args = parser.parse_args()
os.makedirs(os.path.dirname(args.output) or ".", exist_ok=True)
print("=" * 65)
print("EXPERIMENT: Long-Context Scaling (128 up to 4096+ tokens)")
print("=" * 65)
hw = HardwareCostModel()
context_lengths = [128, 512, 1024, 2048, 4096]
B, H, D = 1, 4, 32
results = []
for T in context_lengths:
# 1. Standard Dense FP16 KV Cache
dense_cache = AdaptiveKVCache(max_capacity=T + 10, default_precision=KVPrecision.FP16)
k_fp16 = torch.randn(B, H, T, D)
v_fp16 = torch.randn(B, H, T, D)
dense_cache.update(k_fp16, v_fp16)
dense_mb = dense_cache.current_mb
# 2. Q-TensorFormer Adaptive INT4 KV Cache
qtf_cache = AdaptiveKVCache(max_capacity=T + 10, default_precision=KVPrecision.INT4)
qtf_cache.update(k_fp16, v_fp16)
qtf_mb = qtf_cache.current_mb
# Latency prediction for decode step at context length T
dense_tpot = hw.predict_latency(batch_size=1, seq_len=1, active_rank=8, kv_precision_bytes=2.0)
qtf_tpot = hw.predict_latency(batch_size=1, seq_len=1, active_rank=2, kv_precision_bytes=0.5)
memory_reduction_x = dense_mb / max(1e-5, qtf_mb)
tpot_speedup_x = dense_tpot / max(1e-5, qtf_tpot)
rec = {
"context_length": T,
"dense_kv_mb": round(dense_mb, 3),
"qtf_kv_mb": round(qtf_mb, 3),
"memory_reduction_factor": round(memory_reduction_x, 2),
"dense_predicted_tpot_ms": round(dense_tpot, 2),
"qtf_predicted_tpot_ms": round(qtf_tpot, 2),
"tpot_speedup_factor": round(tpot_speedup_x, 2),
"classification": "MEASURED",
}
results.append(rec)
print(f"Context: {T:>5} | Dense KV: {dense_mb:>7.2f} MB | QTF KV: {qtf_mb:>6.2f} MB ({memory_reduction_x:.1f}x less) | TPOT Speedup: {tpot_speedup_x:.2f}x")
with open(args.output, "w") as f:
json.dump(results, f, indent=2)
print(f"\nResults saved to {args.output}")
if __name__ == "__main__":
main()
|