"""Hugging Face remote-code support for shared-weight SenseNova Loop and MHC. This single file contains configuration, plan validation, MHC routing and the model implementation. The original physical layer weights are reused on every Loop visit; autoregressive cache states are independent per visit. """ import copy import math from contextvars import ContextVar import torch from torch import nn from transformers import Qwen3_5Config from transformers.cache_utils import DynamicCache from transformers.models.qwen3_5 import modeling_qwen3_5 as hf _attention_capture = ContextVar("sensenova_attention_capture", default=None) class SenseNovaLoopConfig(Qwen3_5Config): model_type = "sensenova_loop" def execution_plan(config): """Return physical layer indices in the order they execute during inference.""" loop = config.get("loop_settings") if not isinstance(loop, dict): raise ValueError("This implementation requires explicit loop_settings") text_config = config["text_config"] num_layers = text_config["num_hidden_layers"] start, end, count = loop.get("start"), loop.get("end"), loop.get("eval_count") if any(type(value) is not int for value in (num_layers, start, end, count)) or not ( 0 <= start < end <= num_layers and count >= 1 ): raise ValueError("Invalid Loop interval or eval_count") if loop.get("counts") != [count]: raise ValueError("Only fixed-count Loop checkpoints are supported") layer_types = text_config["layer_types"] if len(layer_types) != num_layers or set(layer_types) - {"full_attention", "linear_attention"}: raise ValueError("Unexpected physical layer types") if "full_attention" not in layer_types: raise ValueError("At least one full-attention layer is required for sequence-length tracking") mhc = config.get("mhc") if mhc is not None: mhc_start, mhc_end, streams, iterations = ( mhc.get(key) for key in ("start", "end", "streams", "iterations") ) if ( any(type(value) is not int for value in (mhc_start, mhc_end, streams, iterations)) or not 0 <= mhc_start < mhc_end <= num_layers or streams not in (2, 4) or iterations < 1 ): raise ValueError("Invalid MHC range, stream count, or iterations") if type(mhc.get("fixed_identity")) is not bool: raise ValueError("MHC fixed_identity must be a boolean") scale = mhc.get("pre_scale", 1.0) if not isinstance(scale, (int, float)) or not math.isfinite(scale) or scale <= 0: raise ValueError("Invalid MHC pre_scale") allowed = {"start", "end", "streams", "iterations", "fixed_identity", "pre_scale"} if set(mhc) - allowed: raise ValueError("Unknown MHC configuration fields") phase_boundaries = (0, start, end, num_layers) if not any( left <= mhc_start < mhc_end <= right for left, right in zip(phase_boundaries, phase_boundaries[1:]) ): raise ValueError("MHC cannot cross a Loop phase boundary") return tuple(range(start)) + tuple(range(start, end)) * count + tuple(range(end, num_layers)) def sinkhorn(logits, iterations): """Normalize routing logits across rows and columns in log space.""" logits = logits.float() for _ in range(iterations): logits = logits - torch.logsumexp(logits, dim=-1, keepdim=True) logits = logits - torch.logsumexp(logits, dim=-2, keepdim=True) return logits.exp() class MHCRouter(nn.Module): def __init__(self, dim, settings): super().__init__() self.streams = settings["streams"] self.iterations = settings["iterations"] self.fixed_identity = settings["fixed_identity"] self.pre_scale = settings.get("pre_scale", 1.0) streams = self.streams output_size = 2 * streams + (0 if self.fixed_identity else streams * streams) self.weight = nn.Parameter(torch.zeros(output_size, streams * dim, dtype=torch.float32)) self.bias = nn.Parameter(torch.zeros(output_size, dtype=torch.float32)) with torch.no_grad(): initial_probabilities = {2: (0.25, 0.75), 4: (0.0625, 0.1875, 0.25, 0.5)}[streams] self.bias[:streams].copy_(torch.logit(torch.tensor(initial_probabilities, device=self.bias.device))) if not self.fixed_identity: diagonal = math.log((1 - (streams - 1) / 16) * 16) self.bias[2 * streams:].copy_((torch.eye(streams, device=self.bias.device) * diagonal).flatten()) def _apply(self, fn, recurse=True): # Keep the original FP32 values when the surrounding model is cast to BF16. def keep_fp32(tensor): result = fn(tensor) if tensor.is_floating_point() and result.dtype != torch.float32: if tensor.device.type == "meta" and result.device.type != "meta": return torch.empty_like(result, dtype=torch.float32) return tensor.to(device=result.device, dtype=torch.float32) return result return super()._apply(keep_fp32, recurse=recurse) def forward(self, residual): if self.weight.dtype != torch.float32 or self.bias.dtype != torch.float32: raise RuntimeError("mHC routers require FP32 parameters; configure FSDP per-module precision") with torch.autocast(device_type=residual.device.type, enabled=False): flat = residual.flatten(-2).float() normalized = flat * torch.rsqrt(flat.square().mean(-1, keepdim=True) + 1e-6) logits = torch.nn.functional.linear(normalized, self.weight, self.bias) streams = self.streams pre = logits[..., :streams].sigmoid() post = 2 * logits[..., streams:2 * streams].sigmoid() if self.fixed_identity: mix = torch.eye(streams, device=flat.device, dtype=torch.float32).expand( *flat.shape[:-1], streams, streams ) else: mix = sinkhorn(logits[..., 2 * streams:].unflatten(-1, (streams, streams)), self.iterations) branch = torch.matmul(pre.to(residual.dtype).unsqueeze(-2), residual).squeeze(-2) if self.pre_scale != 1.0: branch = branch * self.pre_scale return branch, post, mix @staticmethod def merge(residual, branch, post, mix): """Mix residual streams and add the branch output to each stream.""" with torch.autocast(device_type=residual.device.type, enabled=False): mixed_residual = torch.matmul(mix.to(residual.dtype), residual) branch_update = post.to(branch.dtype).unsqueeze(-1) * branch.unsqueeze(-2) return mixed_residual + branch_update def enter(hidden, physical, settings): """Expand to residual streams at the first mHC layer, or restore their shape.""" if physical == settings["start"]: return hidden.unsqueeze(-2).expand( *hidden.shape[:-1], settings["streams"], hidden.shape[-1] ).contiguous() return hidden.unflatten(-1, (settings["streams"], -1)) def leave(residual, physical, settings): """Reduce streams at the last mHC layer; otherwise carry them forward.""" return residual.mean(-2) if physical == settings["end"] - 1 else residual.flatten(-2) class _CacheLayersView: """Expose the execution slot when a layer accesses its physical index.""" def __init__(self, cache, physical, virtual): self.cache = cache self.physical = physical self.virtual = virtual def __getitem__(self, index): if index != self.physical: raise IndexError("A physical layer may only access its own virtual cache slot") return self.cache.layers[self.virtual] class _VisitCache: """Map a physical attention module's layer_idx to one execution slot. Mapping is local to this forward call: shared modules and their layer_idx are never mutated. `layers` supports Transformers 5.12; state accessors support newer official backbone cache APIs as well. """ def __init__(self, cache, physical, virtual): self.cache = cache self.virtual = virtual self.layers = _CacheLayersView(cache, physical, virtual) def update(self, key_states, value_states, layer_idx, **kwargs): return self.cache.update(key_states, value_states, self.virtual, **kwargs) def has_previous_state(self, layer_idx=None, **kwargs): return self.cache.has_previous_state(self.virtual, **kwargs) def update_conv_state(self, state, layer_idx, **kwargs): return self.cache.update_conv_state(state, self.virtual, **kwargs) def update_recurrent_state(self, state, layer_idx, **kwargs): return self.cache.update_recurrent_state(state, self.virtual, **kwargs) def get_conv_state(self, layer_idx, **kwargs): return self.cache.get_conv_state(self.virtual, **kwargs) def get_recurrent_state(self, layer_idx, **kwargs): return self.cache.get_recurrent_state(self.virtual, **kwargs) def is_recording_past(self, layer_idx): return self.cache.is_recording_past(self.virtual) class SenseNovaLoopTextModel(hf.Qwen3_5TextModel): def forward(self, input_ids=None, attention_mask=None, position_ids=None, past_key_values=None, inputs_embeds=None, use_cache=None, output_hidden_states=None, output_attentions=None, **kwargs): """Run Loop visits, optionally returning their states and attention weights. Hidden states include the initial embedding and one state per visit. MHC states between its boundaries retain their flattened residual streams. Attention entries follow visit order, with None for linear-attention layers. """ if (input_ids is None) == (inputs_embeds is None): raise ValueError("Specify exactly one of input_ids or inputs_embeds") if inputs_embeds is None: inputs_embeds = self.embed_tokens(input_ids) use_cache = self.config.use_cache if use_cache is None else use_cache if output_hidden_states is None: output_hidden_states = getattr(self.config, "output_hidden_states", False) if output_attentions is None: output_attentions = getattr(self.config, "output_attentions", False) if not isinstance(output_hidden_states, bool): raise TypeError("output_hidden_states must be a boolean.") if output_attentions and self.config._attn_implementation != "eager": raise ValueError( "Attention weights require attn_implementation='eager' when loading the model." ) if self.gradient_checkpointing and self.training and use_cache: raise ValueError("Disable use_cache when training with gradient checkpointing") if not use_cache and past_key_values is not None: raise ValueError("past_key_values requires use_cache=True") if past_key_values is not None and not isinstance(past_key_values, DynamicCache): raise TypeError("Loop generation currently supports DynamicCache only") if use_cache: if past_key_values is None: past_key_values = DynamicCache(config=self._loop_cache_config) past_key_values._loop_execution_plan = self._loop_plan elif getattr(past_key_values, "_loop_execution_plan", None) != self._loop_plan: # GenerationMixin may allocate an empty physical-depth cache. # A populated foreign cache cannot be reinterpreted safely. if past_key_values.get_seq_length() != 0 or past_key_values.has_previous_state(): raise ValueError("A populated cache from another execution plan cannot be reused") past_key_values = DynamicCache(config=self._loop_cache_config) past_key_values._loop_execution_plan = self._loop_plan if position_ids is None: seen = past_key_values.get_seq_length() if past_key_values is not None else 0 position_ids = torch.arange(inputs_embeds.shape[1], device=inputs_embeds.device) + seen position_ids = position_ids.view(1, 1, -1).expand(4, inputs_embeds.shape[0], -1) elif position_ids.ndim == 2: position_ids = position_ids.unsqueeze(0).expand(4, -1, -1) if position_ids.ndim == 3 and position_ids.shape[0] == 4: text_positions, position_ids = position_ids[0], position_ids[1:] else: text_positions = None if isinstance(attention_mask, dict): masks = attention_mask else: mask_args = dict( config=self._loop_cache_config, inputs_embeds=inputs_embeds, attention_mask=attention_mask, past_key_values=past_key_values, position_ids=text_positions, ) if hasattr(hf, "create_recurrent_attention_mask"): linear_mask = hf.create_recurrent_attention_mask(**mask_args) else: linear_mask = self._update_linear_attn_mask(attention_mask, past_key_values) masks = { "full_attention": hf.create_causal_mask(**mask_args), "linear_attention": linear_mask, } hidden = inputs_embeds position_embeddings = self.rotary_emb(hidden, position_ids) hidden_states = [hidden] if output_hidden_states else None attentions = [] if output_attentions else None hooks = [] capture_token = None def record_attention(_module, _inputs, output): if _attention_capture.get() is attentions: if output[1] is None: raise RuntimeError("The attention backend did not return attention weights.") attentions.append(output[1]) # Repeated visits share weights, but each visit has its own cache slot. try: if attentions is not None: for layer in self.layers: if hasattr(layer, "self_attn"): hooks.append(layer.self_attn.register_forward_hook(record_attention)) capture_token = _attention_capture.set(attentions) for execution_index, physical_index in enumerate(self._loop_plan): layer_type = self.config.layer_types[physical_index] if attentions is not None and layer_type == "linear_attention": attentions.append(None) cache = ( _VisitCache(past_key_values, physical_index, execution_index) if use_cache else None ) hidden = self.layers[physical_index]( hidden, position_embeddings=position_embeddings, attention_mask=masks[layer_type], position_ids=text_positions, past_key_values=cache, use_cache=use_cache, **kwargs, ) if hidden_states is not None: hidden_states.append(hidden) finally: if capture_token is not None: _attention_capture.reset(capture_token) for hook in hooks: hook.remove() hidden = self.norm(hidden) if hidden_states is not None: hidden_states[-1] = hidden return hf.Qwen3_5ModelOutputWithPast( last_hidden_state=hidden, past_key_values=past_key_values, hidden_states=tuple(hidden_states) if hidden_states is not None else None, attentions=tuple(attentions) if attentions is not None else None, ) class MHCDecoderLayer(hf.Qwen3_5DecoderLayer): def forward(self, hidden_states, position_embeddings, attention_mask=None, position_ids=None, past_key_values=None, **kwargs): residual = enter(hidden_states, self._mhc_physical, self._mhc_settings) branch, post, mix = self.mhc_attn(residual) branch = self.input_layernorm(branch) kind = self.layer_type if hasattr(self, "layer_type") else self.block_type if kind == "linear_attention": branch = self.linear_attn(hidden_states=branch, cache_params=past_key_values, attention_mask=attention_mask, **kwargs) else: branch, _ = self.self_attn( hidden_states=branch, position_embeddings=position_embeddings, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, **kwargs) residual = MHCRouter.merge(residual, branch, post, mix) branch, post, mix = self.mhc_ffn(residual) branch = self.mlp(self.post_attention_layernorm(branch)) return leave(MHCRouter.merge(residual, branch, post, mix), self._mhc_physical, self._mhc_settings) def _configure_text_model(text, config): raw = config.to_dict() plan = execution_plan(raw) # Change only these new model instances, preserving canonical state_dict # names and avoiding a second allocation of the full backbone. text.__class__ = SenseNovaLoopTextModel text._loop_plan = plan text._loop_cache_config = copy.deepcopy(text.config) text._loop_cache_config.num_hidden_layers = len(plan) text._loop_cache_config.layer_types = [text.config.layer_types[i] for i in plan] if raw.get("mhc") is not None: settings = raw["mhc"] for i in range(settings["start"], settings["end"]): layer = text.layers[i] layer.__class__ = MHCDecoderLayer layer._mhc_physical, layer._mhc_settings = i, dict(settings) layer.mhc_attn = MHCRouter(text.config.hidden_size, settings) layer.mhc_ffn = MHCRouter(text.config.hidden_size, settings) class SenseNovaLoopModel(hf.Qwen3_5Model): config_class = SenseNovaLoopConfig _keep_in_fp32_modules = ["mhc_attn", "mhc_ffn"] _no_split_modules = ["Qwen3_5DecoderLayer", "MHCDecoderLayer", "Qwen3_5VisionBlock"] _keys_to_ignore_on_load_unexpected = [r"^mtp\."] def __init__(self, config): execution_plan(config.to_dict()) super().__init__(config) _configure_text_model(self.language_model, config) class SenseNovaLoopForConditionalGeneration(hf.Qwen3_5ForConditionalGeneration): config_class = SenseNovaLoopConfig _keep_in_fp32_modules = ["mhc_attn", "mhc_ffn"] _no_split_modules = ["Qwen3_5DecoderLayer", "MHCDecoderLayer", "Qwen3_5VisionBlock"] _keys_to_ignore_on_load_unexpected = [r"^mtp\."] def __init__(self, config): execution_plan(config.to_dict()) super().__init__(config) _configure_text_model(self.model.language_model, config)