"""Integrate dynamic memory into a local Qwen3.5 checkpoint. The original Hugging Face checkpoint is loaded unchanged. Decoder layers are wrapped after loading, and only the separate memory module is trainable by default. """ from __future__ import annotations from dataclasses import asdict, dataclass from pathlib import Path from typing import Any, Optional import json import os import re import time import weakref import torch import torch.nn.functional as F from torch import Tensor, nn from .memory_os_v2 import ( KVBudgetManagerV2, MemoryOSV2, MemoryRouterV2, PagedMemoryBankV2, memory_record_to_dict, ) from .tiered_memory_store_v2 import TieredMemoryStoreV2 # Qwen's local tokenizer already contains this unused FIM control token. It # gives us a reset signal without changing the base vocabulary or resizing a # quantized embedding table. Callers may still pass any other token id. DEFAULT_MEMORY_RESET_TOKEN = "<|fim_prefix|>" def split_memory_candidates(text: str, *, max_candidates: int = 8) -> list[str]: """Split one user turn into bounded fact-sized write candidates. This is only a write-side segmentation aid. The model's internal policy still decides whether each candidate is durable; the read path remains entirely inside the model-owned memory bank. """ text = " ".join(str(text).strip().split()) if not text: return [] # A comma is not a safe write boundary. In natural conversation it is # common for the second half of a sentence to negate or qualify the # first half (for example, "mentioned X, but it is not important"). # Splitting there would make the policy see a misleading fragment and can # permanently store noise. Keep clauses intact and only split at hard # sentence boundaries. parts = [ part.strip() for part in re.split(r"(?<=[。！？!?；;])\s*", text) if part.strip() ] if len(parts) <= 1: return [text] if len(parts) <= max_candidates: return parts # Preserve all text when a long paragraph exceeds the safety bound. The # exact sequence can still be stored in one slot without silent loss. return [text] def looks_like_question(text: str) -> bool: """Reject question-shaped turns from the automatic durable-write path.""" normalized = " ".join(str(text).strip().split()) if not normalized: return False if any(mark in normalized for mark in ("？", "?")): return True return normalized.startswith( ( "请问", "什么", "为什么", "怎么", "如何", "能不能", "是否", "有没有", "如果", "帮我", "解释", ) ) @dataclass class QwenMemoryConfig: memory_slots: int = 16 memory_dim: int = 512 layer_indices: Optional[tuple[int, ...]] = None read_scale: float = 1.0 write_scale: float = 1.0 mode: str = "residual" blend_init: float = 0.0 direct_logit_scale: float = 0.0 write_token_offset: Optional[int] = None broadcast_write: bool = False raw_token_write: bool = False raw_logit_scale: float = 0.0 native_mode: bool = False persistent_memory: bool = False reset_token_id: Optional[int] = None summary_pooling: bool = True # Natural-language episodic memory keeps the exact token sequence of a # selected write inside the model-owned memory state. The model retrieves # and injects these tokens internally; callers never need to replay chat # history or assemble a prompt by hand. natural_language_memory: bool = False text_memory_tokens: int = 256 text_memory_top_k: int = 2 text_memory_threshold: float = 0.40 text_memory_write_threshold: float = 0.5 text_memory_replace_threshold: float = 0.35 # A slightly more permissive boundary for recognizing a same-attribute # update. Unrelated fragments still remain independent V2 records. text_memory_update_overlap_threshold: float = 0.30 text_memory_key_tokens: int = 128 text_memory_overlap_threshold: float = 0.22 automatic_memory: bool = False auto_memory_threshold: float = 0.35 # v2 adds a separately trained natural-language forget decision. Version # 1 keeps the old one-logit policy byte-compatible with existing adapters. automatic_memory_policy_version: int = 1 auto_forget_threshold: float = 0.50 # Natural Memory v2: compact addressed pages sit beside the v1 hot text # bank. V2 is opt-in so existing checkpoints remain byte-compatible. memory_version: int = 1 hierarchical_memory: bool = False memory_router_dim: int = 128 memory_router_heads: int = 8 memory_page_capacity: int = 32 # 32K pages x 32 records/page gives a one-million-record address space. # Pages are allocated lazily; this is a capacity, not a startup tensor. memory_max_pages: int = 32768 memory_hot_pages: int = 8 memory_top_k_pages: int = 4 memory_top_k_records: int = 8 memory_max_hops: int = 3 memory_coarse_index_bits: int = 20 memory_v2_read_threshold: float = 0.65 memory_v2_write_threshold: float = 0.50 # ``embedded`` keeps the portable user snapshot inside the safetensors # package. ``tiered`` additionally opens a durable page store so cold # records can leave RAM while the same model-owned reader remains active. memory_storage_mode: str = "embedded" memory_storage_path: Optional[str] = None memory_resident_pages: int = 256 # Embedded mode keeps every record in process RAM. Only this bounded hot # cache is copied to the execution device; it is not a second store. memory_gpu_cache_records: int = 256 memory_gpu_cache_tokens: int = 131072 # Keep a large VRAM safety margin for the model weights, activations and # KV cache. Memory records that do not fit remain in system RAM. memory_gpu_cache_reserve_mb: int = 2048 memory_gpu_cache_adaptive: bool = True # HF generation can keep the active cache on CPU while the current layer # is executing. This is opt-in because it trades GPU memory for PCIe # traffic and is therefore not always faster on short prompts. kv_offload: bool = False kv_cache_implementation: Optional[str] = None kv_offload_only_non_sliding: bool = True # When a prompt exceeds the hot-window budget, archive the old prefix as # V2 context records and run Qwen only on the recent window. The reader # remains model-owned; callers do not have to assemble a retrieval prompt. auto_compact_context: bool = True context_chunk_tokens: int = 512 context_archive_max_records: int = 1_048_576 kv_budget_tokens: int = 32768 kv_hard_max_tokens: int = 131072 kv_compaction_trigger: float = 0.90 kv_keep_recent_tokens: int = 8192 def __post_init__(self) -> None: if self.mode not in {"residual", "blend", "replace"}: raise ValueError("mode must be one of: residual, blend, replace") if not 0.0 <= self.blend_init <= 1.0: raise ValueError("blend_init must be between 0 and 1") if self.direct_logit_scale < 0.0: raise ValueError("direct_logit_scale must be non-negative") if self.write_token_offset is not None and self.write_token_offset < 1: raise ValueError("write_token_offset must be positive") if self.raw_logit_scale < 0.0: raise ValueError("raw_logit_scale must be non-negative") if self.text_memory_tokens < 1: raise ValueError("text_memory_tokens must be positive") if self.text_memory_top_k < 1: raise ValueError("text_memory_top_k must be positive") if not 0.0 <= self.text_memory_threshold <= 1.0: raise ValueError("text_memory_threshold must be between 0 and 1") if not 0.0 <= self.text_memory_write_threshold <= 1.0: raise ValueError("text_memory_write_threshold must be between 0 and 1") if not -1.0 <= self.text_memory_replace_threshold <= 1.0: raise ValueError("text_memory_replace_threshold must be between -1 and 1") if not -1.0 <= self.text_memory_update_overlap_threshold <= 1.0: raise ValueError("text_memory_update_overlap_threshold must be between -1 and 1") if self.text_memory_key_tokens < 1: raise ValueError("text_memory_key_tokens must be positive") if not 0.0 <= self.text_memory_overlap_threshold <= 1.0: raise ValueError("text_memory_overlap_threshold must be between 0 and 1") if not 0.0 <= self.auto_memory_threshold <= 1.0: raise ValueError("auto_memory_threshold must be between 0 and 1") if self.automatic_memory_policy_version < 1: raise ValueError("automatic_memory_policy_version must be positive") if not 0.0 <= self.auto_forget_threshold <= 1.0: raise ValueError("auto_forget_threshold must be between 0 and 1") if self.memory_version < 1: raise ValueError("memory_version must be positive") if self.memory_router_heads < 1: raise ValueError("memory_router_heads must be positive") if self.memory_router_dim < 8 or self.memory_router_dim % self.memory_router_heads != 0: raise ValueError("memory_router_dim must be divisible by memory_router_heads") if self.memory_page_capacity < 1 or self.memory_max_pages < 1: raise ValueError("memory page limits must be positive") if self.memory_top_k_pages < 1 or self.memory_top_k_records < 1: raise ValueError("memory top-k limits must be positive") if self.memory_max_hops < 1: raise ValueError("memory_max_hops must be positive") if not 0.0 <= self.memory_v2_read_threshold <= 1.0: raise ValueError("memory_v2_read_threshold must be between 0 and 1") if not 0.0 <= self.memory_v2_write_threshold <= 1.0: raise ValueError("memory_v2_write_threshold must be between 0 and 1") if self.memory_storage_mode not in {"embedded", "tiered"}: raise ValueError("memory_storage_mode must be embedded or tiered") if self.memory_resident_pages < self.memory_hot_pages: raise ValueError("memory_resident_pages must be >= memory_hot_pages") if self.memory_gpu_cache_records < 0 or self.memory_gpu_cache_tokens < 0: raise ValueError("memory GPU cache limits must be non-negative") if self.memory_gpu_cache_reserve_mb < 0: raise ValueError("memory_gpu_cache_reserve_mb must be non-negative") if self.kv_cache_implementation is not None and not str(self.kv_cache_implementation).strip(): raise ValueError("kv_cache_implementation must be non-empty when provided") if self.context_chunk_tokens < 1: raise ValueError("context_chunk_tokens must be positive") if self.context_archive_max_records < 1: raise ValueError("context_archive_max_records must be positive") if not 1 <= self.kv_budget_tokens <= self.kv_hard_max_tokens: raise ValueError("kv_budget_tokens must be inside [1, kv_hard_max_tokens]") def resolved_layers(self, num_hidden_layers: int) -> tuple[int, ...]: if self.layer_indices is not None: layers = tuple(sorted(set(self.layer_indices))) else: stride = max(1, num_hidden_layers // 4) layers = tuple(min(num_hidden_layers - 1, stride * i - 1) for i in range(1, 5)) layers = tuple(sorted(set(layers))) if not layers or min(layers) < 0 or max(layers) >= num_hidden_layers: raise ValueError(f"layer_indices must be inside [0, {num_hidden_layers})") return layers def load_memory_config(adapter_dir: str | Path) -> QwenMemoryConfig: """Load the architecture/configuration metadata saved with an adapter.""" metadata_path = Path(adapter_dir) / "memory_config.json" metadata = json.loads(metadata_path.read_text(encoding="utf-8")) saved = metadata.get("memory_config", {}) saved_layers = metadata.get("layer_indices", saved.get("layer_indices")) return QwenMemoryConfig( memory_slots=int(saved.get("memory_slots", 16)), memory_dim=int(saved.get("memory_dim", 512)), layer_indices=tuple(saved_layers) if saved_layers is not None else None, read_scale=float(saved.get("read_scale", 1.0)), write_scale=float(saved.get("write_scale", 1.0)), mode=str(saved.get("mode", "residual")), blend_init=float(saved.get("blend_init", 0.0)), direct_logit_scale=float(saved.get("direct_logit_scale", 0.0)), write_token_offset=( int(saved["write_token_offset"]) if saved.get("write_token_offset") is not None else None ), broadcast_write=bool(saved.get("broadcast_write", False)), raw_token_write=bool(saved.get("raw_token_write", False)), raw_logit_scale=float(saved.get("raw_logit_scale", 0.0)), native_mode=bool(saved.get("native_mode", False)), persistent_memory=bool(saved.get("persistent_memory", False)), reset_token_id=(int(saved["reset_token_id"]) if saved.get("reset_token_id") is not None else None), summary_pooling=bool(saved.get("summary_pooling", True)), natural_language_memory=bool(saved.get("natural_language_memory", False)), text_memory_tokens=int(saved.get("text_memory_tokens", 256)), text_memory_top_k=int(saved.get("text_memory_top_k", 2)), text_memory_threshold=float(saved.get("text_memory_threshold", 0.40)), text_memory_write_threshold=float(saved.get("text_memory_write_threshold", 0.5)), text_memory_replace_threshold=float(saved.get("text_memory_replace_threshold", 0.35)), text_memory_update_overlap_threshold=float( saved.get("text_memory_update_overlap_threshold", 0.30) ), text_memory_key_tokens=int(saved.get("text_memory_key_tokens", 128)), text_memory_overlap_threshold=float(saved.get("text_memory_overlap_threshold", 0.22)), automatic_memory=bool(saved.get("automatic_memory", False)), auto_memory_threshold=float(saved.get("auto_memory_threshold", 0.35)), automatic_memory_policy_version=int(saved.get("automatic_memory_policy_version", 1)), auto_forget_threshold=float(saved.get("auto_forget_threshold", 0.50)), memory_version=int(saved.get("memory_version", 1)), hierarchical_memory=bool(saved.get("hierarchical_memory", False)), memory_router_dim=int(saved.get("memory_router_dim", 128)), memory_router_heads=int(saved.get("memory_router_heads", 8)), memory_page_capacity=int(saved.get("memory_page_capacity", 32)), memory_max_pages=int(saved.get("memory_max_pages", 32768)), memory_hot_pages=int(saved.get("memory_hot_pages", 8)), memory_top_k_pages=int(saved.get("memory_top_k_pages", 4)), memory_top_k_records=int(saved.get("memory_top_k_records", 8)), memory_max_hops=int(saved.get("memory_max_hops", 3)), memory_coarse_index_bits=int(saved.get("memory_coarse_index_bits", 20)), memory_v2_read_threshold=float(saved.get("memory_v2_read_threshold", 0.50)), memory_v2_write_threshold=float(saved.get("memory_v2_write_threshold", 0.50)), memory_storage_mode=str(saved.get("memory_storage_mode", "embedded")), memory_storage_path=(str(saved["memory_storage_path"]) if saved.get("memory_storage_path") else None), memory_resident_pages=int(saved.get("memory_resident_pages", 256)), memory_gpu_cache_records=int(saved.get("memory_gpu_cache_records", 256)), memory_gpu_cache_tokens=int(saved.get("memory_gpu_cache_tokens", 131072)), memory_gpu_cache_reserve_mb=int(saved.get("memory_gpu_cache_reserve_mb", 2048)), memory_gpu_cache_adaptive=bool(saved.get("memory_gpu_cache_adaptive", True)), kv_offload=bool(saved.get("kv_offload", False)), kv_cache_implementation=( str(saved["kv_cache_implementation"]) if saved.get("kv_cache_implementation") is not None else None ), kv_offload_only_non_sliding=bool(saved.get("kv_offload_only_non_sliding", True)), auto_compact_context=bool(saved.get("auto_compact_context", True)), context_chunk_tokens=int(saved.get("context_chunk_tokens", 512)), context_archive_max_records=int(saved.get("context_archive_max_records", 1_048_576)), kv_budget_tokens=int(saved.get("kv_budget_tokens", 32768)), kv_hard_max_tokens=int(saved.get("kv_hard_max_tokens", 131072)), kv_compaction_trigger=float(saved.get("kv_compaction_trigger", 0.90)), kv_keep_recent_tokens=int(saved.get("kv_keep_recent_tokens", 8192)), ) def resolve_memory_reset_token(tokenizer: Any, token: str = DEFAULT_MEMORY_RESET_TOKEN) -> int: """Resolve a tokenizer control token suitable for clearing memory.""" token_id = tokenizer.convert_tokens_to_ids(token) unknown_id = getattr(tokenizer, "unk_token_id", None) if token_id is None or (unknown_id is not None and int(token_id) == int(unknown_id)): raise ValueError( f"reset token {token!r} is not present in this tokenizer; pass --reset-token-id instead" ) return int(token_id) class QwenDynamicMemory(nn.Module): """A compact float32 key-value memory connected to Qwen hidden states.""" def __init__(self, hidden_size: int, config: QwenMemoryConfig) -> None: super().__init__() self.hidden_size = hidden_size self.config = config dim = config.memory_dim self.query = nn.Linear(hidden_size, dim, bias=False) self.key = nn.Linear(dim, dim, bias=False) self.value = nn.Linear(dim, dim, bias=False) self.read_out = nn.Linear(dim, hidden_size, bias=False) self.read_gate = nn.Linear(hidden_size, 1) self.slot_keys = nn.Parameter(torch.randn(config.memory_slots, dim) / dim**0.5) self.write_key = nn.Linear(hidden_size, dim, bias=False) self.write_value = nn.Linear(hidden_size, dim, bias=False) self.write_gate = nn.Linear(hidden_size, config.memory_slots) self.apply(self._init_weights) self.last_read_address: Optional[Tensor] = None self.last_read_relevance: Optional[Tensor] = None self.last_write_address: Optional[Tensor] = None self.last_write_probability: Optional[Tensor] = None @staticmethod def _init_weights(module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) def initial_state(self, batch_size: int, *, device: torch.device) -> Tensor: return torch.zeros( batch_size, self.config.memory_slots, self.config.memory_dim, device=device, dtype=self.slot_keys.dtype, ) def read(self, hidden_states: Tensor, memory: Tensor) -> Tensor: work_dtype = self.slot_keys.dtype hidden = hidden_states.to(work_dtype) memory = memory.to(device=hidden.device, dtype=work_dtype) q = self.query(hidden) k = self.key(memory) v = self.value(memory) scores = torch.matmul(q, k.transpose(-1, -2)) / self.config.memory_dim**0.5 address = scores.softmax(dim=-1) retrieved = torch.matmul(address, v) retrieved = self.read_out(retrieved) gate = torch.sigmoid(self.read_gate(hidden)) self.last_read_address = address self.last_read_relevance = gate * address.max(dim=-1, keepdim=True).values return (self.config.read_scale * gate * retrieved).to(hidden_states.dtype) def update( self, hidden_states: Tensor, memory: Tensor, attention_mask: Optional[Tensor] = None, ) -> Tensor: work_dtype = self.slot_keys.dtype hidden = hidden_states.to(work_dtype) memory = memory.to(device=hidden.device, dtype=work_dtype) if self.config.write_token_offset is None: summary = hidden[:, -1] else: if hidden.shape[1] < self.config.write_token_offset: raise ValueError("write_token_offset exceeds the memory sequence length") summary = hidden[:, -self.config.write_token_offset] proposal = self.write_value(summary) if self.config.broadcast_write: strength = torch.ones( hidden.shape[0], self.config.memory_slots, device=hidden.device, dtype=work_dtype, ) address = strength / float(self.config.memory_slots) else: address = (self.write_key(summary) @ self.slot_keys.t()).softmax(dim=-1) strength = torch.sigmoid(self.write_gate(summary)) * address self.last_write_address = address self.last_write_probability = strength.sum(dim=-1, keepdim=True) strength = (self.config.write_scale * strength).unsqueeze(-1) proposal = proposal[:, None, :].expand(-1, self.config.memory_slots, -1) return memory + strength * (proposal - memory) class NaturalLanguageRetriever(nn.Module): """Trainable pair scorer for query-to-episodic-memory retrieval.""" def __init__(self, hidden_size: int, projection_size: int = 256) -> None: super().__init__() self.query_projection = nn.Linear(hidden_size, projection_size, bias=False) self.key_projection = nn.Linear(hidden_size, projection_size, bias=False) self.pair_scorer = nn.Sequential( nn.Linear(projection_size * 4, projection_size), nn.SiLU(), nn.Linear(projection_size, 1), ) self.apply(self._init_weights) @staticmethod def _init_weights(module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) def forward(self, query: Tensor, keys: Tensor) -> Tensor: query = F.normalize(query.float(), dim=-1) keys = F.normalize(keys.float(), dim=-1) query_projection = self.query_projection(query) key_projection = self.key_projection(keys) if key_projection.ndim == 3: query_projection = query_projection.unsqueeze(1) features = torch.cat( ( query_projection.expand_as(key_projection), key_projection, query_projection.expand_as(key_projection) * key_projection, (query_projection.expand_as(key_projection) - key_projection).abs(), ), dim=-1, ) return self.pair_scorer(features).squeeze(-1) class AutomaticMemoryPolicy(nn.Module): """High-recall controller for deciding whether a user turn is durable.""" def __init__(self, hidden_size: int, projection_size: int = 256) -> None: super().__init__() self.encoder = nn.Sequential( nn.Linear(hidden_size, projection_size), nn.LayerNorm(projection_size), nn.SiLU(), nn.Linear(projection_size, projection_size // 2), nn.SiLU(), ) self.importance = nn.Linear(projection_size // 2, 1) self.apply(self._init_weights) @staticmethod def _init_weights(module: nn.Module) -> None: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) if module.bias is not None: nn.init.zeros_(module.bias) def forward(self, hidden: Tensor) -> Tensor: return self.importance(self.encoder(hidden.float())).squeeze(-1) class AutomaticMemoryPolicyV2(AutomaticMemoryPolicy): """Write/forget policy trained on the same natural-language turn state.""" def __init__(self, hidden_size: int, projection_size: int = 256) -> None: super().__init__(hidden_size, projection_size) self.forget = nn.Linear(projection_size // 2, 1) nn.init.normal_(self.forget.weight, mean=0.0, std=0.02) nn.init.zeros_(self.forget.bias) def _features(self, hidden: Tensor) -> Tensor: return self.encoder(hidden.float()) def forward(self, hidden: Tensor) -> Tensor: return self.importance(self._features(hidden)).squeeze(-1) def forget_logits(self, hidden: Tensor) -> Tensor: return self.forget(self._features(hidden)).squeeze(-1) class NativeQwenDynamicMemory(QwenDynamicMemory): """Learned write/forget controller for model-owned persistent memory.""" def __init__(self, hidden_size: int, config: QwenMemoryConfig) -> None: super().__init__(hidden_size, config) dim = config.memory_dim self.summary_score = nn.Linear(hidden_size, 1, bias=False) self.write_decision = nn.Linear(hidden_size, 1) # Forgetting must depend on both the incoming candidate and what is # already stored. A hidden-only gate can learn "this sentence looks # like a fact", but it cannot learn "this fact conflicts with the # value in the addressed slot". self.forget_gate = nn.Sequential( nn.Linear(hidden_size + (2 * dim), dim), nn.SiLU(), nn.Linear(dim, config.memory_slots), ) nn.init.normal_(self.summary_score.weight, mean=0.0, std=0.02) nn.init.normal_(self.write_decision.weight, mean=0.0, std=0.02) for module in self.forget_gate: if isinstance(module, nn.Linear): nn.init.normal_(module.weight, mean=0.0, std=0.02) nn.init.constant_(self.write_decision.bias, -1.5) nn.init.constant_(self.forget_gate[-1].bias, -2.0) self.last_write_probability: Optional[Tensor] = None self.last_forget_probability: Optional[Tensor] = None # Keep the semantic summary available to the automatic policy. The # older ``last_write_representation`` is optimized for the memory # value path and may discard distinctions such as negation or a # question-shaped turn. self.last_write_summary: Optional[Tensor] = None self.last_write_representation: Optional[Tensor] = None def _summary(self, hidden_states: Tensor, attention_mask: Optional[Tensor]) -> Tensor: work_dtype = self.slot_keys.dtype hidden = hidden_states.to(work_dtype) if self.config.write_token_offset is not None: if hidden.shape[1] < self.config.write_token_offset: raise ValueError("write_token_offset exceeds the memory sequence length") return hidden[:, -self.config.write_token_offset] if not self.config.summary_pooling: return hidden[:, -1] scores = self.summary_score(hidden).squeeze(-1) if attention_mask is not None and attention_mask.ndim == 2: # During cached generation the decoder may expose only the newest # token while the attention mask still covers the full sequence. # Align the mask to the score sequence instead of assuming both # lengths are identical. score_length = scores.shape[1] mask = attention_mask.to(device=hidden.device, dtype=torch.bool) if mask.shape[1] > score_length: mask = mask[:, -score_length:] elif mask.shape[1] < score_length: pad = torch.ones( mask.shape[0], score_length - mask.shape[1], device=mask.device, dtype=mask.dtype, ) mask = torch.cat((mask, pad), dim=1) scores = scores.masked_fill(~mask, torch.finfo(scores.dtype).min) weights = scores.softmax(dim=-1) return (weights.unsqueeze(-1) * hidden).sum(dim=1) def update( self, hidden_states: Tensor, memory: Tensor, attention_mask: Optional[Tensor] = None, ) -> Tensor: work_dtype = self.slot_keys.dtype memory = memory.to(device=hidden_states.device, dtype=work_dtype) summary = self._summary(hidden_states, attention_mask) self.last_write_summary = summary proposal = self.write_value(summary) self.last_write_representation = self.read_out(self.value(proposal)) address = (self.write_key(summary) @ self.slot_keys.t()).softmax(dim=-1) slot_probability = torch.sigmoid(self.write_gate(summary)) write_probability = torch.sigmoid(self.write_decision(summary)) self.last_write_address = address matched_memory = torch.sum(address.unsqueeze(-1) * memory, dim=1) forget_features = torch.cat((summary, proposal, matched_memory), dim=-1) forget_probability = torch.sigmoid(self.forget_gate(forget_features)) strength = self.config.write_scale * write_probability * slot_probability * address erase = strength * forget_probability proposal = proposal[:, None, :] updated = memory * (1.0 - erase.unsqueeze(-1)) updated = updated + strength.unsqueeze(-1) * (proposal - updated) self.last_write_probability = write_probability self.last_forget_probability = forget_probability return updated class _MemoryRuntime: def __init__(self, memory: QwenDynamicMemory) -> None: self.memory = memory self.state: Optional[Tensor] = None self.read_enabled = True self.update_enabled = True self.last_read: Optional[Tensor] = None self.raw_memory: Optional[Tensor] = None self.input_ids: Optional[Tensor] = None self.attention_mask: Optional[Tensor] = None self.output_embeddings: Optional[nn.Module] = None self.reset_mask: Optional[Tensor] = None self.text_token_ids: Optional[Tensor] = None self.text_token_mask: Optional[Tensor] = None self.text_slot_valid: Optional[Tensor] = None self.text_slot_keys: Optional[Tensor] = None self.text_slot_age: Optional[Tensor] = None self.text_write_counter: Optional[Tensor] = None self.text_key_token_ids: Optional[Tensor] = None self.text_key_token_mask: Optional[Tensor] = None self.text_last_written_slot: Optional[Tensor] = None self.text_read_slots: Optional[Tensor] = None self.text_read_relevance: Optional[Tensor] = None self.text_read_overlap: Optional[Tensor] = None self.v2_query_key: Optional[Tensor] = None self.v2_last_decisions: list[Any] = [] self.text_prefix_used: bool = False self.auto_memory_probability: Optional[Tensor] = None self.auto_memory_forget_probability: Optional[Tensor] = None self.context_compaction: Optional[dict[str, Any]] = None # Persistent checkpoints are the durable user state. Evaluation, # branch execution and a new temporary conversation must be able to # reset only their working state without mutating that checkpoint. self.use_persistent_state: bool = True class MemoryLayerAdapter(nn.Module): """Wrap one Qwen decoder layer and optionally replace its token mixer.""" def __init__( self, inner: nn.Module, runtime: _MemoryRuntime, *, read: bool, write: bool, mode: str = "residual", blend_init: float = 0.0, ) -> None: super().__init__() if mode not in {"residual", "blend", "replace"}: raise ValueError("mode must be one of: residual, blend, replace") self.inner = inner self._runtime_ref = weakref.ref(runtime) self.read_enabled = read self.write_enabled = write self.mode = mode if mode == "blend": blend_init = min(max(float(blend_init), 1e-4), 1.0 - 1e-4) logit = torch.logit(torch.tensor(blend_init, dtype=torch.float32)) inner_device = next( (parameter.device for parameter in inner.parameters() if parameter.device.type != "meta"), None, ) if inner_device is not None: logit = logit.to(inner_device) self.blend_logit = nn.Parameter(logit) def _original_token_mixer( self, normalized_hidden: Tensor, position_embeddings: Any, attention_mask: Optional[Tensor], position_ids: Optional[Tensor], past_key_values: Any, kwargs: dict[str, Any], ) -> Tensor: layer_type = getattr(self.inner, "layer_type", None) if layer_type == "linear_attention": output = self.inner.linear_attn( hidden_states=normalized_hidden, cache_params=past_key_values, attention_mask=attention_mask, **kwargs, ) return output[0] if isinstance(output, tuple) else output if layer_type == "full_attention": output = self.inner.self_attn( hidden_states=normalized_hidden, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, position_embeddings=position_embeddings, **kwargs, ) return output[0] if isinstance(output, tuple) else output raise TypeError(f"unsupported Qwen layer_type for surgery: {layer_type!r}") def _prefetch_cache_layer(self, past_key_values: Any) -> None: """Make an offloaded Qwen hybrid-cache layer ready before it is read.""" if past_key_values is None or not getattr(past_key_values, "offloading", False): return component = getattr(self.inner, "self_attn", None) if component is None: component = getattr(self.inner, "linear_attn", None) layer_index = getattr(component, "layer_idx", None) if layer_index is None or not hasattr(past_key_values, "prefetch"): return only_non_sliding = bool(getattr(past_key_values, "only_non_sliding", True)) parameter = next(component.parameters(), None) device = parameter.device if parameter is not None else torch.device("cpu") past_key_values.prefetch(int(layer_index), only_non_sliding) # The stock offload helper uses a separate CUDA stream. Qwen3.5 # reads linear-attention state before it calls ``update_*``; make the # current layer's ownership synchronous here so a CPU state can never # reach a CUDA convolution or attention kernel. if 0 <= int(layer_index) < len(past_key_values.layers): cache_layer = past_key_values.layers[int(layer_index)] for name in ("keys", "values", "conv_states", "recurrent_states"): value = getattr(cache_layer, name, None) if isinstance(value, Tensor) and value.device != device: setattr(cache_layer, name, value.to(device=device, non_blocking=False)) def _memory_read(self, hidden_states: Tensor, runtime: _MemoryRuntime) -> Tensor: if runtime.read_enabled and self.read_enabled and runtime.state is not None: output = runtime.memory.read(hidden_states, runtime.state) else: output = torch.zeros_like(hidden_states) if self.write_enabled: runtime.last_read = output return output def _surgical_forward( self, hidden_states: Tensor, position_embeddings: Any, attention_mask: Optional[Tensor], position_ids: Optional[Tensor], past_key_values: Any, kwargs: dict[str, Any], runtime: _MemoryRuntime, ) -> Tensor: residual = hidden_states normalized_hidden = self.inner.input_layernorm(hidden_states) memory_output = self._memory_read(normalized_hidden, runtime) if self.mode == "replace": token_mixer = memory_output else: original_output = self._original_token_mixer( normalized_hidden, position_embeddings, attention_mask, position_ids, past_key_values, kwargs, ) mix = torch.sigmoid(self.blend_logit).to(dtype=original_output.dtype) # A zero/empty memory must be an exact no-op. This prevents a # trained blend coefficient from attenuating the original Qwen # token mixer on ordinary prompts that have no user memory. if runtime.state is None or not runtime.read_enabled: # Exact zero in the forward pass, but retain a surrogate # gradient so the blend parameter remains trainable in unit # tests and future calibration runs. mix = mix - mix.detach() else: memory_activity = runtime.state.detach().abs().mean() if float(memory_activity) <= 1e-6: mix = mix - mix.detach() token_mixer = (1.0 - mix) * original_output + mix * memory_output hidden_states = residual + token_mixer residual = hidden_states hidden_states = self.inner.post_attention_layernorm(hidden_states) hidden_states = self.inner.mlp(hidden_states) return residual + hidden_states def forward( self, hidden_states: Tensor, position_embeddings: Any = None, attention_mask: Optional[Tensor] = None, position_ids: Optional[Tensor] = None, past_key_values: Any = None, **kwargs: Any, ) -> Any: runtime = self._runtime_ref() if runtime is None: return self.inner( hidden_states, position_embeddings, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, **kwargs, ) self._prefetch_cache_layer(past_key_values) if self.mode == "residual": hidden_states = hidden_states + self._memory_read(hidden_states, runtime) output = self.inner( hidden_states, position_embeddings, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, **kwargs, ) else: output = self._surgical_forward( hidden_states, position_embeddings, attention_mask, position_ids, past_key_values, kwargs, runtime, ) hidden_output = output[0] if isinstance(output, tuple) else output if runtime.update_enabled and self.write_enabled and runtime.state is not None: updated_state = runtime.memory.update( hidden_output, runtime.state, attention_mask=runtime.attention_mask, ) if runtime.reset_mask is not None: reset_mask = runtime.reset_mask.to(device=updated_state.device, dtype=torch.bool) updated_state = torch.where( reset_mask[:, None, None], torch.zeros_like(updated_state), updated_state, ) runtime.reset_mask = None runtime.state = updated_state if ( runtime.memory.config.raw_token_write and not runtime.read_enabled and runtime.input_ids is not None ): if runtime.output_embeddings is None: raise RuntimeError("raw_token_write requires output embeddings") offset = runtime.memory.config.write_token_offset or 1 if runtime.attention_mask is not None and runtime.attention_mask.ndim == 2: last_valid = runtime.attention_mask.sum(dim=-1).to(torch.long) - 1 positions = (last_valid - (offset - 1)).clamp_min(0) token_ids = runtime.input_ids.gather(1, positions[:, None]).squeeze(1) else: token_ids = runtime.input_ids[:, -offset] # Store the output-projection row, not the input embedding row. # Qwen checkpoints are allowed to untie these matrices, and the # generation-time pointer is consumed through output_embeddings. runtime.raw_memory = runtime.output_embeddings.weight[token_ids].detach() return output class QwenDynamicMemoryOutput: """Proxy for a normal HF output with one added field: ``memory``.""" def __init__(self, base_output: Any, memory: Optional[Tensor]) -> None: self.base_output = base_output self.memory = memory def __getattr__(self, name: str) -> Any: return getattr(self.base_output, name) def __getitem__(self, key: Any) -> Any: return self.base_output[key] class QwenDynamicMemoryModel(nn.Module): """Qwen3.5 with caller-owned persistent dynamic memory state.""" def __init__( self, base_model: nn.Module, memory_config: Optional[QwenMemoryConfig] = None, *, freeze_backbone: bool = True, ) -> None: super().__init__() self.base_model = base_model text_config = base_model.config.text_config self.memory_config = memory_config or QwenMemoryConfig() self.layer_indices = self.memory_config.resolved_layers(text_config.num_hidden_layers) memory_type = NativeQwenDynamicMemory if self.memory_config.native_mode else QwenDynamicMemory self.memory = memory_type(text_config.hidden_size, self.memory_config) self.text_retriever = ( NaturalLanguageRetriever(text_config.hidden_size) if self.memory_config.natural_language_memory else None ) policy_type = ( AutomaticMemoryPolicyV2 if self.memory_config.automatic_memory_policy_version >= 2 else AutomaticMemoryPolicy ) self.memory_policy = ( policy_type(text_config.hidden_size) if self.memory_config.natural_language_memory and self.memory_config.automatic_memory else None ) self.memory_router_v2: Optional[MemoryRouterV2] = None self.memory_os_v2: Optional[MemoryOSV2] = None self._memory_router_v2_ready = False if self.memory_config.hierarchical_memory or self.memory_config.memory_version >= 2: self.memory_router_v2 = MemoryRouterV2( text_config.hidden_size, router_dim=self.memory_config.memory_router_dim, num_heads=self.memory_config.memory_router_heads, max_hops=self.memory_config.memory_max_hops, ) self.memory_os_v2 = self._new_memory_os_v2(text_config.hidden_size) self._text_retriever_ready = False self._memory_policy_ready = False self.runtime = _MemoryRuntime(self.memory) self.runtime.output_embeddings = self.base_model.get_output_embeddings() self._persistent_memory: Optional[Tensor] = None self.register_buffer("persistent_memory", torch.empty(0), persistent=True) self.register_buffer( "persistent_text_token_ids", torch.empty(0, dtype=torch.long), persistent=True, ) self.register_buffer( "persistent_text_token_mask", torch.empty(0, dtype=torch.bool), persistent=True, ) self.register_buffer( "persistent_text_slot_valid", torch.empty(0, dtype=torch.bool), persistent=True, ) self.register_buffer( "persistent_text_slot_keys", torch.empty(0), persistent=True, ) self.register_buffer( "persistent_text_slot_age", torch.empty(0, dtype=torch.long), persistent=True, ) self.register_buffer( "persistent_text_write_counter", torch.empty(0, dtype=torch.long), persistent=True, ) self.register_buffer( "persistent_text_key_token_ids", torch.empty(0, dtype=torch.long), persistent=True, ) self.register_buffer( "persistent_text_key_token_mask", torch.empty(0, dtype=torch.bool), persistent=True, ) if freeze_backbone: for parameter in self.base_model.parameters(): parameter.requires_grad_(False) language_model = self.base_model.model.language_model self._memory_adapters: list[MemoryLayerAdapter] = [] for layer_index in self.layer_indices: layer = language_model.layers[layer_index] adapter = MemoryLayerAdapter( layer, self.runtime, read=True, write=layer_index == self.layer_indices[-1], mode=self.memory_config.mode, blend_init=self.memory_config.blend_init, ) language_model.layers[layer_index] = adapter self._memory_adapters.append(adapter) self.memory.to(self._find_layer_device()) if self.memory_router_v2 is not None: self.memory_router_v2.to(self._find_layer_device()) if self.text_retriever is not None: self.text_retriever.to(self._find_layer_device()) if self.memory_policy is not None: self.memory_policy.to(self._find_layer_device()) if self.memory_config.persistent_memory: self.persistent_memory = self.memory.initial_state( 1, device=self._find_layer_device(), ) self._persistent_memory = self.persistent_memory.detach() if hasattr(self.base_model.config, "use_cache"): self.base_model.config.use_cache = False if hasattr(self.base_model.config, "text_config"): self.base_model.config.text_config.use_cache = False def _new_memory_os_v2(self, hidden_size: int) -> MemoryOSV2: """Create a V2 store from the configured page and candidate limits.""" if self.memory_router_v2 is None: raise RuntimeError("cannot create V2 memory without a router") tier_store = None if self.memory_config.memory_storage_mode == "tiered": if not self.memory_config.memory_storage_path: raise ValueError("tiered memory requires memory_storage_path") tier_store = TieredMemoryStoreV2( self.memory_config.memory_storage_path, key_dim=self.memory_config.memory_router_dim, page_capacity=self.memory_config.memory_page_capacity, ) bank = PagedMemoryBankV2( hidden_size, page_capacity=self.memory_config.memory_page_capacity, max_pages=self.memory_config.memory_max_pages, hot_pages=self.memory_config.memory_hot_pages, top_k_pages=self.memory_config.memory_top_k_pages, top_k_records=self.memory_config.memory_top_k_records, max_hops=self.memory_config.memory_max_hops, router=self.memory_router_v2, key_dim=self.memory_config.memory_router_dim, coarse_index_bits=self.memory_config.memory_coarse_index_bits, tier_store=tier_store, max_resident_pages=self.memory_config.memory_resident_pages, runtime_device=self._find_layer_device(), gpu_cache_records=self.memory_config.memory_gpu_cache_records, gpu_cache_tokens=self.memory_config.memory_gpu_cache_tokens, gpu_cache_reserve_mb=self.memory_config.memory_gpu_cache_reserve_mb, gpu_cache_adaptive=self.memory_config.memory_gpu_cache_adaptive, ) return MemoryOSV2( hidden_size, router=self.memory_router_v2, bank=bank, kv_budget=KVBudgetManagerV2( max_tokens=self.memory_config.kv_budget_tokens, hard_max_tokens=self.memory_config.kv_hard_max_tokens, compaction_trigger=self.memory_config.kv_compaction_trigger, keep_recent_tokens=self.memory_config.kv_keep_recent_tokens, ), read_threshold=self.memory_config.memory_v2_read_threshold, write_threshold=self.memory_config.memory_v2_write_threshold, ) def _ensure_text_memory(self, batch_size: int, *, device: torch.device) -> None: """Create the fixed-size model-owned natural-language memory bank.""" if not self.memory_config.natural_language_memory: return shape = ( batch_size, self.memory_config.memory_slots, self.memory_config.text_memory_tokens, ) if tuple(self.persistent_text_token_ids.shape) != shape: self.persistent_text_token_ids = torch.zeros( shape, dtype=torch.long, device=device, ) self.persistent_text_token_mask = torch.zeros( shape, dtype=torch.bool, device=device, ) self.persistent_text_slot_valid = torch.zeros( batch_size, self.memory_config.memory_slots, dtype=torch.bool, device=device, ) hidden_size = self.memory.hidden_size self.persistent_text_slot_keys = torch.zeros( batch_size, self.memory_config.memory_slots, hidden_size, dtype=torch.float32, device=device, ) self.persistent_text_slot_age = torch.full( (batch_size, self.memory_config.memory_slots), -1, dtype=torch.long, device=device, ) self.persistent_text_write_counter = torch.zeros( batch_size, dtype=torch.long, device=device, ) key_shape = ( batch_size, self.memory_config.memory_slots, self.memory_config.text_memory_key_tokens, ) self.persistent_text_key_token_ids = torch.zeros( key_shape, dtype=torch.long, device=device, ) self.persistent_text_key_token_mask = torch.zeros( key_shape, dtype=torch.bool, device=device, ) def _bind_text_memory(self, batch_size: int, *, device: torch.device) -> None: """Bind the current text bank to the per-call runtime.""" if not self.memory_config.natural_language_memory: self.runtime.text_token_ids = None self.runtime.text_token_mask = None self.runtime.text_slot_valid = None return if self.memory_config.persistent_memory and self.runtime.use_persistent_state: self._ensure_text_memory(1, device=device) else: self._ensure_text_memory(batch_size, device=device) if not (self.memory_config.persistent_memory and self.runtime.use_persistent_state): shape = ( batch_size, self.memory_config.memory_slots, self.memory_config.text_memory_tokens, ) key_shape = ( batch_size, self.memory_config.memory_slots, self.memory_config.text_memory_key_tokens, ) if ( self.runtime.text_token_ids is None or tuple(self.runtime.text_token_ids.shape) != shape or self.runtime.text_key_token_ids is None or tuple(self.runtime.text_key_token_ids.shape) != key_shape ): self.runtime.text_token_ids = torch.zeros(shape, dtype=torch.long, device=device) self.runtime.text_token_mask = torch.zeros(shape, dtype=torch.bool, device=device) self.runtime.text_slot_valid = torch.zeros( batch_size, self.memory_config.memory_slots, dtype=torch.bool, device=device ) self.runtime.text_slot_keys = torch.zeros( batch_size, self.memory_config.memory_slots, self.memory.hidden_size, dtype=torch.float32, device=device ) self.runtime.text_slot_age = torch.full( (batch_size, self.memory_config.memory_slots), -1, dtype=torch.long, device=device ) self.runtime.text_write_counter = torch.zeros(batch_size, dtype=torch.long, device=device) self.runtime.text_key_token_ids = torch.zeros(key_shape, dtype=torch.long, device=device) self.runtime.text_key_token_mask = torch.zeros(key_shape, dtype=torch.bool, device=device) return if self.persistent_text_token_ids.shape[0] == batch_size: self.runtime.text_token_ids = self.persistent_text_token_ids self.runtime.text_token_mask = self.persistent_text_token_mask self.runtime.text_slot_valid = self.persistent_text_slot_valid self.runtime.text_slot_keys = self.persistent_text_slot_keys self.runtime.text_slot_age = self.persistent_text_slot_age self.runtime.text_write_counter = self.persistent_text_write_counter self.runtime.text_key_token_ids = self.persistent_text_key_token_ids self.runtime.text_key_token_mask = self.persistent_text_key_token_mask return if self.persistent_text_token_ids.shape[0] == 1: self.runtime.text_token_ids = self.persistent_text_token_ids.expand( batch_size, -1, -1 ).clone() self.runtime.text_token_mask = self.persistent_text_token_mask.expand( batch_size, -1, -1 ).clone() self.runtime.text_slot_valid = self.persistent_text_slot_valid.expand( batch_size, -1 ).clone() self.runtime.text_slot_keys = self.persistent_text_slot_keys.expand( batch_size, -1, -1 ).clone() self.runtime.text_slot_age = self.persistent_text_slot_age.expand( batch_size, -1 ).clone() self.runtime.text_write_counter = self.persistent_text_write_counter.expand( batch_size ).clone() self.runtime.text_key_token_ids = self.persistent_text_key_token_ids.expand( batch_size, -1, -1 ).clone() self.runtime.text_key_token_mask = self.persistent_text_key_token_mask.expand( batch_size, -1, -1 ).clone() return raise ValueError( "natural-language memory batch does not match the loaded persistent bank" ) def _clear_text_memory(self, mask: Optional[Tensor] = None) -> None: if self.runtime.text_token_ids is None: return if mask is None: self.runtime.text_token_ids.zero_() self.runtime.text_token_mask.zero_() self.runtime.text_slot_valid.zero_() self.runtime.text_slot_keys.zero_() self.runtime.text_slot_age.fill_(-1) self.runtime.text_write_counter.zero_() self.runtime.text_key_token_ids.zero_() self.runtime.text_key_token_mask.zero_() self.runtime.text_last_written_slot = None self.runtime.text_prefix_used = False return mask = mask.to(device=self.runtime.text_token_ids.device, dtype=torch.bool) self.runtime.text_token_ids[mask] = 0 self.runtime.text_token_mask[mask] = False self.runtime.text_slot_valid[mask] = False self.runtime.text_slot_keys[mask] = 0 self.runtime.text_slot_age[mask] = -1 self.runtime.text_write_counter[mask] = 0 self.runtime.text_key_token_ids[mask] = 0 self.runtime.text_key_token_mask[mask] = False self.runtime.text_last_written_slot = None @torch.no_grad() def _encode_model_key(self, input_ids: Tensor, attention_mask: Tensor) -> Tensor: """Encode a memory key with the frozen original Qwen representation.""" previous_read = self.runtime.read_enabled previous_update = self.runtime.update_enabled self.runtime.read_enabled = False self.runtime.update_enabled = False try: embedding_layer = self.base_model.get_input_embeddings() input_ids = input_ids.to(embedding_layer.weight.device) attention_mask = attention_mask.to(input_ids.device) output = self.base_model( input_ids=input_ids, attention_mask=attention_mask, use_cache=False, output_hidden_states=True, ) hidden = output.hidden_states[-1].float() weights = attention_mask.to(device=hidden.device, dtype=hidden.dtype).unsqueeze(-1) key = (hidden * weights).sum(dim=1) / weights.sum(dim=1).clamp_min(1.0) return F.normalize(key, dim=-1) finally: self.runtime.read_enabled = previous_read self.runtime.update_enabled = previous_update @torch.no_grad() def compact_context_for_kv( self, input_ids: Tensor, attention_mask: Optional[Tensor] = None, *, archive: bool = True, chunk_tokens: Optional[int] = None, ) -> tuple[Tensor, Tensor, dict[str, Any]]: """Archive the old prefix and return a bounded hot-context window. This is the model-side boundary between the exact working KV and the lossy long-term memory path. It never sends the current token to all pages: each archived chunk gets one address, and future reads still use the V2 coarse-page -> exact-record -> Top-K route. The method accepts already-tokenized input because it is also used by generation. The archived token sequence is kept losslessly inside a V2 record; its semantic address is produced by the frozen Qwen hidden representation. A caller may disable ``archive`` for measurement, in which case the method only reports the retention plan. """ if input_ids.ndim != 2: raise ValueError("input_ids must have shape [batch, sequence]") if attention_mask is None: attention_mask = torch.ones_like(input_ids) if attention_mask.shape != input_ids.shape: raise ValueError("attention_mask must have the same shape as input_ids") budget = self.memory_os_v2.kv_budget if self.memory_os_v2 is not None else None if ( self.memory_os_v2 is None or not self.memory_config.auto_compact_context or budget is None ): return input_ids, attention_mask, { "compacted": False, "reason": "disabled", "archived_records": 0, "original_tokens": [int(item) for item in attention_mask.sum(dim=-1).tolist()], "retained_tokens": [int(item) for item in attention_mask.sum(dim=-1).tolist()], } chunk_size = int(chunk_tokens or self.memory_config.context_chunk_tokens) if chunk_size < 1: raise ValueError("chunk_tokens must be positive") hot_limit = int(budget.max_tokens) pending: list[dict[str, Any]] = [] pending_previous: list[Optional[int]] = [] retained: list[Tensor] = [] original_lengths: list[int] = [] retained_lengths: list[int] = [] for batch_index in range(input_ids.shape[0]): valid = input_ids[batch_index][attention_mask[batch_index].to(dtype=torch.bool)] original_lengths.append(int(valid.numel())) if valid.numel() <= hot_limit: retained.append(valid) retained_lengths.append(int(valid.numel())) continue archived_prefix = valid[:-hot_limit] retained_window = valid[-hot_limit:] retained.append(retained_window) retained_lengths.append(int(retained_window.numel())) previous_index: Optional[int] = None for start in range(0, int(archived_prefix.numel()), chunk_size): end = min(start + chunk_size, int(archived_prefix.numel())) if len(pending) >= self.memory_config.context_archive_max_records: raise RuntimeError( "context archive record limit reached; increase " "context_archive_max_records before compacting this input" ) chunk = archived_prefix[start:end].detach() chunk_mask = torch.ones_like(chunk, dtype=torch.bool) key = self._encode_model_key(chunk.unsqueeze(0), chunk_mask.unsqueeze(0))[0] pending.append( { "text": f"context_chunk:{batch_index}:{start}:{end}", "key": key, "summary": key, "memory_type": "context_chunk", # Context chunks are ordered evidence, not competing # values of one fact. Leaving the fact-conflict # fields empty prevents ordinary version resolution # from superseding every earlier chunk in a document. "entity": "", "attribute": "", "value": "", "importance": 0.55, "confidence": 0.80, "source": "kv_compaction", "evidence": [ f"batch:{batch_index}", f"token_range:{start}:{end}", f"token_count:{end - start}", ], "token_ids": chunk, "token_mask": chunk_mask, "trusted": True, "force": True, } ) pending_previous.append(previous_index) previous_index = len(pending) - 1 results: list[tuple[Any, str]] = [] if pending and archive: results = self.memory_os_v2.write_batch(pending) # Add a forward chain after the atomic insert. This gives the # multi-hop reader a deterministic path through adjacent chunks # without making page allocation depend on future record IDs. for result_index, previous_index in enumerate(pending_previous): if previous_index is None: continue record = results[result_index][0] previous_record = results[previous_index][0] if previous_record.record_id not in record.related_ids: record.related_ids.append(previous_record.record_id) self.memory_os_v2.bank._store_record(record) max_retained = max(retained_lengths, default=0) compacted_ids = torch.zeros( (input_ids.shape[0], max_retained), dtype=input_ids.dtype, device=input_ids.device ) compacted_mask = torch.zeros( (input_ids.shape[0], max_retained), dtype=attention_mask.dtype, device=input_ids.device ) for batch_index, values in enumerate(retained): if values.numel() == 0: continue compacted_ids[batch_index, : values.numel()] = values.to(input_ids.device) compacted_mask[batch_index, : values.numel()] = 1 return compacted_ids, compacted_mask, { "compacted": any(original > retained for original, retained in zip(original_lengths, retained_lengths)), "archived_records": len(results), "archived_tokens": sum(original - retained for original, retained in zip(original_lengths, retained_lengths)), "original_tokens": original_lengths, "retained_tokens": retained_lengths, "hot_limit": hot_limit, "archive_requested": bool(archive), } @torch.no_grad() def _rank_v2_text_matches( self, query_key: Tensor, query_token_ids: Tensor, ) -> list[tuple[float, int, Any]]: """Rank active V2 text records without materializing their payloads.""" if self.memory_os_v2 is None: return [] active = [ record for record in self.memory_os_v2.bank.records.values() if record.status == "active" ] if not active: return [] keys = torch.stack([record.key for record in active], dim=0) scores = self.memory_os_v2.bank._score_candidates(query_key, keys).reshape(-1).cpu() query_unique = torch.unique(query_token_ids.detach().reshape(-1).cpu()) ranked: list[tuple[float, int, Any]] = [] for index, record in enumerate(active): shared = 0 if record.token_ids is not None and query_unique.numel() > 0: own = torch.unique(record.token_ids.detach().reshape(-1).cpu()) shared = int(torch.isin(query_unique, own).sum().item()) ranked.append((float(scores[index].item()), shared, record)) ranked.sort(key=lambda item: (item[0], item[1], item[2].last_access), reverse=True) return ranked @torch.no_grad() def _forget_text_memory_by_key( self, key_input_ids: Tensor, key_attention_mask: Tensor, ) -> Tensor: """Erase the addressed natural-language record for a learned forget turn.""" erased = torch.zeros( key_input_ids.shape[0], dtype=torch.bool, device=key_input_ids.device ) if ( self.runtime.text_slot_valid is None or self.runtime.text_key_token_ids is None or self.runtime.text_key_token_mask is None or not bool(self.runtime.text_slot_valid.any()) ): return erased query_keys = self._encode_model_key(key_input_ids, key_attention_mask) for batch_index in range(key_input_ids.shape[0]): valid_slots = self.runtime.text_slot_valid[batch_index] if not bool(valid_slots.any()): continue query_ids = key_input_ids[batch_index][key_attention_mask[batch_index].bool()] query_ids = query_ids[: self.memory_config.text_memory_key_tokens] old_ids = self.runtime.text_key_token_ids[batch_index] old_mask = self.runtime.text_key_token_mask[batch_index] shared = ( (old_ids[:, :, None] == query_ids[None, None, :]) & old_mask[:, :, None] & valid_slots[:, None, None] ).any(dim=-1).sum(dim=-1).float() old_lengths = old_mask.sum(dim=-1).float() overlap = shared / torch.sqrt( old_lengths * float(max(1, query_ids.numel())) ).clamp_min(1.0) overlap = overlap.masked_fill(~valid_slots, -1.0) best_overlap, best_slot = overlap.max(dim=-1) best_score = best_overlap if self.text_retriever is not None and self._text_retriever_ready: slot_keys = F.normalize( self.runtime.text_slot_keys[batch_index].to(query_keys.device, dtype=query_keys.dtype), dim=-1, ) learned = torch.sigmoid( self.text_retriever(query_keys[batch_index].unsqueeze(0), slot_keys.unsqueeze(0))[0] ).masked_fill(~valid_slots.to(query_keys.device), -1.0) learned_score, learned_slot = learned.max(dim=-1) if float(learned_score) >= float(best_score): best_score = learned_score best_slot = learned_slot if float(best_score) < self.memory_config.text_memory_overlap_threshold and not ( self.text_retriever is not None and self._text_retriever_ready and float(best_score) >= 0.65 ): continue slot = int(best_slot.item()) self._clear_legacy_memory_slot(slot) if self.memory_os_v2 is not None and self.runtime.use_persistent_state: ranked = self._rank_v2_text_matches( query_keys[batch_index], query_ids, ) for score, shared, record in ranked: # The hot-bank matcher has already located the target. # V2 records created as independent fragments have no # slot index, so use the frozen-Qwen address plus token # evidence to retract only the matching old record. same_hot_slot = record.slot_index == slot if same_hot_slot or (score >= 0.60 and shared >= 2): self.memory_os_v2.retract_record(record.record_id) if not same_hot_slot and score < 0.78: break erased[batch_index] = True return erased def _write_text_memory( self, input_ids: Optional[Tensor], attention_mask: Optional[Tensor], text_input_ids: Optional[Tensor] = None, text_attention_mask: Optional[Tensor] = None, key_input_ids: Optional[Tensor] = None, key_attention_mask: Optional[Tensor] = None, storage_input_ids: Optional[Tensor] = None, storage_attention_mask: Optional[Tensor] = None, force_write: bool = False, memory_text: Optional[str] = None, ) -> None: """Commit one conversational write into the selected internal slot. The learned controller decides whether a turn is a durable fact and which slot it belongs to. The exact token sequence is then copied into the model-owned bank so multi-token names, values and punctuation are not lossy-compressed into a single output token. """ # The bank is consumed as an internal prefix before the next Qwen # chat turn. Store the valid system-message form for generation; # keep the exact user fact separately in ``text_key_token_ids`` for # duplicate detection and retrieval overlap. A raw fact placed # before ``<|im_start|>user`` is not a reliable conditioning channel. storage_ids = ( text_input_ids if text_input_ids is not None else (storage_input_ids if storage_input_ids is not None else input_ids) ) storage_mask = ( text_attention_mask if text_input_ids is not None else (storage_attention_mask if storage_input_ids is not None else attention_mask) ) key_ids = key_input_ids if key_input_ids is not None else storage_ids key_mask = key_attention_mask if key_input_ids is not None else storage_mask if ( not self.memory_config.natural_language_memory or storage_ids is None or self.runtime.text_token_ids is None or self.runtime.text_slot_keys is None or self.runtime.text_key_token_ids is None ): return self.runtime.text_last_written_slot = torch.full( (storage_ids.shape[0],), -1, dtype=torch.long, device=self.runtime.text_slot_valid.device, ) address = self.memory.last_write_address write_probability = self.memory.last_write_probability if address is None or key_ids is None: return if write_probability is None: write_probability = address.max(dim=-1, keepdim=True).values if force_write: should_write = torch.ones( storage_ids.shape[0], dtype=torch.bool, device=storage_ids.device ) elif ( self.memory_config.automatic_memory and self._memory_policy_ready and self.runtime.auto_memory_probability is not None ): # High recall is intentional: the policy is trained to recognize # durable user information. It is the final automatic decision; # the older native write gate is not OR-ed here because it was # trained for continuous-state updates and can be over-eager on # question-shaped inputs. should_write = self.runtime.auto_memory_probability >= self.memory_config.auto_memory_threshold else: should_write = write_probability.squeeze(-1) >= self.memory_config.text_memory_write_threshold forget_probability = self.runtime.auto_memory_forget_probability if ( not force_write and forget_probability is not None and bool((forget_probability >= self.memory_config.auto_forget_threshold).any()) ): forget_mask = forget_probability >= self.memory_config.auto_forget_threshold forget_key_mask = ( key_mask if key_mask is not None else torch.ones_like(key_ids, dtype=torch.bool) ) erased = self._forget_text_memory_by_key(key_ids, forget_key_mask) should_write = should_write & ~(forget_mask & erased) if not force_write and memory_text: if looks_like_question(memory_text): should_write = torch.zeros_like(should_write, dtype=torch.bool) if storage_mask is None: storage_mask = torch.ones_like(storage_ids, dtype=torch.bool) else: storage_mask = storage_mask.to(device=storage_ids.device, dtype=torch.bool) if key_mask is None: key_mask = torch.ones_like(key_ids, dtype=torch.bool) else: key_mask = key_mask.to(device=key_ids.device, dtype=torch.bool) text_keys = self._encode_model_key(key_ids, key_mask) max_tokens = self.memory_config.text_memory_tokens for batch_index in range(storage_ids.shape[0]): if not bool(should_write[batch_index]): continue valid_ids = storage_ids[batch_index][storage_mask[batch_index]] if valid_ids.numel() == 0: continue valid_key_ids = key_ids[batch_index][key_mask[batch_index]] valid_key_ids = valid_key_ids[: self.memory_config.text_memory_key_tokens].detach().to( device=self.runtime.text_key_token_ids.device, dtype=torch.long, ) valid_ids = valid_ids[:max_tokens].detach().to( device=self.runtime.text_token_ids.device, dtype=torch.long, ) key = text_keys[batch_index] valid_slots = self.runtime.text_slot_valid[batch_index] if bool(valid_slots.any()): old_key_ids = self.runtime.text_key_token_ids[batch_index] old_key_mask = self.runtime.text_key_token_mask[batch_index] equal = old_key_ids[:, :, None] == valid_key_ids[None, None, :] overlap_count = ( equal & old_key_mask[:, :, None] & torch.ones( 1, 1, valid_key_ids.numel(), dtype=torch.bool, device=old_key_ids.device, ) ).any(dim=-1).sum(dim=-1).float() old_lengths = old_key_mask.sum(dim=-1).float() similarities = overlap_count / torch.sqrt( old_lengths * float(max(1, valid_key_ids.numel())) ).clamp_min(1.0) masked_similarities = similarities.masked_fill( ~valid_slots, torch.finfo(similarities.dtype).min, ) best_similarity, best_slot = masked_similarities.max(dim=-1) # Exact token equality is a stronger duplicate signal than # any learned semantic score. It makes repeated natural # language facts idempotent, even when the retriever has not # seen that exact name/value during training. key_length = valid_key_ids.numel() exact_match = ( (old_key_ids[:, :key_length] == valid_key_ids[None, :]) | (~old_key_mask[:, :key_length]) ).all(dim=-1) & valid_slots & ( old_key_mask.sum(dim=-1) == valid_key_ids.numel() ) exact_slots = exact_match.nonzero(as_tuple=False).flatten() if exact_slots.numel() > 0: best_similarity = torch.tensor(1.0, device=key.device) best_slot = exact_slots[0] if self.text_retriever is not None and self._text_retriever_ready: if exact_slots.numel() == 0: learned_similarity = torch.sigmoid( self.text_retriever( key.unsqueeze(0), self.runtime.text_slot_keys[batch_index].unsqueeze(0), )[0] ) learned_similarity = learned_similarity.masked_fill( ~valid_slots, torch.finfo(learned_similarity.dtype).min, ) learned_best, learned_slot = learned_similarity.max(dim=-1) # Shared words such as "我" and "请记住" are common in # many facts. Once the retriever is trained, use its # semantic same-attribute score for replacement so a new # fruit fact cannot evict a work-location fact merely # because both contain the same pronouns. if float(learned_best) >= 0.65: best_similarity = learned_best best_slot = learned_slot # Keep the exact token-overlap score when the learned # retriever is uncertain. Overwriting it with -1 # made obvious same-attribute updates look unrelated # merely because the paraphrase was outside the # retriever's training distribution. else: best_similarity = torch.tensor(-1.0, device=key.device) best_slot = torch.tensor(0, dtype=torch.long, device=key.device) if float(best_similarity) >= self.memory_config.text_memory_update_overlap_threshold: slot = int(best_slot.item()) else: free_slots = (~valid_slots).nonzero(as_tuple=False).flatten() if free_slots.numel() > 0: slot = int(free_slots[0].item()) else: slot = int(self.runtime.text_slot_age[batch_index].argmin().item()) self.runtime.text_token_ids[batch_index, slot].zero_() self.runtime.text_token_mask[batch_index, slot].zero_() self.runtime.text_token_ids[batch_index, slot, : valid_ids.numel()].copy_(valid_ids) self.runtime.text_token_mask[batch_index, slot, : valid_ids.numel()] = True self.runtime.text_slot_valid[batch_index, slot] = True self.runtime.text_slot_keys[batch_index, slot].copy_(key.to(self.runtime.text_slot_keys.device)) self.runtime.text_key_token_ids[batch_index, slot].zero_() self.runtime.text_key_token_mask[batch_index, slot].zero_() self.runtime.text_key_token_ids[batch_index, slot, : valid_key_ids.numel()].copy_(valid_key_ids) self.runtime.text_key_token_mask[batch_index, slot, : valid_key_ids.numel()] = True self.runtime.text_write_counter[batch_index] += 1 self.runtime.text_slot_age[batch_index, slot] = self.runtime.text_write_counter[batch_index] self.runtime.text_last_written_slot[batch_index] = slot if self.memory_os_v2 is not None and self.runtime.use_persistent_state: confidence = float(write_probability[batch_index].reshape(-1)[0].item()) if force_write: confidence = 1.0 # The 16-slot bank is only the hot cache. New, unrelated # conversational fragments must remain independent V2 # records instead of superseding whatever happened to be in # the same hot slot. Reuse a V2 slot index only when the # hot-bank matcher already established that this is a # duplicate/update of an existing fact. v2_slot_index = ( slot if float(best_similarity) >= self.memory_config.text_memory_update_overlap_threshold else -1 ) if v2_slot_index >= 0 and self.memory_os_v2 is not None: # A semantic replacement should retire the matching # active record, while unrelated fragments remain active # even if they once occupied the same hot-cache slot. for score, shared, record in self._rank_v2_text_matches( key, valid_key_ids, ): if record.slot_index == slot or (score >= 0.60 and shared >= 2): self.memory_os_v2.retract_record(record.record_id) if record.slot_index != slot and score < 0.78: break self.memory_os_v2.write( text=memory_text or f"memory_slot:{slot}", key=key, summary=key, memory_type="episodic_text", importance=max(confidence, 0.5), confidence=confidence, source="explicit" if force_write else "automatic", evidence=[f"token_count:{int(valid_ids.numel())}"], slot_index=v2_slot_index, token_ids=valid_ids, token_mask=torch.ones_like(valid_ids, dtype=torch.bool), trusted=force_write or confidence >= self.memory_config.memory_v2_write_threshold, force=force_write, ) def _probe_text_retrieval( self, input_ids: Tensor, attention_mask: Tensor, ) -> tuple[Optional[Tensor], Optional[Tensor]]: """Compute query-dependent similarity against internal text keys.""" if ( not self.memory_config.natural_language_memory or self.runtime.text_slot_valid is None or self.runtime.text_slot_keys is None or self.runtime.text_key_token_ids is None or self.runtime.text_key_token_mask is None or not bool(self.runtime.text_slot_valid.any()) ): return None, None query_key = self._encode_model_key(input_ids, attention_mask) self.runtime.v2_query_key = query_key.detach() slot_keys = F.normalize( self.runtime.text_slot_keys.to(device=query_key.device, dtype=query_key.dtype), dim=-1, ) if self.text_retriever is not None and self._text_retriever_ready: learned_logits = self.text_retriever(query_key, slot_keys) dense_similarity = torch.sigmoid(learned_logits) else: dense_similarity = torch.einsum("bh,bsh->bs", query_key, slot_keys) # Dense similarity alone is too permissive for short Chinese queries: # common tokens such as "我" and "的" can make unrelated memories # look relevant. Combine it with exact token overlap from the # internally stored key, yielding a conservative hybrid score. key_token_ids = self.runtime.text_key_token_ids.to(device=input_ids.device) key_token_mask = self.runtime.text_key_token_mask.to(device=input_ids.device, dtype=torch.bool) query_token_ids = input_ids.to(device=key_token_ids.device) query_token_mask = attention_mask.to(device=key_token_ids.device, dtype=torch.bool) equal = query_token_ids[:, None, :, None] == key_token_ids[:, :, None, :] equal &= query_token_mask[:, None, :, None] & key_token_mask[:, :, None, :] shared_tokens = equal.any(dim=-1).sum(dim=-1).float() query_lengths = query_token_mask.sum(dim=-1, keepdim=True).float() key_lengths = key_token_mask.sum(dim=-1).float() overlap = shared_tokens / torch.sqrt(query_lengths * key_lengths).clamp_min(1.0) if self.text_retriever is not None and self._text_retriever_ready: scores = 0.85 * dense_similarity + 0.15 * overlap.clamp(0.0, 1.0) else: dense_score = (dense_similarity + 1.0).clamp(0.0, 2.0) * 0.5 scores = 0.4 * dense_score + 0.6 * overlap.clamp(0.0, 1.0) self.runtime.text_read_overlap = overlap scores = scores.masked_fill( ~self.runtime.text_slot_valid.to(device=scores.device), torch.finfo(scores.dtype).min, ) return scores, scores.max(dim=-1).values def _build_text_prefix( self, input_ids: Tensor, attention_mask: Tensor, query_text: str = "", ) -> tuple[Optional[Tensor], Optional[Tensor], int]: """Retrieve text memory and build an internal prefix for Qwen.""" address, relevance = self._probe_text_retrieval(input_ids, attention_mask) self.runtime.text_read_slots = None self.runtime.v2_last_decisions = [] self.runtime.text_read_relevance = relevance self.runtime.text_prefix_used = False if ( self.memory_os_v2 is not None and self._memory_router_v2_ready and self.runtime.v2_query_key is None ): self.runtime.v2_query_key = self._encode_model_key(input_ids, attention_mask).detach() # V2 is the scalable address path. It first selects a bounded set of # pages through the coarse index, then reranks only records in those # pages. The legacy fixed bank remains a compatibility fallback for # old checkpoints that do not contain a trained V2 router. if ( self.memory_os_v2 is not None and self._memory_router_v2_ready and self.runtime.v2_query_key is not None ): v2_prefix_parts: list[Tensor] = [] v2_prefix_masks: list[Tensor] = [] v2_slots: list[list[int]] = [] for batch_index in range(input_ids.shape[0]): records, decision = self.memory_os_v2.read( query_key=self.runtime.v2_query_key[batch_index], query_text=query_text, query_token_ids=input_ids[batch_index], top_k_pages=self.memory_config.memory_top_k_pages, top_k_records=self.memory_config.memory_top_k_records, max_hops=self.memory_config.memory_max_hops, ) self.runtime.v2_last_decisions.append(asdict(decision)) ids_parts: list[Tensor] = [] slot_list: list[int] = [] for record in records: ids: Optional[Tensor] = None if record.token_ids is not None: cached_ids, cached_mask = self.memory_os_v2.bank.gpu_record_payload(record) if cached_ids is not None: if cached_mask is not None and cached_mask.numel() == cached_ids.numel(): ids = cached_ids[cached_mask] else: ids = cached_ids elif ( self.runtime.text_token_ids is not None and self.runtime.text_token_mask is not None and 0 <= record.slot_index < self.runtime.text_token_ids.shape[1] ): slot = record.slot_index slot_mask = self.runtime.text_token_mask[batch_index, slot] ids = self.runtime.text_token_ids[batch_index, slot][slot_mask].detach().cpu() if ids is not None and ids.numel() > 0: ids_parts.append(ids) if record.slot_index >= 0: slot_list.append(record.slot_index) if ids_parts: v2_prefix_parts.append(torch.cat(ids_parts, dim=0)) v2_prefix_masks.append(torch.ones_like(v2_prefix_parts[-1], dtype=torch.bool)) else: v2_prefix_parts.append(torch.zeros(0, dtype=torch.long, device=input_ids.device)) v2_prefix_masks.append(torch.zeros(0, dtype=torch.bool, device=input_ids.device)) v2_slots.append(slot_list) v2_prefix_length = max((part.numel() for part in v2_prefix_parts), default=0) if v2_prefix_length > 0: prefix_ids = torch.zeros( input_ids.shape[0], v2_prefix_length, dtype=input_ids.dtype, device=input_ids.device ) prefix_mask = torch.zeros( input_ids.shape[0], v2_prefix_length, dtype=attention_mask.dtype, device=input_ids.device ) for batch_index, (ids, mask) in enumerate(zip(v2_prefix_parts, v2_prefix_masks)): prefix_ids[batch_index, : ids.numel()] = ids.to(input_ids.device) prefix_mask[batch_index, : mask.numel()] = mask.to(attention_mask.dtype) self.runtime.text_read_slots = torch.tensor( [slots + [-1] * max(0, self.memory_config.memory_top_k_records - len(slots)) for slots in v2_slots], dtype=torch.long, device=input_ids.device, ) self.runtime.text_prefix_used = True return prefix_ids, prefix_mask, v2_prefix_length if address is None or self.runtime.text_slot_valid is None: return None, None, 0 valid = self.runtime.text_slot_valid scores = address.masked_fill(~valid, torch.finfo(address.dtype).min) top_k = min(self.memory_config.text_memory_top_k, scores.shape[-1]) top_scores, top_slots = scores.topk(top_k, dim=-1) selected = top_scores >= self.memory_config.text_memory_threshold # If exactly one stored memory exists, its relevance still has to pass # the threshold; this prevents unrelated questions from receiving an # arbitrary memory snippet and reduces hallucinated personal facts. if relevance is not None: selected &= relevance[:, None] >= self.memory_config.text_memory_threshold if ( self.runtime.text_read_overlap is not None and not (self.text_retriever is not None and self._text_retriever_ready) ): selected_overlap = self.runtime.text_read_overlap.gather(1, top_slots) selected &= selected_overlap >= self.memory_config.text_memory_overlap_threshold self.runtime.text_read_slots = top_slots if not bool(selected.any()): return None, None, 0 self.runtime.text_prefix_used = True bank_ids = self.runtime.text_token_ids bank_mask = self.runtime.text_token_mask if bank_ids is None or bank_mask is None: return None, None, 0 prefix_parts: list[Tensor] = [] prefix_masks: list[Tensor] = [] for batch_index in range(input_ids.shape[0]): ids_parts: list[Tensor] = [] mask_parts: list[Tensor] = [] for rank in range(top_k): if not bool(selected[batch_index, rank]): continue slot = int(top_slots[batch_index, rank].item()) slot_mask = bank_mask[batch_index, slot] ids_parts.append(bank_ids[batch_index, slot][slot_mask]) mask_parts.append(torch.ones_like(ids_parts[-1], dtype=torch.bool)) if ids_parts: prefix_parts.append(torch.cat(ids_parts, dim=0)) prefix_masks.append(torch.cat(mask_parts, dim=0)) else: prefix_parts.append(torch.zeros(0, dtype=torch.long, device=input_ids.device)) prefix_masks.append(torch.zeros(0, dtype=torch.bool, device=input_ids.device)) prefix_length = max((part.numel() for part in prefix_parts), default=0) if prefix_length == 0: return None, None, 0 prefix_ids = torch.zeros( input_ids.shape[0], prefix_length, dtype=input_ids.dtype, device=input_ids.device ) prefix_mask = torch.zeros( input_ids.shape[0], prefix_length, dtype=attention_mask.dtype, device=input_ids.device ) for batch_index, (ids, mask) in enumerate(zip(prefix_parts, prefix_masks)): prefix_ids[batch_index, : ids.numel()] = ids.to(input_ids.device) prefix_mask[batch_index, : mask.numel()] = mask.to(attention_mask.dtype) return prefix_ids, prefix_mask, prefix_length def _find_layer_device(self) -> torch.device: for index in self.layer_indices: layer = self.base_model.model.language_model.layers[index] # In blend mode the adapter owns a trainable scalar. Inspect the # original layer first so a CPU-created scalar cannot mislead the # device choice for input_ids and the memory module. layer_for_device = getattr(layer, "inner", layer) for parameter in layer_for_device.parameters(): if parameter.device.type != "meta": return parameter.device return self.base_model.get_input_embeddings().weight.device @property def trainable_parameters(self): parameters = [parameter for parameter in self.memory.parameters() if parameter.requires_grad] if self.text_retriever is not None: parameters.extend( parameter for parameter in self.text_retriever.parameters() if parameter.requires_grad ) if self.memory_policy is not None: parameters.extend( parameter for parameter in self.memory_policy.parameters() if parameter.requires_grad ) if self.memory_router_v2 is not None: parameters.extend( parameter for parameter in self.memory_router_v2.parameters() if parameter.requires_grad ) for adapter in self._memory_adapters: blend_logit = getattr(adapter, "blend_logit", None) if blend_logit is not None and blend_logit.requires_grad: parameters.append(blend_logit) return iter(parameters) @torch.no_grad() def read_hierarchical_memory( self, query_key: Tensor, *, query_text: str = "", query_token_ids: Optional[Tensor] = None, top_k_pages: Optional[int] = None, top_k_records: Optional[int] = None, max_hops: Optional[int] = None, ) -> tuple[list[Any], Any]: """Read V2 memory through its bounded page and record router.""" if self.memory_os_v2 is None: raise RuntimeError("hierarchical memory is disabled in memory_config") return self.memory_os_v2.read( query_key=query_key, query_text=query_text, query_token_ids=query_token_ids, top_k_pages=top_k_pages, top_k_records=top_k_records, max_hops=max_hops, ) @torch.no_grad() def write_hierarchical_memory(self, **kwargs: Any) -> tuple[Any, str]: """Write one versioned V2 record owned by the model checkpoint.""" if self.memory_os_v2 is None: raise RuntimeError("hierarchical memory is disabled in memory_config") return self.memory_os_v2.write(**kwargs) def memory_v2_stats(self) -> dict[str, Any]: if self.memory_os_v2 is None: return {"enabled": False} output = dict(self.memory_os_v2.stats()) output["enabled"] = True output["router_ready"] = self._memory_router_v2_ready return output def list_memory_records( self, *, query_text: str = "", status: str = "active", limit: int = 100, offset: int = 0, ) -> list[dict[str, Any]]: """Return a safe, tensor-free view of model-owned memory records.""" if self.memory_os_v2 is None: return [] return [ memory_record_to_dict(record) for record in self.memory_os_v2.list_records( query_text=query_text, status=status, limit=limit, offset=offset, ) ] def get_memory_record(self, record_id: str) -> dict[str, Any]: if self.memory_os_v2 is None: raise KeyError(record_id) record = self.memory_os_v2.bank._resolve_record(record_id) return memory_record_to_dict(record) def _clear_legacy_memory_slot(self, slot_index: int) -> None: """Remove a V1 compatibility copy when a V2 record is edited/retracted.""" if slot_index < 0: return fields = ( ("persistent_text_token_ids", "text_token_ids", 0), ("persistent_text_token_mask", "text_token_mask", 0), ("persistent_text_slot_valid", "text_slot_valid", 0), ("persistent_text_slot_keys", "text_slot_keys", 0), ("persistent_text_slot_age", "text_slot_age", -1), ("persistent_text_key_token_ids", "text_key_token_ids", 0), ("persistent_text_key_token_mask", "text_key_token_mask", 0), ) for persistent_name, runtime_name, fill_value in fields: persistent = getattr(self, persistent_name, None) if isinstance(persistent, Tensor) and persistent.ndim >= 2 and slot_index < persistent.shape[1]: persistent[:, slot_index].fill_(fill_value) runtime = getattr(self.runtime, runtime_name, None) if isinstance(runtime, Tensor) and runtime.ndim >= 2 and slot_index < runtime.shape[1]: runtime[:, slot_index].fill_(fill_value) @torch.no_grad() def edit_memory_record( self, record_id: str, *, text: Optional[str] = None, entity: Optional[str] = None, attribute: Optional[str] = None, value: Optional[str] = None, importance: Optional[float] = None, confidence: Optional[float] = None, evidence: Optional[list[str]] = None, token_ids: Optional[Tensor] = None, token_mask: Optional[Tensor] = None, ) -> dict[str, Any]: """Edit memory by creating a versioned successor, never overwriting history.""" if self.memory_os_v2 is None: raise KeyError(record_id) old = self.memory_os_v2.bank._resolve_record(record_id) key = None if token_ids is not None: ids = token_ids.detach().reshape(1, -1).to(self._find_layer_device()) mask = ( token_mask.detach().reshape(1, -1).to(self._find_layer_device(), dtype=torch.long) if token_mask is not None else torch.ones_like(ids) ) key = self._encode_model_key(ids, mask)[0] self._clear_legacy_memory_slot(old.slot_index) record = self.memory_os_v2.edit_record( record_id, text=text, key=key, summary=key, entity=entity, attribute=attribute, value=value, importance=importance, confidence=confidence, evidence=evidence, token_ids=token_ids, token_mask=token_mask, ) return memory_record_to_dict(record) @torch.no_grad() def retract_memory_record(self, record_id: str) -> dict[str, Any]: """Retract a record and its compatibility-slot copy without erasing history.""" if self.memory_os_v2 is None: raise KeyError(record_id) record = self.memory_os_v2.bank._resolve_record(record_id) self._clear_legacy_memory_slot(record.slot_index) self.memory_os_v2.retract_record(record_id) return memory_record_to_dict(record) def audit_memory(self) -> dict[str, Any]: if self.memory_os_v2 is None: return {"healthy": True, "issues": [], "issue_count": 0, "stats": {"enabled": False}} return self.memory_os_v2.audit() def export_memory_records(self, *, limit: int = 10000, offset: int = 0) -> dict[str, Any]: records = self.list_memory_records(status="all", limit=limit, offset=offset) return { "format_version": 1, "records": records, "offset": offset, "limit": limit, "returned": len(records), "stats": self.memory_v2_stats(), } def flush_memory_storage(self) -> None: """Flush the optional warm/cold tier without rewriting model weights.""" if self.memory_os_v2 is not None: self.memory_os_v2.flush_storage() def close_memory_storage(self) -> None: """Close the optional durable page store before process shutdown.""" if self.memory_os_v2 is not None: self.memory_os_v2.close_storage() def clear_hierarchical_memory(self) -> None: """Clear the durable V2 page store while keeping router weights.""" if self.memory_router_v2 is None: return if self.memory_os_v2 is not None and self.memory_os_v2.bank.tier_store is not None: old_store = self.memory_os_v2.bank.tier_store old_store.clear() old_store.close() self.memory_os_v2 = self._new_memory_os_v2(self.memory.hidden_size) def reset_memory(self, batch_size: Optional[int] = None, *, device: Optional[torch.device] = None) -> None: """Explicitly clear the durable memory state. Use :meth:`reset_runtime_memory` for a temporary conversation or an evaluation case. Keeping these operations separate prevents a test or a branch from silently deleting a user's embedded checkpoint. """ self.runtime.use_persistent_state = True self.clear_hierarchical_memory() if self.memory_config.persistent_memory: device = device or self._find_layer_device() batch_size = batch_size or 1 self.persistent_memory = self.memory.initial_state(batch_size, device=device) self._persistent_memory = self.persistent_memory.detach() if self.memory_config.natural_language_memory: self._ensure_text_memory(batch_size, device=device) self._clear_text_memory() self.runtime.state = self._persistent_memory self.runtime.last_read = None self.runtime.auto_memory_probability = None self.runtime.raw_memory = None self.runtime.reset_mask = None self.runtime.text_token_ids = self.persistent_text_token_ids self.runtime.text_token_mask = self.persistent_text_token_mask self.runtime.text_slot_valid = self.persistent_text_slot_valid self.runtime.text_slot_keys = self.persistent_text_slot_keys self.runtime.text_slot_age = self.persistent_text_slot_age self.runtime.text_write_counter = self.persistent_text_write_counter self.runtime.text_key_token_ids = self.persistent_text_key_token_ids self.runtime.text_key_token_mask = self.persistent_text_key_token_mask self.runtime.text_last_written_slot = None return if batch_size is None: self._persistent_memory = None self.runtime.state = None self.runtime.last_read = None self.runtime.auto_memory_probability = None self.runtime.raw_memory = None self.runtime.text_token_ids = None self.runtime.text_token_mask = None self.runtime.text_slot_valid = None self.runtime.text_slot_keys = None self.runtime.text_slot_age = None self.runtime.text_write_counter = None self.runtime.text_key_token_ids = None self.runtime.text_key_token_mask = None self.runtime.text_last_written_slot = None return device = device or self._find_layer_device() self._persistent_memory = self.memory.initial_state(batch_size, device=device) if self.memory_config.natural_language_memory: self._ensure_text_memory(batch_size, device=device) self._clear_text_memory() self.runtime.state = self._persistent_memory self.runtime.last_read = None self.runtime.auto_memory_probability = None self.runtime.raw_memory = None self.runtime.text_token_ids = self.persistent_text_token_ids self.runtime.text_token_mask = self.persistent_text_token_mask self.runtime.text_slot_valid = self.persistent_text_slot_valid self.runtime.text_slot_keys = self.persistent_text_slot_keys self.runtime.text_slot_age = self.persistent_text_slot_age self.runtime.text_write_counter = self.persistent_text_write_counter self.runtime.text_key_token_ids = self.persistent_text_key_token_ids self.runtime.text_key_token_mask = self.persistent_text_key_token_mask self.runtime.text_last_written_slot = None def reset_runtime_memory( self, batch_size: Optional[int] = None, *, device: Optional[torch.device] = None, ) -> None: """Reset only the active conversation, preserving durable state. In persistent mode this creates an ephemeral empty bank and prevents the generation path from rebinding the embedded user checkpoint. It is used by evaluation, memory branches and temporary sessions. """ device = device or self._find_layer_device() batch_size = batch_size or 1 self.runtime.use_persistent_state = False self.runtime.state = self.memory.initial_state(batch_size, device=device) self._persistent_memory = self.runtime.state.detach() self.runtime.last_read = None self.runtime.auto_memory_probability = None self.runtime.raw_memory = None self.runtime.v2_query_key = None self.runtime.v2_last_decisions = [] self.runtime.reset_mask = None self.runtime.text_last_written_slot = None if self.memory_config.natural_language_memory: self._bind_text_memory(batch_size, device=device) self._clear_text_memory() def _reset_mask(self, input_ids: Optional[Tensor]) -> Optional[Tensor]: token_id = self.memory_config.reset_token_id if token_id is None or input_ids is None: return None mask = (input_ids == token_id).any(dim=-1) return mask if bool(mask.any()) else None def save_runtime_memory( self, path: str | Path, memory_state: Optional[Tensor] = None, ) -> Path: """Persist one user's runtime memory, without saving model weights.""" state = memory_state if memory_state is not None else self.runtime.state if state is None: state = self._persistent_memory if state is None: raise ValueError("no runtime memory state is available to save") if state.ndim != 3 or tuple(state.shape[1:]) != ( self.memory_config.memory_slots, self.memory_config.memory_dim, ): raise ValueError(f"unexpected runtime memory shape: {tuple(state.shape)}") raw_memory = self.runtime.raw_memory if raw_memory is not None: raw_memory = raw_memory.detach().cpu() payload = { "format_version": 1, "memory_state": state.detach().cpu(), "raw_memory": raw_memory, "model_hidden_size": self.memory.hidden_size, "memory_config": asdict(self.memory_config), "layer_indices": list(self.layer_indices), } if self.memory_config.natural_language_memory and self.runtime.text_token_ids is not None: payload.update( { "text_token_ids": self.runtime.text_token_ids.detach().cpu(), "text_token_mask": self.runtime.text_token_mask.detach().cpu(), "text_slot_valid": self.runtime.text_slot_valid.detach().cpu(), "text_slot_keys": self.runtime.text_slot_keys.detach().cpu(), "text_slot_age": self.runtime.text_slot_age.detach().cpu(), "text_write_counter": self.runtime.text_write_counter.detach().cpu(), "text_key_token_ids": self.runtime.text_key_token_ids.detach().cpu(), "text_key_token_mask": self.runtime.text_key_token_mask.detach().cpu(), } ) output_path = Path(path) output_path.parent.mkdir(parents=True, exist_ok=True) torch.save(payload, output_path) return output_path def load_runtime_memory( self, path: str | Path, *, device: Optional[torch.device] = None, ) -> Tensor: """Load a user's runtime memory into this freshly created model.""" payload = torch.load(Path(path), map_location="cpu", weights_only=True) if isinstance(payload, Tensor): state = payload raw_memory = None elif isinstance(payload, dict): state = payload.get("memory_state", payload.get("state")) raw_memory = payload.get("raw_memory") else: state = None raw_memory = None if not isinstance(state, Tensor): raise ValueError("runtime memory file does not contain a tensor memory_state") expected_shape = ( self.memory_config.memory_slots, self.memory_config.memory_dim, ) if state.ndim != 3 or tuple(state.shape[1:]) != expected_shape: raise ValueError( f"runtime memory shape {tuple(state.shape)} does not match {expected_shape}" ) saved_hidden_size = payload.get("model_hidden_size") if isinstance(payload, dict) else None if saved_hidden_size is not None and int(saved_hidden_size) != self.memory.hidden_size: raise ValueError( f"runtime memory hidden size {saved_hidden_size} does not match {self.memory.hidden_size}" ) target_device = device or self._find_layer_device() state = state.to(device=target_device, dtype=self.memory.slot_keys.dtype) self._persistent_memory = state.detach() if self.memory_config.persistent_memory: self.persistent_memory = self._persistent_memory self.runtime.state = self._persistent_memory self.runtime.last_read = None self.runtime.raw_memory = None if self.memory_config.natural_language_memory: token_ids = payload.get("text_token_ids") if isinstance(payload, dict) else None token_mask = payload.get("text_token_mask") if isinstance(payload, dict) else None slot_valid = payload.get("text_slot_valid") if isinstance(payload, dict) else None slot_keys = payload.get("text_slot_keys") if isinstance(payload, dict) else None slot_age = payload.get("text_slot_age") if isinstance(payload, dict) else None write_counter = payload.get("text_write_counter") if isinstance(payload, dict) else None key_token_ids = payload.get("text_key_token_ids") if isinstance(payload, dict) else None key_token_mask = payload.get("text_key_token_mask") if isinstance(payload, dict) else None if not all( isinstance(item, Tensor) for item in ( token_ids, token_mask, slot_valid, slot_keys, slot_age, write_counter, key_token_ids, key_token_mask, ) ): raise ValueError("natural-language memory file is missing its indexed text bank") expected_text_shape = ( state.shape[0], self.memory_config.memory_slots, self.memory_config.text_memory_tokens, ) if tuple(token_ids.shape) != expected_text_shape or tuple(token_mask.shape) != expected_text_shape: raise ValueError("natural-language text bank shape does not match memory config") if tuple(slot_valid.shape) != expected_text_shape[:2]: raise ValueError("natural-language slot validity shape does not match memory config") if tuple(slot_keys.shape) != (state.shape[0], self.memory_config.memory_slots, self.memory.hidden_size): raise ValueError("natural-language slot key shape does not match memory config") if tuple(slot_age.shape) != expected_text_shape[:2] or tuple(write_counter.shape) != (state.shape[0],): raise ValueError("natural-language slot age shape does not match memory config") expected_key_shape = ( state.shape[0], self.memory_config.memory_slots, self.memory_config.text_memory_key_tokens, ) if tuple(key_token_ids.shape) != expected_key_shape or tuple(key_token_mask.shape) != expected_key_shape: raise ValueError("natural-language key token shape does not match memory config") self.persistent_text_token_ids = token_ids.to(device=target_device, dtype=torch.long) self.persistent_text_token_mask = token_mask.to(device=target_device, dtype=torch.bool) self.persistent_text_slot_valid = slot_valid.to(device=target_device, dtype=torch.bool) self.persistent_text_slot_keys = slot_keys.to(device=target_device, dtype=torch.float32) self.persistent_text_slot_age = slot_age.to(device=target_device, dtype=torch.long) self.persistent_text_write_counter = write_counter.to(device=target_device, dtype=torch.long) self.persistent_text_key_token_ids = key_token_ids.to(device=target_device, dtype=torch.long) self.persistent_text_key_token_mask = key_token_mask.to(device=target_device, dtype=torch.bool) self.runtime.text_token_ids = self.persistent_text_token_ids self.runtime.text_token_mask = self.persistent_text_token_mask self.runtime.text_slot_valid = self.persistent_text_slot_valid self.runtime.text_slot_keys = self.persistent_text_slot_keys self.runtime.text_slot_age = self.persistent_text_slot_age self.runtime.text_write_counter = self.persistent_text_write_counter self.runtime.text_key_token_ids = self.persistent_text_key_token_ids self.runtime.text_key_token_mask = self.persistent_text_key_token_mask if raw_memory is not None: if not isinstance(raw_memory, Tensor): raise ValueError("runtime memory raw_memory must be a tensor or null") if raw_memory.ndim == 1: raw_memory = raw_memory.unsqueeze(0) if tuple(raw_memory.shape) != (state.shape[0], self.memory.hidden_size): raise ValueError(f"unexpected raw_memory shape: {tuple(raw_memory.shape)}") output_weight = self.base_model.get_output_embeddings().weight self.runtime.raw_memory = raw_memory.to( device=output_weight.device, dtype=output_weight.dtype, ).detach() return self._persistent_memory def save_memory_adapter(self, output_dir: str | Path) -> None: path = Path(output_dir) path.mkdir(parents=True, exist_ok=True) torch.save(self.memory.state_dict(), path / "memory.pt") if self.text_retriever is not None and self._text_retriever_ready: torch.save(self.text_retriever.state_dict(), path / "text_retriever.pt") if self.memory_policy is not None and self._memory_policy_ready: torch.save(self.memory_policy.state_dict(), path / "memory_policy.pt") metadata = { "hidden_size": self.memory.hidden_size, "memory_config": asdict(self.memory_config), "layer_indices": list(self.layer_indices), "text_retriever_ready": self._text_retriever_ready, "memory_policy_ready": self._memory_policy_ready, } (path / "memory_config.json").write_text(json.dumps(metadata, indent=2), encoding="utf-8") surgery_state = { str(index): adapter.blend_logit.detach().cpu() for index, adapter in zip(self.layer_indices, self._memory_adapters) if hasattr(adapter, "blend_logit") } torch.save({"blend_logits": surgery_state}, path / "surgery.pt") if self.memory_config.persistent_memory and self.persistent_memory.numel() > 0: payload = { "format_version": 1, "memory_state": self.persistent_memory.detach().cpu(), } if self.memory_config.natural_language_memory: payload.update( { "text_token_ids": self.persistent_text_token_ids.detach().cpu(), "text_token_mask": self.persistent_text_token_mask.detach().cpu(), "text_slot_valid": self.persistent_text_slot_valid.detach().cpu(), "text_slot_keys": self.persistent_text_slot_keys.detach().cpu(), "text_slot_age": self.persistent_text_slot_age.detach().cpu(), "text_write_counter": self.persistent_text_write_counter.detach().cpu(), "text_key_token_ids": self.persistent_text_key_token_ids.detach().cpu(), "text_key_token_mask": self.persistent_text_key_token_mask.detach().cpu(), } ) torch.save( payload, path / "persistent_memory.pt", ) def save_persistent_memory_checkpoint(self, output_dir: str | Path) -> None: """Save controller weights and the current user memory together. This produces a compact model-owned memory checkpoint. The large frozen Qwen shards remain untouched; the learned reader/controller and the user's current memory are stored in the adapter package. """ state = self.runtime.state if self.runtime.state is not None else self._persistent_memory if state is None: raise ValueError("no runtime memory state is available to checkpoint") was_persistent = self.memory_config.persistent_memory self.memory_config.persistent_memory = True self.persistent_memory = state.detach().clone() self._persistent_memory = self.persistent_memory.detach() if self.memory_config.natural_language_memory: text_fields = ( self.runtime.text_token_ids, self.runtime.text_token_mask, self.runtime.text_slot_valid, self.runtime.text_slot_keys, self.runtime.text_slot_age, self.runtime.text_write_counter, self.runtime.text_key_token_ids, self.runtime.text_key_token_mask, ) if not all(isinstance(item, Tensor) for item in text_fields): raise ValueError("natural-language memory bank is not initialized") target_device = self._find_layer_device() self.persistent_text_token_ids = text_fields[0].detach().clone().to(target_device, dtype=torch.long) self.persistent_text_token_mask = text_fields[1].detach().clone().to(target_device, dtype=torch.bool) self.persistent_text_slot_valid = text_fields[2].detach().clone().to(target_device, dtype=torch.bool) self.persistent_text_slot_keys = text_fields[3].detach().clone().to(target_device, dtype=torch.float32) self.persistent_text_slot_age = text_fields[4].detach().clone().to(target_device, dtype=torch.long) self.persistent_text_write_counter = text_fields[5].detach().clone().to(target_device, dtype=torch.long) self.persistent_text_key_token_ids = text_fields[6].detach().clone().to(target_device, dtype=torch.long) self.persistent_text_key_token_mask = text_fields[7].detach().clone().to(target_device, dtype=torch.bool) try: self.save_memory_adapter(output_dir) path = Path(output_dir) metadata_path = path / "memory_config.json" metadata = json.loads(metadata_path.read_text(encoding="utf-8")) metadata["memory_config"]["persistent_memory"] = True metadata["memory_config"]["checkpoint_contains_user_memory"] = True metadata_path.write_text(json.dumps(metadata, indent=2), encoding="utf-8") finally: self.memory_config.persistent_memory = was_persistent def load_memory_adapter(self, adapter_dir: str | Path, *, strict: bool = True) -> None: path = Path(adapter_dir) state = torch.load(path / "memory.pt", map_location=self._find_layer_device(), weights_only=True) self.memory.load_state_dict(state, strict=strict) metadata_ready = None policy_ready = None metadata_path = path / "memory_config.json" if metadata_path.exists(): metadata = json.loads(metadata_path.read_text(encoding="utf-8")) if isinstance(metadata, dict) and "text_retriever_ready" in metadata: metadata_ready = bool(metadata["text_retriever_ready"]) if isinstance(metadata, dict) and "memory_policy_ready" in metadata: policy_ready = bool(metadata["memory_policy_ready"]) if metadata_ready and self.text_retriever is None: raise ValueError( "adapter contains a trained text retriever but natural_language_memory is disabled" ) retriever_path = path / "text_retriever.pt" if self.text_retriever is not None and metadata_ready and not retriever_path.exists(): raise FileNotFoundError( f"adapter declares a trained text retriever but is missing {retriever_path}" ) should_load_retriever = ( self.text_retriever is not None and retriever_path.exists() and metadata_ready is not False ) if should_load_retriever: retriever_state = torch.load( retriever_path, map_location=self._find_layer_device(), weights_only=True, ) self.text_retriever.load_state_dict(retriever_state, strict=True) self._text_retriever_ready = True policy_path = path / "memory_policy.pt" if policy_ready and self.memory_policy is None: raise ValueError( "adapter contains an automatic memory policy but automatic_memory is disabled" ) if self.memory_policy is not None and policy_ready is not None: if policy_ready and not policy_path.exists(): raise FileNotFoundError( f"adapter declares an automatic memory policy but is missing {policy_path}" ) if policy_ready: policy_state = torch.load( policy_path, map_location=self._find_layer_device(), weights_only=True, ) self.memory_policy.load_state_dict(policy_state, strict=strict) self._memory_policy_ready = True surgery_path = path / "surgery.pt" if surgery_path.exists(): surgery_state = torch.load(surgery_path, map_location="cpu", weights_only=True) blend_logits = surgery_state.get("blend_logits", {}) for index, adapter in zip(self.layer_indices, self._memory_adapters): if hasattr(adapter, "blend_logit") and str(index) in blend_logits: adapter.blend_logit.data.copy_(blend_logits[str(index)].to(adapter.blend_logit.device)) persistent_path = path / "persistent_memory.pt" if persistent_path.exists(): payload = torch.load(persistent_path, map_location="cpu", weights_only=True) state = payload.get("memory_state", payload.get("state")) if isinstance(payload, dict) else payload if not isinstance(state, Tensor): raise ValueError("persistent_memory.pt does not contain a tensor memory_state") expected_shape = (self.memory_config.memory_slots, self.memory_config.memory_dim) if state.ndim != 3 or tuple(state.shape[1:]) != expected_shape: raise ValueError( f"persistent memory shape {tuple(state.shape)} does not match {expected_shape}" ) target_device = self._find_layer_device() self.memory_config.persistent_memory = True self.persistent_memory = state.to( device=target_device, dtype=self.memory.slot_keys.dtype, ) self._persistent_memory = self.persistent_memory.detach() self.runtime.state = self._persistent_memory if self.memory_config.natural_language_memory: token_ids = payload.get("text_token_ids") if isinstance(payload, dict) else None token_mask = payload.get("text_token_mask") if isinstance(payload, dict) else None slot_valid = payload.get("text_slot_valid") if isinstance(payload, dict) else None slot_keys = payload.get("text_slot_keys") if isinstance(payload, dict) else None slot_age = payload.get("text_slot_age") if isinstance(payload, dict) else None write_counter = payload.get("text_write_counter") if isinstance(payload, dict) else None key_token_ids = payload.get("text_key_token_ids") if isinstance(payload, dict) else None key_token_mask = payload.get("text_key_token_mask") if isinstance(payload, dict) else None if not all( isinstance(item, Tensor) for item in ( token_ids, token_mask, slot_valid, slot_keys, slot_age, write_counter, key_token_ids, key_token_mask, ) ): raise ValueError("persistent adapter is missing its indexed natural-language bank") expected_text_shape = ( state.shape[0], self.memory_config.memory_slots, self.memory_config.text_memory_tokens, ) if tuple(token_ids.shape) != expected_text_shape or tuple(token_mask.shape) != expected_text_shape: raise ValueError("persistent text bank shape does not match memory config") if tuple(slot_valid.shape) != expected_text_shape[:2]: raise ValueError("persistent slot validity shape does not match memory config") if tuple(slot_keys.shape) != (state.shape[0], self.memory_config.memory_slots, self.memory.hidden_size): raise ValueError("persistent slot key shape does not match memory config") if tuple(slot_age.shape) != expected_text_shape[:2] or tuple(write_counter.shape) != (state.shape[0],): raise ValueError("persistent slot age shape does not match memory config") expected_key_shape = ( state.shape[0], self.memory_config.memory_slots, self.memory_config.text_memory_key_tokens, ) if tuple(key_token_ids.shape) != expected_key_shape or tuple(key_token_mask.shape) != expected_key_shape: raise ValueError("persistent key token shape does not match memory config") self.persistent_text_token_ids = token_ids.to(device=target_device, dtype=torch.long) self.persistent_text_token_mask = token_mask.to(device=target_device, dtype=torch.bool) self.persistent_text_slot_valid = slot_valid.to(device=target_device, dtype=torch.bool) self.persistent_text_slot_keys = slot_keys.to(device=target_device, dtype=torch.float32) self.persistent_text_slot_age = slot_age.to(device=target_device, dtype=torch.long) self.persistent_text_write_counter = write_counter.to(device=target_device, dtype=torch.long) self.persistent_text_key_token_ids = key_token_ids.to(device=target_device, dtype=torch.long) self.persistent_text_key_token_mask = key_token_mask.to(device=target_device, dtype=torch.bool) self.runtime.text_token_ids = self.persistent_text_token_ids self.runtime.text_token_mask = self.persistent_text_token_mask self.runtime.text_slot_valid = self.persistent_text_slot_valid self.runtime.text_slot_keys = self.persistent_text_slot_keys self.runtime.text_slot_age = self.persistent_text_slot_age self.runtime.text_write_counter = self.persistent_text_write_counter self.runtime.text_key_token_ids = self.persistent_text_key_token_ids self.runtime.text_key_token_mask = self.persistent_text_key_token_mask @torch.no_grad() def _load_persistent_memory_payload(self, payload: dict[str, Any]) -> None: """Load a runtime/persistent memory payload shared by .pt and safetensors.""" state = payload.get("memory_state", payload.get("state")) if not isinstance(state, Tensor): raise ValueError("persistent memory payload does not contain a tensor memory_state") expected_shape = (self.memory_config.memory_slots, self.memory_config.memory_dim) if state.ndim != 3 or tuple(state.shape[1:]) != expected_shape: raise ValueError( f"persistent memory shape {tuple(state.shape)} does not match {expected_shape}" ) target_device = self._find_layer_device() self.memory_config.persistent_memory = True self.runtime.use_persistent_state = True self.persistent_memory = state.to(device=target_device, dtype=self.memory.slot_keys.dtype) self._persistent_memory = self.persistent_memory.detach() self.runtime.state = self._persistent_memory if not self.memory_config.natural_language_memory: return fields = { "persistent_text_token_ids": ("text_token_ids", torch.long), "persistent_text_token_mask": ("text_token_mask", torch.bool), "persistent_text_slot_valid": ("text_slot_valid", torch.bool), "persistent_text_slot_keys": ("text_slot_keys", torch.float32), "persistent_text_slot_age": ("text_slot_age", torch.long), "persistent_text_write_counter": ("text_write_counter", torch.long), "persistent_text_key_token_ids": ("text_key_token_ids", torch.long), "persistent_text_key_token_mask": ("text_key_token_mask", torch.bool), } if not all(isinstance(payload.get(key), Tensor) for key, _ in fields.values()): raise ValueError("persistent memory payload is missing its indexed natural-language bank") expected_text_shape = ( state.shape[0], self.memory_config.memory_slots, self.memory_config.text_memory_tokens, ) expected_key_shape = ( state.shape[0], self.memory_config.memory_slots, self.memory_config.text_memory_key_tokens, ) expected_shapes = { "text_token_ids": expected_text_shape, "text_token_mask": expected_text_shape, "text_slot_valid": expected_text_shape[:2], "text_slot_keys": (state.shape[0], self.memory_config.memory_slots, self.memory.hidden_size), "text_slot_age": expected_text_shape[:2], "text_write_counter": (state.shape[0],), "text_key_token_ids": expected_key_shape, "text_key_token_mask": expected_key_shape, } for key, _ in fields.values(): if tuple(payload[key].shape) != expected_shapes[key]: raise ValueError(f"persistent field {key} shape does not match memory config") for attribute, (key, dtype) in fields.items(): setattr(self, attribute, payload[key].to(device=target_device, dtype=dtype)) runtime_fields = { "text_token_ids": self.persistent_text_token_ids, "text_token_mask": self.persistent_text_token_mask, "text_slot_valid": self.persistent_text_slot_valid, "text_slot_keys": self.persistent_text_slot_keys, "text_slot_age": self.persistent_text_slot_age, "text_write_counter": self.persistent_text_write_counter, "text_key_token_ids": self.persistent_text_key_token_ids, "text_key_token_mask": self.persistent_text_key_token_mask, } for attribute, value in runtime_fields.items(): setattr(self.runtime, attribute, value) def _export_memory_os_v2_checkpoint(self) -> tuple[dict[str, Tensor], dict[str, Any]]: """Split a V2 bank into tensors plus JSON-safe checkpoint metadata.""" if self.memory_os_v2 is None: return {}, {} payload = self.memory_os_v2.export_payload() tensors: dict[str, Tensor] = {} metadata = dict(payload) def pack_record(item: dict[str, Any], prefix: str) -> dict[str, Any]: item = dict(item) for field in ("key", "summary", "token_ids", "token_mask"): value = item.pop(field, None) if isinstance(value, Tensor): name = f"dynamic_memory.v2.{prefix}.{field}" tensors[name] = value.detach().cpu().contiguous() item[f"{field}_ref"] = name return item metadata["records"] = [ pack_record(item, f"records.{index}") for index, item in enumerate(payload.get("records", [])) ] metadata["quarantine"] = [ pack_record(item, f"quarantine.{index}") for index, item in enumerate(payload.get("quarantine", [])) ] metadata["pages"] = [] for index, item in enumerate(payload.get("pages", [])): item = dict(item) for field in ("key", "summary"): value = item.pop(field, None) if isinstance(value, Tensor): name = f"dynamic_memory.v2.pages.{index}.{field}" tensors[name] = value.detach().cpu().contiguous() item[f"{field}_ref"] = name metadata["pages"].append(item) return tensors, metadata def _load_memory_os_v2_checkpoint( self, metadata_text: str, weights: dict[str, Tensor], ) -> None: """Rehydrate the V2 page store from safetensors metadata and tensors.""" if self.memory_router_v2 is None: return payload = json.loads(metadata_text) def unpack_record(item: dict[str, Any]) -> dict[str, Any]: item = dict(item) for field in ("key", "summary", "token_ids", "token_mask"): ref = item.pop(f"{field}_ref", None) item[field] = weights.get(ref) if ref else None return item payload["records"] = [unpack_record(item) for item in payload.get("records", [])] payload["quarantine"] = [unpack_record(item) for item in payload.get("quarantine", [])] unpacked_pages = [] for item in payload.get("pages", []): item = dict(item) for field in ("key", "summary"): ref = item.pop(f"{field}_ref", None) item[field] = weights.get(ref) if ref else None unpacked_pages.append(item) payload["pages"] = unpacked_pages tier_store = None if self.memory_config.memory_storage_mode == "tiered": if not self.memory_config.memory_storage_path: raise ValueError("tiered memory requires memory_storage_path") tier_store = TieredMemoryStoreV2( self.memory_config.memory_storage_path, key_dim=self.memory_config.memory_router_dim, page_capacity=self.memory_config.memory_page_capacity, ) self.memory_os_v2 = MemoryOSV2.from_payload( payload, router=self.memory_router_v2, tier_store=tier_store, max_resident_pages=self.memory_config.memory_resident_pages, runtime_device=self._find_layer_device(), gpu_cache_records=self.memory_config.memory_gpu_cache_records, gpu_cache_tokens=self.memory_config.memory_gpu_cache_tokens, gpu_cache_reserve_mb=self.memory_config.memory_gpu_cache_reserve_mb, gpu_cache_adaptive=self.memory_config.memory_gpu_cache_adaptive, ) # The package configuration is authoritative for runtime safety and # capacity gates. Older snapshots may contain looser values. self.memory_os_v2.read_threshold = self.memory_config.memory_v2_read_threshold self.memory_os_v2.write_threshold = self.memory_config.memory_v2_write_threshold bank = self.memory_os_v2.bank bank.page_capacity = self.memory_config.memory_page_capacity bank.max_pages = self.memory_config.memory_max_pages bank.hot_pages = self.memory_config.memory_hot_pages bank.top_k_pages = self.memory_config.memory_top_k_pages bank.top_k_records = self.memory_config.memory_top_k_records bank.max_hops = self.memory_config.memory_max_hops bank._rebuild_coarse_index(self.memory_config.memory_coarse_index_bits) # The package configuration is authoritative for runtime safety gates. # Older snapshots may contain a looser threshold in their serialized # payload; do not let that silently re-enable noisy recalls after an # upgrade. self.memory_os_v2.read_threshold = self.memory_config.memory_v2_read_threshold self.memory_os_v2.write_threshold = self.memory_config.memory_v2_write_threshold @torch.no_grad() def load_embedded_memory_weights(self, merged_dir: str | Path) -> None: """Load memory modules and optional user state from one safetensors package.""" path = Path(merged_dir) manifest = json.loads((path / "memory_merge.json").read_text(encoding="utf-8")) from safetensors import safe_open weights_path = path / str(manifest.get("memory_weights", "memory.safetensors")) tensor_prefix = str(manifest.get("tensor_prefix", "dynamic_memory.")) with safe_open(str(weights_path), framework="pt", device="cpu") as handle: file_metadata = handle.metadata() or {} weights = { key: handle.get_tensor(key) for key in handle.keys() if key.startswith(tensor_prefix) } memory_prefix = f"{tensor_prefix}memory." memory_state = { key.removeprefix(memory_prefix): value for key, value in weights.items() if key.startswith(memory_prefix) } if not memory_state: raise ValueError(f"merged memory file has no {memory_prefix}* tensors: {weights_path}") self.memory.load_state_dict(memory_state, strict=True) retriever_prefix = f"{tensor_prefix}text_retriever." retriever_state = { key.removeprefix(retriever_prefix): value for key, value in weights.items() if key.startswith(retriever_prefix) } if retriever_state: if self.text_retriever is None: raise ValueError("merged package contains a text retriever but it is disabled") self.text_retriever.load_state_dict(retriever_state, strict=True) self._text_retriever_ready = True policy_prefix = f"{tensor_prefix}memory_policy." policy_state = { key.removeprefix(policy_prefix): value for key, value in weights.items() if key.startswith(policy_prefix) } if policy_state: if self.memory_policy is None: raise ValueError("merged package contains a memory policy but it is disabled") # A v2 candidate may be loaded on top of an older embedded package: # the old shard has only the write head, while the candidate adds # a trained forget head. The adapter load immediately afterwards # is strict and supplies the complete v2 policy. self.memory_policy.load_state_dict( policy_state, strict=self.memory_config.automatic_memory_policy_version < 2, ) self._memory_policy_ready = True router_prefix = f"{tensor_prefix}memory_router_v2." router_state = { key.removeprefix(router_prefix): value for key, value in weights.items() if key.startswith(router_prefix) } if router_state: if self.memory_router_v2 is None: raise ValueError("merged package contains a V2 router but hierarchical memory is disabled") self.memory_router_v2.load_state_dict(router_state, strict=True) self._memory_router_v2_ready = True for index, adapter in zip(self.layer_indices, self._memory_adapters): key = f"{tensor_prefix}blend_logits.{index}" if hasattr(adapter, "blend_logit") and key in weights: adapter.blend_logit.data.copy_(weights[key].to(adapter.blend_logit.device)) persistent_prefix = f"{tensor_prefix}persistent." persistent_payload = { key.removeprefix(persistent_prefix): value for key, value in weights.items() if key.startswith(persistent_prefix) } if persistent_payload: self._load_persistent_memory_payload(persistent_payload) v2_payload = file_metadata.get("memory_os_v2_payload") if v2_payload and self.memory_os_v2 is not None: self._load_memory_os_v2_checkpoint(v2_payload, weights) @torch.no_grad() def save_embedded_memory_weights(self, merged_dir: str | Path) -> None: """Atomically write the current memory snapshot back into the main shard. The merged package keeps the original Qwen shard as the base payload and appends ``dynamic_memory.*`` tensors to the same safetensors file. This intentionally rewrites the package's second shard, so callers should use it as a deliberate weight-backed persistence mode rather than on every token. """ path = Path(merged_dir) manifest = json.loads((path / "memory_merge.json").read_text(encoding="utf-8")) from safetensors.torch import save_file from .merge_memory_weights import _write_merged_shard memory_weights_name = str( manifest.get("memory_weights", "model.safetensors-00002-of-00002.safetensors") ) target_shard = path / memory_weights_name if not target_shard.exists(): raise FileNotFoundError(f"merged package target shard is missing: {target_shard}") tensors: dict[str, Tensor] = {} tensors.update( { f"dynamic_memory.memory.{key}": value.detach().cpu().contiguous() for key, value in self.memory.state_dict().items() } ) if self.text_retriever is not None and self._text_retriever_ready: tensors.update( { f"dynamic_memory.text_retriever.{key}": value.detach().cpu().contiguous() for key, value in self.text_retriever.state_dict().items() } ) if self.memory_router_v2 is not None: tensors.update( { f"dynamic_memory.memory_router_v2.{key}": value.detach().cpu().contiguous() for key, value in self.memory_router_v2.state_dict().items() } ) if self.memory_policy is not None and self._memory_policy_ready: tensors.update( { f"dynamic_memory.memory_policy.{key}": value.detach().cpu().contiguous() for key, value in self.memory_policy.state_dict().items() } ) for index, adapter in zip(self.layer_indices, self._memory_adapters): if hasattr(adapter, "blend_logit"): tensors[f"dynamic_memory.blend_logits.{index}"] = ( adapter.blend_logit.detach().cpu().contiguous() ) if self.memory_config.persistent_memory and self.persistent_memory.numel() > 0: state = self.persistent_memory else: state = self.runtime.state if self.runtime.state is not None else self._persistent_memory if state is None: raise ValueError("cannot write embedded memory without a runtime memory state") tensors["dynamic_memory.persistent.memory_state"] = state.detach().cpu().contiguous() if self.memory_config.natural_language_memory: if self.memory_config.persistent_memory and self.persistent_text_token_ids.numel() > 0: runtime_fields = { "text_token_ids": self.persistent_text_token_ids, "text_token_mask": self.persistent_text_token_mask, "text_slot_valid": self.persistent_text_slot_valid, "text_slot_keys": self.persistent_text_slot_keys, "text_slot_age": self.persistent_text_slot_age, "text_write_counter": self.persistent_text_write_counter, "text_key_token_ids": self.persistent_text_key_token_ids, "text_key_token_mask": self.persistent_text_key_token_mask, } else: runtime_fields = { "text_token_ids": self.runtime.text_token_ids, "text_token_mask": self.runtime.text_token_mask, "text_slot_valid": self.runtime.text_slot_valid, "text_slot_keys": self.runtime.text_slot_keys, "text_slot_age": self.runtime.text_slot_age, "text_write_counter": self.runtime.text_write_counter, "text_key_token_ids": self.runtime.text_key_token_ids, "text_key_token_mask": self.runtime.text_key_token_mask, } if not all(isinstance(value, Tensor) for value in runtime_fields.values()): raise ValueError("cannot write embedded memory without a complete text memory bank") tensors.update( { f"dynamic_memory.persistent.{key}": value.detach().cpu().contiguous() for key, value in runtime_fields.items() } ) v2_metadata: dict[str, Any] = {} if self.memory_os_v2 is not None: v2_tensors, v2_metadata = self._export_memory_os_v2_checkpoint() tensors.update(v2_tensors) extra_path = path / f".{memory_weights_name}.memory-update.tmp" staging_path = path / f".{memory_weights_name}.merged-update.tmp" def _sync_model_index(shard_name: str) -> None: """Keep the HF shard index aligned with the embedded memory shard.""" index_path = path / "model.safetensors.index.json" if not index_path.exists(): return index = json.loads(index_path.read_text(encoding="utf-8")) weight_map = index.get("weight_map") if not isinstance(weight_map, dict): return changed = False for key, value in list(weight_map.items()): if str(key).startswith("dynamic_memory.") and value != shard_name: weight_map[key] = shard_name changed = True if changed: index_path.write_text( json.dumps(index, indent=2, ensure_ascii=False), encoding="utf-8" ) try: metadata = {"format": "qwen_dynamic_memory_embedded_v2"} if v2_metadata: metadata["memory_os_v2_payload"] = json.dumps( v2_metadata, ensure_ascii=False, separators=(",", ":") ) save_file(tensors, str(extra_path), metadata=metadata) if str(manifest.get("shard_mode", "combined")) == "extra_only": # The extra-only layout keeps the official Qwen shards # untouched. Windows may keep the currently loaded shard # open, so rotate to a new small shard when replacement is # denied instead of requiring the chat process to exit. try: os.replace(extra_path, target_shard) except PermissionError: rotated_name = f"{target_shard.stem}.runtime-{time.time_ns()}.safetensors" rotated_path = path / rotated_name os.replace(extra_path, rotated_path) manifest["memory_weights"] = rotated_name manifest["checkpoint_contains_user_memory"] = True (path / "memory_merge.json").write_text( json.dumps(manifest, indent=2, ensure_ascii=False), encoding="utf-8" ) _sync_model_index(str(manifest.get("memory_weights", memory_weights_name))) return base_dir = Path(str(manifest["base_model"])) base_shard = base_dir / memory_weights_name if not base_shard.exists(): raise FileNotFoundError(f"merged package base shard is missing: {base_shard}") _write_merged_shard(base_shard, extra_path, staging_path) os.replace(staging_path, target_shard) _sync_model_index(memory_weights_name) finally: if extra_path.exists(): extra_path.unlink() if staging_path.exists(): staging_path.unlink() def forward( self, *args: Any, memory_state: Optional[Tensor] = None, update_memory: bool = True, read_memory: bool = True, detach_memory: bool = False, return_memory: bool = True, memory_text_input_ids: Optional[Tensor] = None, memory_text_attention_mask: Optional[Tensor] = None, memory_key_input_ids: Optional[Tensor] = None, memory_key_attention_mask: Optional[Tensor] = None, memory_query_input_ids: Optional[Tensor] = None, memory_query_attention_mask: Optional[Tensor] = None, memory_storage_input_ids: Optional[Tensor] = None, memory_storage_attention_mask: Optional[Tensor] = None, force_memory_write: bool = False, memory_text: Optional[str] = None, **kwargs: Any, ) -> Any: input_ids = kwargs.get("input_ids") if input_ids is None and args: input_ids = args[0] batch_size = input_ids.shape[0] if input_ids is not None else None reset_mask = self._reset_mask(input_ids) if reset_mask is not None and self.memory_os_v2 is not None: # The reset token is a model-level durable-memory clear signal, # not merely a transient recurrent-state reset. self.clear_hierarchical_memory() if memory_state is not None: state = memory_state elif self.memory_config.persistent_memory and self.persistent_memory.numel() > 0: state = self.persistent_memory elif self._persistent_memory is not None: state = self._persistent_memory elif batch_size is not None: state = self.memory.initial_state(batch_size, device=self._find_layer_device()) else: state = None if state is not None and batch_size is not None and state.shape[0] != batch_size: if self.memory_config.persistent_memory and state.shape[0] == 1: state = state.expand(batch_size, -1, -1).clone() else: raise ValueError(f"memory batch {state.shape[0]} does not match input batch {batch_size}") if reset_mask is not None and state is not None: state = state.clone() state[reset_mask] = 0 if batch_size is not None: self._bind_text_memory(batch_size, device=self._find_layer_device()) if reset_mask is not None: self._clear_text_memory(reset_mask) self.runtime.state = state self.runtime.read_enabled = read_memory self.runtime.update_enabled = update_memory self.runtime.last_read = None self.runtime.reset_mask = reset_mask self.runtime.input_ids = input_ids self.runtime.attention_mask = kwargs.get("attention_mask") kwargs.setdefault("use_cache", False) base_output = self.base_model(*args, **kwargs) if ( not self.memory_config.natural_language_memory and self.memory_config.direct_logit_scale > 0.0 and read_memory and self.runtime.last_read is not None and hasattr(base_output, "logits") ): read = self.runtime.last_read output_embeddings = self.base_model.get_output_embeddings() memory_logits = output_embeddings(read.to(output_embeddings.weight.device)) logits = base_output.logits + self.memory_config.direct_logit_scale * memory_logits.to( base_output.logits.device ) base_output.logits = logits labels = kwargs.get("labels") if labels is not None: base_output.loss = F.cross_entropy( logits[..., :-1, :].contiguous().view(-1, logits.shape[-1]), labels[..., 1:].contiguous().view(-1), ignore_index=-100, ) if ( not self.memory_config.natural_language_memory and self.memory_config.raw_logit_scale > 0.0 and read_memory and self.runtime.raw_memory is not None and hasattr(base_output, "logits") ): output_embeddings = self.base_model.get_output_embeddings() raw_logits = output_embeddings(self.runtime.raw_memory.to(output_embeddings.weight.device)) raw_logits = raw_logits[:, None, :].expand(-1, base_output.logits.shape[1], -1) base_output.logits = base_output.logits + self.memory_config.raw_logit_scale * raw_logits.to( base_output.logits.device ) labels = kwargs.get("labels") if labels is not None: base_output.loss = F.cross_entropy( base_output.logits[..., :-1, :].contiguous().view(-1, base_output.logits.shape[-1]), labels[..., 1:].contiguous().view(-1), ignore_index=-100, ) new_state = self.runtime.state if new_state is not None and detach_memory: new_state = new_state.detach() self._persistent_memory = new_state.detach() if new_state is not None else None if self.memory_config.persistent_memory and self.runtime.use_persistent_state and new_state is not None: self.persistent_memory = new_state.detach() if self.memory_config.natural_language_memory and update_memory: self.runtime.auto_memory_probability = None self.runtime.auto_memory_forget_probability = None if ( self.memory_policy is not None and self.memory_config.automatic_memory and self._memory_policy_ready and isinstance( getattr( self.memory, "last_write_summary", getattr(self.memory, "last_write_representation", None), ), Tensor, ) ): policy_input = getattr( self.memory, "last_write_summary", self.memory.last_write_representation, ) policy_logits = self.memory_policy(policy_input) self.runtime.auto_memory_probability = torch.sigmoid(policy_logits).detach() if hasattr(self.memory_policy, "forget_logits"): forget_logits = self.memory_policy.forget_logits(policy_input) self.runtime.auto_memory_forget_probability = torch.sigmoid(forget_logits).detach() self._write_text_memory( input_ids, kwargs.get("attention_mask"), memory_text_input_ids, memory_text_attention_mask, memory_key_input_ids, memory_key_attention_mask, memory_storage_input_ids, memory_storage_attention_mask, force_memory_write, memory_text, ) self.runtime.state = new_state return QwenDynamicMemoryOutput(base_output, new_state) if return_memory else base_output @torch.no_grad() def generate( self, *args: Any, memory_state: Optional[Tensor] = None, update_memory: bool = True, memory_query_input_ids: Optional[Tensor] = None, memory_query_attention_mask: Optional[Tensor] = None, memory_query_text: Optional[str] = None, **kwargs: Any, ) -> Any: """Use the native HF generation loop with memory-aware decoder layers.""" input_ids = kwargs.get("input_ids") if input_ids is None and args: input_ids = args[0] self.runtime.context_compaction = None reset_mask = self._reset_mask(input_ids) if reset_mask is not None and self.memory_os_v2 is not None: # Clear before any context compaction so the reset marker can # never cause the old bank to be archived again in this call. self.clear_hierarchical_memory() if input_ids is not None and kwargs.get("attention_mask") is None: kwargs["attention_mask"] = torch.ones_like(input_ids) if input_ids is not None and reset_mask is None: compacted_ids, compacted_mask, compaction_info = self.compact_context_for_kv( input_ids, kwargs.get("attention_mask"), ) input_ids = compacted_ids kwargs["input_ids"] = compacted_ids kwargs["attention_mask"] = compacted_mask self.runtime.context_compaction = compaction_info elif input_ids is not None: self.runtime.context_compaction = { "compacted": False, "reason": "reset_token", "archived_records": 0, } # This is the public prompt returned by generate. When compaction # occurred it intentionally contains the bounded hot window, not a # second copy of the potentially enormous source prompt. original_input_ids = input_ids original_prompt_length = int(input_ids.shape[1]) if input_ids is not None else 0 if memory_state is not None: self.runtime.state = memory_state elif ( self.memory_config.persistent_memory and self.runtime.use_persistent_state and self.persistent_memory.numel() > 0 ): self.runtime.state = self.persistent_memory elif self._persistent_memory is not None and input_ids is not None and self._persistent_memory.shape[0] == input_ids.shape[0]: self.runtime.state = self._persistent_memory elif input_ids is not None: self.runtime.state = self.memory.initial_state(input_ids.shape[0], device=self._find_layer_device()) if reset_mask is not None and self.runtime.state is not None: state = self.runtime.state.clone() if state.shape[0] == 1 and reset_mask.shape[0] > 1: state = state.expand(reset_mask.shape[0], -1, -1).clone() state[reset_mask] = 0 self.runtime.state = state if input_ids is not None: if self.memory_config.natural_language_memory: self._bind_text_memory(input_ids.shape[0], device=self._find_layer_device()) if reset_mask is not None: self._clear_text_memory(reset_mask) if "attention_mask" not in kwargs or kwargs["attention_mask"] is None: kwargs["attention_mask"] = torch.ones_like(input_ids) prefix_ids: Optional[Tensor] = None prefix_mask: Optional[Tensor] = None prefix_length = 0 retrieval_ids = memory_query_input_ids if memory_query_input_ids is not None else input_ids retrieval_mask = ( memory_query_attention_mask if memory_query_attention_mask is not None else kwargs["attention_mask"] ) if ( self.memory_config.natural_language_memory and self.runtime.text_slot_valid is not None and retrieval_ids is not None and retrieval_mask is not None ): prefix_ids, prefix_mask, prefix_length = self._build_text_prefix( retrieval_ids.to(device=input_ids.device), retrieval_mask.to(device=input_ids.device), query_text=memory_query_text or "", ) if prefix_ids is not None and prefix_mask is not None and prefix_length > 0: input_ids = torch.cat((prefix_ids, input_ids), dim=1) kwargs["input_ids"] = input_ids kwargs["attention_mask"] = torch.cat( (prefix_mask, kwargs["attention_mask"]), dim=1 ) # Internal memory tokens are hidden from the public wrapper # API, but remain visible to the actual Qwen forward pass. self.runtime.input_ids = input_ids # In natural-language mode, the retrieved text prefix is the complete # read path. Running the older continuous residual at the same time # would distort Qwen's native chat control tokens. With no prefix the # memory path is an exact no-op and the base model remains untouched. self.runtime.read_enabled = not self.memory_config.natural_language_memory self.runtime.update_enabled = update_memory self.runtime.last_read = None self.runtime.reset_mask = reset_mask self.runtime.input_ids = input_ids if self.memory_config.kv_offload or self.memory_config.kv_cache_implementation is not None: # Transformers 5.x owns the offload lifecycle. Supplying an # explicit DynamicCache lets us honor the non-sliding-layer # setting, which the shorthand ``cache_implementation`` does not # expose for dynamic caches. kwargs["use_cache"] = True cache_implementation = self.memory_config.kv_cache_implementation if cache_implementation is None: cache_implementation = "offloaded" if ( cache_implementation == "offloaded" and kwargs.get("past_key_values") is None ): try: from transformers.cache_utils import DynamicCache, LinearAttentionCacheLayerMixin class _NaturalMemoryOffloadedCache(DynamicCache): """Prefetch the current layer before every cache update. Transformers 5.9 prefetches the next layer as part of ``Cache.update``. Qwen3.5's hybrid linear/full stack can enter the same layer again on the next decoding step before that next-layer prefetch runs, leaving a CPU tensor to concatenate with CUDA states. The current-layer prefetch makes the ownership explicit for both attention and recurrent cache updates. """ def update(self, key_states, value_states, layer_idx, *args, **extra): if self.offloading: self.prefetch(layer_idx, self.only_non_sliding) return super().update(key_states, value_states, layer_idx, *args, **extra) def update_conv_state(self, conv_states, layer_idx, **extra): if self.offloading: self.prefetch(layer_idx, self.only_non_sliding) value = super().update_conv_state(conv_states, layer_idx, **extra) if self.offloading: self.offload(layer_idx, self.only_non_sliding) return value def update_recurrent_state(self, recurrent_states, layer_idx, **extra): if self.offloading: self.prefetch(layer_idx, self.only_non_sliding) value = super().update_recurrent_state(recurrent_states, layer_idx, **extra) if self.offloading: self.offload(layer_idx, self.only_non_sliding) return value def offload(self, layer_idx, only_non_sliding=True): # Qwen3.5 linear-attention state is tiny and is # read before update. Keep it on the execution # device; offload only the full-attention KV that # dominates memory usage. if isinstance(self.layers[layer_idx], LinearAttentionCacheLayerMixin): return return super().offload(layer_idx, only_non_sliding) decoder_config = ( self.base_model.config.get_text_config(decoder=True) if hasattr(self.base_model.config, "get_text_config") else getattr(self.base_model.config, "text_config", self.base_model.config) ) kwargs["past_key_values"] = _NaturalMemoryOffloadedCache( config=decoder_config, offloading=True, offload_only_non_sliding=self.memory_config.kv_offload_only_non_sliding, ) except (ImportError, TypeError, AttributeError): # Older Transformers releases can still understand the # public generation option even if their DynamicCache # constructor has a different signature. kwargs["cache_implementation"] = cache_implementation elif kwargs.get("past_key_values") is None: kwargs["cache_implementation"] = cache_implementation # Natural-language memory is consumed as an internal text prefix. Do # not also add a constant token-level logit bias, because that would # repeat the first value token at every generation position. if ( not self.memory_config.natural_language_memory and (self.memory_config.direct_logit_scale > 0.0 or self.memory_config.raw_logit_scale > 0.0) ): from transformers import LogitsProcessor, LogitsProcessorList runtime = self.runtime output_embeddings = self.base_model.get_output_embeddings() prompt_length = int(input_ids.shape[1]) if input_ids is not None else None class MemoryLogitsProcessor(LogitsProcessor): def __call__(self, input_ids: Tensor, scores: Tensor) -> Tensor: first_generated_token = prompt_length is None or input_ids.shape[1] == prompt_length if first_generated_token and runtime.raw_memory is not None and raw_scale > 0.0: raw = runtime.raw_memory.to(output_embeddings.weight.device) scores = scores + raw_scale * output_embeddings(raw).to(scores.device) if runtime.last_read is None: return scores last_read = runtime.last_read[:, -1].to(output_embeddings.weight.device) memory_logits = output_embeddings(last_read).to(scores.device) return scores + self_scale * memory_logits self_scale = self.memory_config.direct_logit_scale raw_scale = self.memory_config.raw_logit_scale existing = kwargs.get("logits_processor") if existing is None: kwargs["logits_processor"] = LogitsProcessorList([MemoryLogitsProcessor()]) else: kwargs["logits_processor"] = LogitsProcessorList( list(existing) + [MemoryLogitsProcessor()] ) if input_ids is not None and args: args = () generated = self.base_model.generate(*args, **kwargs) if ( original_input_ids is not None and input_ids is not None and input_ids.shape[1] != original_prompt_length ): generated = torch.cat( ( original_input_ids, generated[:, input_ids.shape[1] :], ), dim=1, ) if self.runtime.state is not None: self._persistent_memory = self.runtime.state.detach() if self.memory_config.persistent_memory and self.runtime.use_persistent_state: self.persistent_memory = self._persistent_memory return generated def train(self, mode: bool = True) -> "QwenDynamicMemoryModel": self.memory.train(mode) if self.memory_router_v2 is not None: self.memory_router_v2.train(mode) self.base_model.eval() return self def load_qwen_base( model_path: str | Path, *, load_in_4bit: bool = True, device_map: str | dict[str, Any] = "auto", ): """Load the unmodified local Qwen checkpoint with the benchmark defaults.""" import transformers load_kwargs: dict[str, Any] = { "device_map": device_map, "low_cpu_mem_usage": True, "dtype": torch.bfloat16, } if load_in_4bit: load_kwargs["quantization_config"] = transformers.BitsAndBytesConfig( load_in_4bit=True, bnb_4bit_quant_type="nf4", bnb_4bit_use_double_quant=True, bnb_4bit_compute_dtype=torch.bfloat16, ) return transformers.AutoModelForImageTextToText.from_pretrained(str(model_path), **load_kwargs) def load_qwen_dynamic( model_path: str | Path, *, memory_config: Optional[QwenMemoryConfig] = None, load_in_4bit: bool = True, device_map: str | dict[str, Any] = "auto", freeze_backbone: bool = True, ) -> QwenDynamicMemoryModel: """Load Qwen and attach either an explicit adapter or an embedded merge package.""" model_path = Path(model_path) embedded_manifest = model_path / "memory_merge.json" if memory_config is None and embedded_manifest.exists(): embedded_config = model_path / "memory_config.json" if not embedded_config.exists(): raise FileNotFoundError( f"embedded memory package is missing {embedded_config}" ) memory_config = load_memory_config(model_path) if memory_config is not None and memory_config.memory_storage_mode == "tiered": storage_path = Path(memory_config.memory_storage_path) if memory_config.memory_storage_path else None if storage_path is None: storage_path = model_path / "memory_pages.sqlite" elif not storage_path.is_absolute(): storage_path = model_path / storage_path memory_config.memory_storage_path = str(storage_path) base_model = load_qwen_base( model_path, load_in_4bit=load_in_4bit, device_map=device_map, ) model = QwenDynamicMemoryModel( base_model, memory_config=memory_config, freeze_backbone=freeze_backbone, ) if embedded_manifest.exists(): model.load_embedded_memory_weights(model_path) return model def load_tokenizer(model_path: str | Path): import transformers tokenizer = transformers.AutoTokenizer.from_pretrained(str(model_path), use_fast=True) if tokenizer.pad_token is None: tokenizer.pad_token = tokenizer.eos_token return tokenizer