Upload agentframe_quad.py with huggingface_hub

agentframe_quad.py

@@ -0,0 +1,477 @@
+"""
+AgentFrame 四层融合实现
+=======================
+认知层 × 路由层 × 存储层 × 物理层
+
+设计: agentframe_four_layer_blueprint.md
+实现: numpy (可 CPU 验证), 接口兼容 torch 迁移
+
+层:
+  L1 MetaCog      — 认知层: 任务分解/置信度/信息缺口
+  L2 LandmarkRouter — 路由层: landmark 摘要 + 分层软max + Q-Cal
+  L3 AbsorbedMLA  — 存储层: 吸收式 MLA + 分层量化 (35.6×)
+  L4 KVPager      — 物理层: 热/温/冷三级换页 + 预测性驱逐
+"""
+import numpy as np
+from dataclasses import dataclass, field
+from typing import Optional, Literal
+
+
+# ============================================================
+# L3 存储层: 吸收式 MLA + 分层量化 (核心, 自研)
+# ============================================================
+@dataclass
+class CompressedKV:
+    """压缩后的 KV 块 (存储层输出, 物理层/路由层操作对象)"""
+    chunk_id: int
+    latent: np.ndarray          # [576] 潜在向量 (kv_lora_rank 512 + k_pe 64)
+    quant_bits: int             # 16/8/4
+    heat: float = 0.5           # 热度 (物理层)
+    last_access: float = 0.0    # 最后访问时间
+    size_bytes: int = 0         # 实际字节数
+
+
+class AbsorbedMLA:
+    """
+    吸收式 MLA: 只缓存 576 维潜在向量, 不展开 KV
+    270KB → 30.4KB (8.9×) → INT8 15.2KB → INT4 7.6KB (35.6×)
+    """
+    KV_LORA_RANK = 512
+    K_ROPE = 64
+    DIM = KV_LORA_RANK + K_ROPE  # 576
+
+    def __init__(self, n_layers=27, quant_bits=4, n_ch=32):
+        self.n_layers = n_layers
+        self.quant_bits = quant_bits
+        self.n_ch = n_ch
+        self.chunks = {}  # chunk_id -> CompressedKV
+
+    def encode(self, hidden_states: np.ndarray) -> CompressedKV:
+        """压缩: [seq, d] → 潜在向量 [576]"""
+        # 模拟 kv_a_proj_with_mqa: 压缩到 512 + 64
+        if hidden_states.ndim == 2:
+            h = hidden_states.mean(axis=0)
+        else:
+            h = hidden_states
+        # 投影到 576 维 (实际是学习矩阵, 这里用确定性映射)
+        latent = np.tanh(h[:self.DIM] if len(h) >= self.DIM else np.pad(h, (0, self.DIM - len(h))))
+        return self._quantize(latent)
+
+    def _quantize(self, latent: np.ndarray) -> CompressedKV:
+        """per-channel 非对称量化"""
+        d = latent.shape[-1]
+        ch = d // self.n_ch
+        tc = latent.reshape(self.n_ch, ch)
+        if self.quant_bits < 16:
+            max_val = 2 ** self.quant_bits - 1
+            tmin = tc.min(axis=-1, keepdims=True)
+            tmax = tc.max(axis=-1, keepdims=True)
+            scale = (tmax - tmin) / max_val
+            q = np.clip(np.round((tc - tmin) / (scale + 1e-8)), 0, max_val)
+            deq = q * scale + tmin
+            bytes_per = self.quant_bits / 8
+        else:
+            deq = tc
+            bytes_per = 2
+        size = int(d * bytes_per)
+        kv = CompressedKV(
+            chunk_id=len(self.chunks),
+            latent=deq.reshape(-1),
+            quant_bits=self.quant_bits,
+            size_bytes=size,
+        )
+        self.chunks[kv.chunk_id] = kv
+        return kv
+
+    def bytes_per_token(self) -> float:
+        """每 token 每层字节 (27 层总)"""
+        per_layer = self.DIM * (self.quant_bits / 8) if self.quant_bits < 16 else self.DIM * 2
+        return per_layer * self.n_layers
+
+    def memory_bytes(self, seq_len: int) -> int:
+        """N 层总占用"""
+        return self.bytes_per_token() * seq_len * self.n_layers
+
+
+# ============================================================
+# L2 路由层: Landmark 摘要 + 分层软max + Q-Cal
+# ============================================================
+@dataclass
+class ChunkSelection:
+    """路由层输出"""
+    chunk_ids: list
+    scores: list
+    weights: list  # 分层软max 权重
+
+
+class LandmarkRouter:
+    """
+    HiLS 思路: 块摘要 = Attn(q'c, Kc, Kc) 加权和 + 熵偏置
+    分数 = (q̂·k'c)/√d + b'c, q̂ = q + W_up W_down h (Q-Cal)
+    """
+    def __init__(self, chunk_size=64, top_k=32, d_model=576, rank=16, seed=42):
+        self.chunk_size = chunk_size
+        self.top_k = top_k
+        self.d_model = d_model
+        rng = np.random.default_rng(seed)
+        # Q-Cal 低秩校准 (仅 0.6% 参数)
+        self.W_up = rng.normal(0, 0.02, (d_model, rank)) * 0.1
+        self.W_down = rng.normal(0, 0.02, (rank, d_model)) * 0.1
+        # landmark 可学习 query (每层共享一个, 简化)
+        self.landmark_q = rng.normal(0, 0.1, d_model)
+
+    def build_summary(self, chunk_latents: np.ndarray) -> tuple:
+        """
+        块摘要: k'c = Σ p_j k_j (注意力加权和)
+                b'c = -Σ p_j log p_j (熵偏置)
+        """
+        # 用 landmark query 对块内 latent 打分
+        scores = self.landmark_q @ chunk_latents.T / np.sqrt(self.d_model)
+        p = np.exp(scores - scores.max())
+        p /= p.sum()
+        k_prime = p @ chunk_latents          # 加权和 → 块摘要
+        entropy = -(p * np.log(p + 1e-9)).sum()  # 熵偏置
+        return k_prime, entropy
+
+    def q_calibrate(self, query: np.ndarray, hidden: np.ndarray) -> np.ndarray:
+        """低秩校准: q̂ = q + W_up W_down h"""
+        delta = self.W_up @ (self.W_down @ hidden)
+        return query + delta
+
+    def route(self, query: np.ndarray, hidden: np.ndarray,
+              chunk_summaries: dict) -> ChunkSelection:
+        """
+        打分 + top-K:
+        ŝ_i,c = (q̂ᵢ·k'c)/√d + b'c
+        """
+        q_hat = self.q_calibrate(query, hidden)
+        scores = {}
+        for cid, (k_prime, bias) in chunk_summaries.items():
+            s = q_hat @ k_prime / np.sqrt(self.d_model) + bias
+            scores[cid] = s
+
+        # top-K
+        ranked = sorted(scores.items(), key=lambda x: -x[1])
+        selected = ranked[:self.top_k]
+
+        # 分层软max: 分数标准化防饱和
+        raw = np.array([s for _, s in selected])
+        # z-score 标准化 + 温度
+        if raw.std() > 1e-8:
+            norm = (raw - raw.mean()) / (raw.std() + 1e-8)
+        else:
+            norm = raw - raw.mean()
+        exp_s = np.exp(norm)
+        weights = exp_s / exp_s.sum()
+
+        return ChunkSelection(
+            chunk_ids=[cid for cid, _ in selected],
+            scores=[s for _, s in selected],
+            weights=weights.tolist(),
+        )
+
+    def gqa_group_select(self, group_scores: np.ndarray) -> list:
+        """GQA 组内 max 聚合 (任一头重要即选中)"""
+        group_max = np.max(group_scores, axis=0)
+        return np.argsort(group_max)[-self.top_k:].tolist()
+
+
+# ============================================================
+# L4 物理层: KV 换页 + 预测性驱逐
+# ============================================================
+class KVPager:
+    """
+    三级存储: 显存(热) / 内存(温) / 磁盘(冷, mmap)
+    预测性驱逐: 注意力衰减 + 访问频率 + 语义相关
+    """
+    def __init__(self, vram_limit_mb=10240, ram_limit_mb=32768):
+        self.vram_limit = vram_limit_mb * 1024 * 1024
+        self.ram_limit = ram_limit_mb * 1024 * 1024
+        self.vram = {}   # chunk_id -> CompressedKV
+        self.ram = {}
+        self.disk = {}   # 模拟 mmap
+        self.vram_used = 0
+        self.ram_used = 0
+        self.access_log = []  # (chunk_id, time, attention_score)
+
+    def place(self, kv: CompressedKV, now: float = 0.0):
+        """放置: 按热度分层"""
+        if kv.heat >= 0.8:
+            if self.vram_used + kv.size_bytes <= self.vram_limit:
+                self.vram[kv.chunk_id] = kv
+                self.vram_used += kv.size_bytes
+                return
+            kv.heat = 0.6
+        if kv.heat >= 0.4:
+            if self.ram_used + kv.size_bytes <= self.ram_limit:
+                self.ram[kv.chunk_id] = kv
+                self.ram_used += kv.size_bytes
+                return
+        self.disk[kv.chunk_id] = kv
+
+    def access(self, chunk_id: int, attention_score: float, now: float):
+        """访问: 更新热度 + 日志"""
+        self.access_log.append((chunk_id, now, attention_score))
+        kv = self.vram.get(chunk_id) or self.ram.get(chunk_id) or self.disk.get(chunk_id)
+        if kv:
+            kv.heat = min(1.0, kv.heat + 0.15)
+            kv.last_access = now
+            # 磁盘 → 提升
+            if chunk_id in self.disk and kv.heat >= 0.5:
+                self.disk.pop(chunk_id)
+                kv.heat = 0.6
+                self.place(kv, now)
+
+    def eviction_score(self, kv: CompressedKV, now: float) -> float:
+        """预测性驱逐分数 (越高越该驱逐)"""
+        # 注意力衰减: 最近分数越低越该走
+        logs = [a for cid, t, a in self.access_log if cid == kv.chunk_id]
+        if not logs:
+            return 0.5
+        attn_decay = 1.0 - np.mean(logs[-5:])  # 最近注意力均值
+        # 访问频率: 越频繁越留
+        freq = len(logs) / max(now, 1.0)
+        access_factor = min(freq * 5, 1.0)
+        # 时间衰减: 越久没访问越该走
+        time_decay = min((now - kv.last_access) / 100.0, 1.0)
+        return 0.5 * attn_decay + 0.3 * (1 - access_factor) + 0.2 * time_decay
+
+    def evict(self, now: float, target_layer: str = "vram"):
+        """驱逐: 选分数最高的换出"""
+        pool = self.vram if target_layer == "vram" else self.ram
+        if not pool:
+            return None
+        victim_id = max(pool, key=lambda cid: self.eviction_score(pool[cid], now))
+        victim = pool.pop(victim_id)
+        if target_layer == "vram":
+            self.vram_used -= victim.size_bytes
+            victim.heat = 0.5
+            self.place(victim, now)  # 降级到 RAM/DISK
+        else:
+            self.ram_used -= victim.size_bytes
+            victim.heat = 0.2
+            self.disk[victim_id] = victim
+        return victim_id
+
+    def stats(self) -> dict:
+        return {
+            "vram": f"{len(self.vram)} 块 / {self.vram_used/1024/1024:.1f}MB",
+            "ram": f"{len(self.ram)} 块 / {self.ram_used/1024/1024:.1f}MB",
+            "disk": f"{len(self.disk)} 块",
+        }
+
+
+# ============================================================
+# L1 认知层: 元认知控制器
+# ============================================================
+@dataclass
+class SubTask:
+    """子任务"""
+    id: int
+    name: str
+    info_needs: list          # 需要的信息类型
+    confidence: float = 0.0   # 完成置信度
+
+
+@dataclass
+class RetrievalDirective:
+    """认知层 → 物理层指令"""
+    required_chunks: list
+    priority: Literal["hot", "warm", "cold"]
+    reason: str
+    confidence: float
+
+
+class MetaCog:
+    """
+    认知层: 任务分解 + 置信度追踪 + 信息缺口分析 + 记忆提升
+    """
+    def __init__(self, confidence_threshold=0.7):
+        self.confidence_threshold = confidence_threshold
+        self.task_stack = []
+        self.long_term = {}      # 长期记忆 (key -> 摘要)
+        self.episodic = []       # 情景记忆
+        self._task_id = 0
+
+    def decompose(self, task: str) -> list[SubTask]:
+        """任务分解: 规则+关键词 (生产可换 LLM)"""
+        # 简单规则: 按句子/分号切分
+        parts = [p.strip() for p in task.replace('；', ';').replace('。', ';').split(';') if p.strip()]
+        subs = []
+        for p in parts:
+            self._task_id += 1
+            # 信息需求推断: 含"分析/对比/总结"等需要历史信息
+            needs = []
+            for kw, info in [("对比", "comparison"), ("分析", "method"),
+                             ("历史", "history"), ("数据", "data"),
+                             ("代码", "code"), ("总结", "result"),
+                             ("稀疏", "sparse")]:
+                if kw in p:
+                    needs.append(info)
+            if not needs:
+                needs = ["data"]  # 默认需要数据
+            subs.append(SubTask(id=self._task_id, name=p, info_needs=needs))
+        self.task_stack = subs
+        return subs
+
+    def track_confidence(self, response_likelihood: float) -> float:
+        """置信度追踪 (0-1)"""
+        return response_likelihood
+
+    def analyze_gap(self, subtask: SubTask, context_keys: set) -> list:
+        """信息缺口: 需要但上下文没有的信息"""
+        gaps = [need for need in subtask.info_needs if need not in context_keys]
+        return gaps
+
+    def build_directive(self, subtask: SubTask, chunk_map: dict) -> RetrievalDirective:
+        """生成检索指令 (认知层 → 路由/物理层)"""
+        needed = []
+        for need in subtask.info_needs:
+            for cid, meta in chunk_map.items():
+                if need in meta.get("tags", []):
+                    needed.append(cid)
+        return RetrievalDirective(
+            required_chunks=needed,
+            priority="hot" if len(needed) <= 3 else "warm",
+            reason=f"子任务 '{subtask.name}' 需要 {subtask.info_needs}",
+            confidence=subtask.confidence,
+        )
+
+    def promote(self, chunk_id: int, importance: float):
+        """记忆提升: 高频+重要 → 长期"""
+        if importance > 0.8:
+            self.long_term[chunk_id] = {"importance": importance, "promoted_at": len(self.episodic)}
+        self.episodic.append({"chunk": chunk_id, "importance": importance})
+
+
+# ============================================================
+# 四层编排 (端到端)
+# ============================================================
+class QuadLayerAgent:
+    """四层融合 Agent"""
+    def __init__(self, n_layers=27, quant_bits=4, top_k=32, seed=42):
+        self.metacog = MetaCog()
+        self.router = LandmarkRouter(top_k=top_k, seed=seed)
+        self.store = AbsorbedMLA(n_layers=n_layers, quant_bits=quant_bits)
+        self.pager = KVPager()
+        self.chunk_meta = {}  # chunk_id -> {tags, ...}
+        self.summaries = {}   # chunk_id -> (k_prime, bias)
+        self.now = 0.0
+
+    def ingest(self, text_chunks: list[dict]):
+        """
+        摄取知识: text_chunks = [{"text": "...", "tags": ["data"]}, ...]
+        每块 → 压缩为 CompressedKV → 建摘要 → 放置
+        """
+        for i, chunk in enumerate(text_chunks):
+            # 模拟 hidden states
+            h = np.random.default_rng(i).normal(0, 1, self.store.DIM)
+            kv = self.store.encode(h)
+            kv.heat = 0.9  # 新知识 = 热
+            self.pager.place(kv, self.now)
+            # 建 landmark 摘要
+            k_prime, bias = self.router.build_summary(kv.latent.reshape(1, -1))
+            self.summaries[kv.chunk_id] = (k_prime, bias)
+            self.chunk_meta[kv.chunk_id] = {"tags": chunk.get("tags", []), "text": chunk.get("text", "")}
+
+    def run_task(self, task: str, query: np.ndarray):
+        """执行任务: 认知 → 路由 → 存储 → 输出"""
+        print(f"\n{'='*56}")
+        print(f"任务: {task}")
+        print(f"{'='*56}")
+
+        # L1 认知: 分解
+        subtasks = self.metacog.decompose(task)
+        print(f"[认知层] 分解为 {len(subtasks)} 个子任务")
+
+        # L1 认知: 信息缺口 → 检索指令
+        directive = self.metacog.build_directive(subtasks[0], self.chunk_meta)
+        print(f"[认知层] 检索指令: 需要 {directive.required_chunks} ({directive.reason})")
+
+        # L4 物理: 预取
+        for cid in directive.required_chunks:
+            if cid in self.pager.disk:
+                kv = self.pager.disk.pop(cid)
+                kv.heat = 0.6
+                self.pager.place(kv, self.now)
+        print(f"[物理层] 预取完成 → {self.pager.stats()}")
+
+        # L2 路由: landmark 检索
+        selection = self.router.route(query, query, self.summaries)
+        print(f"[路由层] top-{len(selection.chunk_ids)} 块: {selection.chunk_ids[:5]}...")
+        print(f"[路由层] 权重分布: {[f'{w:.2f}' for w in selection.weights[:5]]}...")
+
+        # L4 物理: 访问更新
+        for cid, score in zip(selection.chunk_ids, selection.scores):
+            self.pager.access(cid, abs(score), self.now)
+            self.metacog.promote(cid, abs(score))
+
+        # L3 存储: 汇总被选中块的压缩表示
+        gathered = []
+        for cid in selection.chunk_ids:
+            kv = (self.pager.vram.get(cid) or self.pager.ram.get(cid) or self.pager.disk.get(cid))
+            if kv:
+                gathered.append(kv.latent)
+        context = np.mean(gathered, axis=0) if gathered else np.zeros(self.store.DIM)
+        print(f"[存储层] 聚合 {len(gathered)} 块压缩表示 (每块 {self.store.bytes_per_token():.1f}B)")
+
+        # 置信度
+        conf = self.metacog.track_confidence(0.85)
+        print(f"[认知层] 置信度: {conf:.2f}")
+
+        # 驱逐维护
+        if self.pager.vram_used > self.pager.vram_limit * 0.9:
+            victim = self.pager.evict(self.now, "vram")
+            print(f"[物理层] 驱逐 {victim} → 降级")
+
+        self.now += 1
+        return context, selection
+
+
+# ============================================================
+# 验证
+# ============================================================
+if __name__ == "__main__":
+    print("=" * 56)
+    print("AgentFrame 四层融合 — 端到端验证")
+    print("=" * 56)
+
+    # 创建 Agent
+    agent = QuadLayerAgent(n_layers=27, quant_bits=4, top_k=8)
+
+    # 注入知识库 (12 块, 带标签)
+    knowledge = [
+        {"text": "KV 缓存 270KB/token 展开存储", "tags": ["data", "comparison"]},
+        {"text": "吸收式 MLA 缓存 576 维潜在向量", "tags": ["data", "method"]},
+        {"text": "INT8 思考链误差 0.011", "tags": ["data", "comparison"]},
+        {"text": "INT4 工具结果误差 0.079", "tags": ["data"]},
+        {"text": "per-channel 非对称量化 27 倍精度", "tags": ["method"]},
+        {"text": "L40S 实测 138 万 token 上下文", "tags": ["data", "result", "comparison"]},
+        {"text": "HiLS 分层软max 端到端块选择", "tags": ["method", "sparse"]},
+        {"text": "landmark token 块摘要检索", "tags": ["method", "sparse"]},
+        {"text": "Q-Cal 低秩校准 0.6% 参数", "tags": ["method"]},
+        {"text": "GQA 组内 max 聚合", "tags": ["method", "sparse"]},
+        {"text": "KV 换页 mmap 冷存储", "tags": ["method", "paging"]},
+        {"text": "预测性驱逐注意力信号", "tags": ["method", "paging"]},
+    ]
+    agent.ingest(knowledge)
+    print(f"[存储层] 注入 {len(knowledge)} 块知识")
+    print(f"[存储层] 每 token 缓存: {agent.store.bytes_per_token():.1f}B (原 270KB → 35.6×)")
+    print(f"[物理层] 初始: {agent.pager.stats()}")
+
+    # 执行任务
+    rng = np.random.default_rng(0)
+    query = rng.normal(0, 1, 576)
+
+    agent.run_task("对比压缩方案；分析稀疏路由；总结实测数据", query)
+
+    # 容量对比
+    seq_len = 100_000
+    original = 270 * 1024 * seq_len * 27
+    compressed = agent.store.memory_bytes(seq_len)
+    print(f"\n{'='*56}")
+    print(f"容量对比 (100K token × 27 层):")
+    print(f"  原始展开: {original/1024**3:.1f} GB")
+    print(f"  四层后:   {compressed/1024**3:.3f} GB")
+    print(f"  压缩比:   {original/compressed:.1f}×")
+    print(f"{'='*56}")