Commit Β·
badf4f1
1
Parent(s): 9437df5
Fix Unicode encoding on Windows terminals
Browse files- demo.py +1 -1
- toy_flashmemory_inference.py +2 -2
demo.py
CHANGED
|
@@ -126,7 +126,7 @@ def main():
|
|
| 126 |
idx = keep_topk[b].nonzero(as_tuple=True)[0].tolist()
|
| 127 |
print(f" row {b}: kept chunk indices = {idx}")
|
| 128 |
|
| 129 |
-
print("\n[demo] done
|
| 130 |
|
| 131 |
|
| 132 |
if __name__ == "__main__":
|
|
|
|
| 126 |
idx = keep_topk[b].nonzero(as_tuple=True)[0].tolist()
|
| 127 |
print(f" row {b}: kept chunk indices = {idx}")
|
| 128 |
|
| 129 |
+
print("\n[demo] done -- forward + scoring + selection all ran.")
|
| 130 |
|
| 131 |
|
| 132 |
if __name__ == "__main__":
|
toy_flashmemory_inference.py
CHANGED
|
@@ -236,7 +236,7 @@ def main():
|
|
| 236 |
B, N = args.batch, args.n_chunks
|
| 237 |
|
| 238 |
# ββ 1. Load retriever ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 239 |
-
print(f"FlashMemory DS-V4
|
| 240 |
print(f"[load] {args.ckpt}")
|
| 241 |
retriever = FlashMemoryRetriever.from_checkpoint(
|
| 242 |
args.ckpt, device=device, max_position=args.max_position
|
|
@@ -298,7 +298,7 @@ def main():
|
|
| 298 |
avg_keep = sum(keep_ratios) / max(len(keep_ratios), 1)
|
| 299 |
print("-" * 60)
|
| 300 |
print(f"[done] {args.steps} tokens, {cycle} cycles, "
|
| 301 |
-
f"avg keep/cycle: {avg_keep:.1%}
|
| 302 |
print(f"[note] Dropped chunks are masked to -inf in attention (= KV not recalled to GPU). "
|
| 303 |
f"Production swap engine not included in this release.")
|
| 304 |
|
|
|
|
| 236 |
B, N = args.batch, args.n_chunks
|
| 237 |
|
| 238 |
# ββ 1. Load retriever ββββββββββββββββββββββββββββββββββββββββββββββββββββββ
|
| 239 |
+
print(f"FlashMemory DS-V4 -- toy sparse-decode loop")
|
| 240 |
print(f"[load] {args.ckpt}")
|
| 241 |
retriever = FlashMemoryRetriever.from_checkpoint(
|
| 242 |
args.ckpt, device=device, max_position=args.max_position
|
|
|
|
| 298 |
avg_keep = sum(keep_ratios) / max(len(keep_ratios), 1)
|
| 299 |
print("-" * 60)
|
| 300 |
print(f"[done] {args.steps} tokens, {cycle} cycles, "
|
| 301 |
+
f"avg keep/cycle: {avg_keep:.1%} => ~{1 - avg_keep:.0%} CSA KV dropped")
|
| 302 |
print(f"[note] Dropped chunks are masked to -inf in attention (= KV not recalled to GPU). "
|
| 303 |
f"Production swap engine not included in this release.")
|
| 304 |
|