Natural Memory NM2.1: 记忆路由器分叉、数据集缺陷修复与全轴评测证据

- 引入 MemoryRouterXL 与 v5/v6 流式多线程训练/编码管线
- 修复 prepare_memory_router_dataset 候选池重建缺陷（mega 家族 3568x 加速，输出逐字节相同）
- 修复 v5 被破坏的拒答与多跳标签（train 未知样本 319 -> 16319，multi_hop 平均正例 1.00 -> 2.00）
- 同存储预算下 V2-128 v6 逐轴 22/22 通过：Top-1 41.12% -> 94.62%，未知拒答 0.00% -> 100.00%
- 记录三条被实测推翻的显然优化（logits_to_keep=1 反而慢 55%、XL 容量未带来收益）
- 记忆手术跨架构可移植性 14/14，读写关闭时与原生模型逐位相同

tasks.py

@@ -0,0 +1,56 @@
+"""Synthetic streaming tasks for testing dynamic learning behavior."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import torch
+from torch import Tensor
+
+
+@dataclass
+class AssociativeBatch:
+    learn_chunks: list[Tensor]
+    query_input: Tensor
+    query_labels: Tensor
+    expected: Tensor
+
+
+def _rand_tokens(batch_size: int, low: int, high: int, device: torch.device) -> Tensor:
+    return torch.randint(low, high, (batch_size,), device=device)
+
+
+def sample_associative_batch(
+    *,
+    batch_size: int,
+    vocab_size: int,
+    device: torch.device,
+    overwrite: bool = False,
+) -> AssociativeBatch:
+    """Generate key-value observations followed by a separated query.
+
+    The query cannot see the learning chunks through attention. It can only
+    answer by using the returned dynamic memory state.
+    """
+
+    key_low, key_high = 4, vocab_size // 2
+    value_low, value_high = vocab_size // 2, vocab_size
+    keys = _rand_tokens(batch_size, key_low, key_high, device)
+    value = _rand_tokens(batch_size, value_low, value_high, device)
+    learn_chunks = [torch.stack((keys, value), dim=1)]
+
+    expected = value
+    if overwrite:
+        replacement = _rand_tokens(batch_size, value_low, value_high, device)
+        learn_chunks.append(torch.stack((keys, replacement), dim=1))
+        expected = replacement
+
+    query_input = torch.stack((keys, expected), dim=1)
+    query_labels = torch.full_like(query_input, -100)
+    query_labels[:, 1] = expected
+    return AssociativeBatch(
+        learn_chunks=learn_chunks,
+        query_input=query_input,
+        query_labels=query_labels,
+        expected=expected,
+    )