"""GS1: native hybrid text/vision decision model with attention residuals.""" import hashlib import json import math from dataclasses import dataclass from pathlib import Path from typing import Literal import torch from accelerate import init_empty_weights from pydantic import BaseModel, ConfigDict, Field, field_validator from safetensors.torch import load_file from torch import nn from torch.nn import functional as F from torch.nn.attention import SDPBackend, sdpa_kernel from torch.utils.checkpoint import checkpoint from transformers import AutoConfig, AutoModel, Qwen3_5Config from transformers.models.qwen3_5.modeling_qwen3_5 import ( Qwen3_5Attention, Qwen3_5DynamicCache, Qwen3_5GatedDeltaNet, Qwen3_5MLP, Qwen3_5PreTrainedModel, Qwen3_5RMSNorm, Qwen3_5TextRotaryEmbedding, Qwen3_5VisionModel, apply_rotary_pos_emb, ) from transformers.utils import ModelOutput INFERENCE_LIMITS = dict( max_context=65536, max_questions=512, max_request_tokens=131072, max_branches=5120, batch_tokens=16384, max_attention_bytes=1073741824, max_logit_elements=16777216, ) class ModelConfig(BaseModel): model_config = ConfigDict(extra="forbid", strict=True, allow_inf_nan=False) class DecisionConfig(ModelConfig): attnres: bool = True attnres_layers: list[int] | None = None attnres_eps: float = Field(default=1e-6, gt=0) bidirectional: bool = True bidirectional_layers: list[int] | None = None cross_attention_heads: int = Field(default=8, ge=1) cross_attention_layers: int = Field(default=2, ge=1) projection_size: int = Field(default=256, ge=2) norm_eps: float = Field(default=1e-5, gt=0) @field_validator("attnres_layers", "bidirectional_layers") @classmethod def selected_layers(cls, value): if value is not None: if not value or any(i < 0 for i in value) or len(set(value)) != len(value): raise ValueError("Layer selection must contain distinct nonnegative indices") return sorted(value) return value class CalibrationConfig(ModelConfig): mode: Literal["identity", "scalar_temperature"] = "identity" log_temperature_bound: float = Field(default=math.log(20), gt=0) def named_seed(seed, name): payload = json.dumps( ["gs1.rng", seed, name], ensure_ascii=False, separators=(",", ":"), sort_keys=True, allow_nan=False ) return int.from_bytes(hashlib.sha256(payload.encode()).digest()[:8], "big") % (2**63) def _read_weights(directory): index = directory / "model.safetensors.index.json" if index.is_file(): weight_map = json.loads(index.read_text())["weight_map"] files = sorted(set(weight_map.values())) else: files = sorted(p.name for p in directory.glob("*.safetensors")) weight_map = None if not files: raise ValueError("Checkpoint contains no safetensors weights") if any(Path(name).is_absolute() or ".." in Path(name).parts for name in files): raise ValueError("Checkpoint shard is outside its directory") # Hub snapshots can link to cache blobs outside their snapshot directory. state = {} for name in files: shard = load_file(directory / name, device="cpu") if state.keys() & shard.keys(): raise ValueError("Duplicate checkpoint parameter") state.update(shard) if weight_map is not None and state.keys() != weight_map.keys(): raise ValueError("Checkpoint index differs from its parameter contents") return state class GS1Config(Qwen3_5Config): model_type = "gs1" def __init__( self, decision=None, calibration=None, initialization_seed=12, head_compute_dtype="float32", inference_limits=None, **kwargs, ): kwargs.pop("model_type", None) kwargs.pop("architectures", None) super().__init__(**kwargs) self.decision = DecisionConfig.model_validate(decision or {}).model_dump() self.calibration = CalibrationConfig.model_validate(calibration or {}).model_dump() self.initialization_seed = initialization_seed self.inference_limits = {**INFERENCE_LIMITS, **(inference_limits or {})} if self.inference_limits.keys() != INFERENCE_LIMITS.keys() or any( type(value) is not int or value < 1 for value in self.inference_limits.values() ): raise ValueError("Invalid GS1 inference limits") if head_compute_dtype not in {"float32", "bfloat16"}: raise ValueError("head_compute_dtype must be float32 or bfloat16") self.head_compute_dtype = head_compute_dtype self.architectures = ["GS1Model"] self.tie_word_embeddings = False self.text_config.use_cache = False self.text_config._attn_implementation = "sdpa" self.vision_config._attn_implementation = "sdpa" @dataclass class GS1Output(ModelOutput): logits: torch.Tensor | None = None last_hidden_state: torch.Tensor | None = None @dataclass class GS1TextOutput(ModelOutput): last_hidden_state: torch.Tensor | None = None @dataclass(frozen=True) class GDNState: """Native GDN prefix, including the last complete scan-chunk boundary.""" conv: torch.Tensor recurrent: torch.Tensor boundary_conv: torch.Tensor | None boundary_recurrent: torch.Tensor | None tail: torch.Tensor position: int class GS1DecoderLayer(nn.Module): """A hybrid decoder block; only its full-attention input can read depth history.""" def __init__(self, config, index): super().__init__() self.layer_type = config.layer_types[index] self.config = config self.index = index if self.layer_type == "full_attention": self.self_attn = Qwen3_5Attention(config, index) elif self.layer_type == "linear_attention": self.linear_attn = Qwen3_5GatedDeltaNet(config, index) else: raise ValueError(f"Unsupported decoder layer type: {self.layer_type}") self.mlp = Qwen3_5MLP(config, config.intermediate_size) self.input_layernorm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.post_attention_layernorm = Qwen3_5RMSNorm(config.hidden_size, eps=config.rms_norm_eps) def forward(self, carry, positional, mask, read, terminal_query, valid_lengths=None, cohorts=None): if self.layer_type == "full_attention": mixed, _ = self.self_attn( self.input_layernorm(read), position_embeddings=positional, attention_mask=mask ) elif valid_lengths is not None and min(valid_lengths) != carry.shape[1]: # GDN scans only real tokens. Equal-length cohorts retain native # batched kernels and their original chunk boundaries; FA and FFN # still execute over the complete decoder batch. mixed = torch.zeros_like(carry) if cohorts is None: cohorts = tuple( ( length, torch.tensor( [i for i, value in enumerate(valid_lengths) if value == length], device=carry.device, ), ) for length in sorted(set(valid_lengths)) ) for length, rows in cohorts: value = carry.index_select(0, rows)[:, -length:] normalized = self.input_layernorm(value) if terminal_query: cache = Qwen3_5DynamicCache(self.config) cache.last_linear_layer = self.index context = self.linear_attn( normalized[:, :-1], cache_params=cache, cache_position=torch.arange(length - 1, device=carry.device), ) query, _ = _linear_context( self.linear_attn, normalized[:, -1:], (cache.conv_states[self.index], cache.recurrent_states[self.index]), self.config, ) value = torch.cat((context, query), 1) else: value = self.linear_attn(normalized) mixed = mixed.index_copy( 0, rows, F.pad(value, (0, 0, carry.shape[1] - length, 0)).to(mixed.dtype) ) elif terminal_query: # Keep the trained context scan and differentiable one-token continuation. cache = Qwen3_5DynamicCache(self.config) cache.last_linear_layer = self.index context = carry[:, :-1] context_mixed = self.linear_attn( self.input_layernorm(context), cache_params=cache, cache_position=torch.arange(context.shape[1], device=carry.device), ) query_mixed, _ = _linear_context( self.linear_attn, self.input_layernorm(carry[:, -1:]), (cache.conv_states[self.index], cache.recurrent_states[self.index]), self.config, ) context = context + context_mixed query = carry[:, -1:] + query_mixed context_ffn = self.mlp(self.post_attention_layernorm(context)) query_ffn = self.mlp(self.post_attention_layernorm(query)) return ( torch.cat((context + context_ffn, query + query_ffn), 1), torch.cat((context_mixed, query_mixed), 1), torch.cat((context_ffn, query_ffn), 1), ) else: mixed = self.linear_attn(self.input_layernorm(carry)) hidden = carry + mixed ffn = self.mlp(self.post_attention_layernorm(hidden)) return hidden + ffn, mixed, ffn class CrossAttentionDecisionHead(nn.Module): """Read state at backbone width; project only the final dot-product vectors.""" def __init__(self, hidden_size, decision=None, seed=12): super().__init__() self.config = DecisionConfig.model_validate(decision or {}) if hidden_size % self.config.cross_attention_heads: raise ValueError("cross_attention_heads must divide the backbone hidden size") width = self.config.projection_size eps = self.config.norm_eps self.bf16_compute = False with torch.random.fork_rng(devices=[]): torch.manual_seed(named_seed(seed, "initialization/cross_attention_head")) self.memory_norm = nn.LayerNorm(hidden_size, eps=eps) self.memory_projection = nn.Linear(hidden_size, hidden_size, bias=False) self.question_seed = nn.Linear(hidden_size, hidden_size, bias=False) self.question_attention = nn.MultiheadAttention( hidden_size, self.config.cross_attention_heads, batch_first=True ) self.decision_projection = nn.Linear(hidden_size, hidden_size, bias=False) self.question_projection = nn.Linear(hidden_size, hidden_size, bias=False) self.question_evidence_projection = nn.Linear(hidden_size, hidden_size, bias=False) self.context_norm = nn.LayerNorm(hidden_size, eps=eps) self.option_projection = nn.Linear(hidden_size, hidden_size, bias=False) self.condition_projection = nn.Linear(hidden_size, hidden_size, bias=False) self.option_norm = nn.LayerNorm(hidden_size, eps=eps) self.routing_attention = nn.ModuleList( nn.MultiheadAttention(hidden_size, self.config.cross_attention_heads, batch_first=True) for _ in range(self.config.cross_attention_layers) ) self.routing_attention_norm = nn.ModuleList( nn.LayerNorm(hidden_size, eps=eps) for _ in range(self.config.cross_attention_layers) ) self.routing_ffn = nn.ModuleList( nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.GELU(), nn.Linear(hidden_size, hidden_size), ) for _ in range(self.config.cross_attention_layers) ) self.routing_output_norm = nn.ModuleList( nn.LayerNorm(hidden_size, eps=eps) for _ in range(self.config.cross_attention_layers) ) self.query_mlp = nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.GELU(), nn.Linear(hidden_size, width) ) self.key_mlp = nn.Sequential( nn.Linear(hidden_size, hidden_size), nn.GELU(), nn.Linear(hidden_size, width) ) self.query_norm = nn.LayerNorm(width, eps=eps) self.key_norm = nn.LayerNorm(width, eps=eps) def prepare_memory(self, state): """Prepare immutable state memory and each attention module's own K/V.""" dtype = torch.bfloat16 if self.bf16_compute else self.memory_projection.weight.dtype with torch.autocast(device_type=state.device.type, dtype=torch.bfloat16, enabled=self.bf16_compute): memory = self.memory_projection(self.memory_norm(state.to(dtype))).unsqueeze(0) projections = [] for attention in (self.question_attention, *self.routing_attention): width = attention.embed_dim bias = attention.in_proj_bias[width:] if attention.in_proj_bias is not None else None key, value = F.linear(memory, attention.in_proj_weight[width:], bias).chunk(2, -1) shape = (1, memory.shape[1], attention.num_heads, width // attention.num_heads) projections.append((key.reshape(shape).transpose(1, 2), value.reshape(shape).transpose(1, 2))) return memory, tuple(projections) @staticmethod def _attend(attention, query, memory, projected=None): if projected is None: return attention(query, memory, memory, need_weights=False)[0] width = attention.embed_dim bias = attention.in_proj_bias[:width] if attention.in_proj_bias is not None else None query = F.linear(query, attention.in_proj_weight[:width], bias) batch, length = query.shape[:2] query = query.reshape(batch, length, attention.num_heads, width // attention.num_heads).transpose( 1, 2 ) value = F.scaled_dot_product_attention(query, *projected, dropout_p=0.0) value = value.transpose(1, 2).reshape(batch, length, width) return F.linear(value, attention.out_proj.weight, attention.out_proj.bias) def forward(self, state, question, decision, options, *, prepared_memory=None): if state.ndim != 2 or question.ndim != 1 or decision.ndim != 1 or options.ndim != 2: raise ValueError( "Cross-attention readout expects state [T,D], question/decision [D], options [K,D]" ) if not len(state) or not len(options): raise ValueError("Cross-attention readout requires state and candidates") dtype = torch.bfloat16 if self.bf16_compute else self.memory_projection.weight.dtype with torch.autocast(device_type=state.device.type, dtype=torch.bfloat16, enabled=self.bf16_compute): state = state.to(dtype) question = question.to(dtype) decision = decision.to(dtype) options = options.to(dtype) memory, projected = ( (self.memory_projection(self.memory_norm(state)).unsqueeze(0), None) if prepared_memory is None else prepared_memory ) question_query = self.question_seed(question)[None, None] question_evidence = self._attend( self.question_attention, question_query, memory, None if projected is None else projected[0] )[0, 0] context = self.context_norm( self.decision_projection(decision) + self.question_projection(question) + self.question_evidence_projection(question_evidence) ) routed = self.option_norm(self.option_projection(options) + self.condition_projection(context)) for layer_index, (attention, attn_norm, ffn, output_norm) in enumerate( zip( self.routing_attention, self.routing_attention_norm, self.routing_ffn, self.routing_output_norm, strict=True, ) ): evidence = self._attend( attention, routed[None], memory, None if projected is None else projected[layer_index + 1] )[0] updated = attn_norm(routed + evidence) routed = output_norm(updated + ffn(updated)) query = self.query_norm(self.query_mlp(context)) keys = self.key_norm(self.key_mlp(routed)) with torch.autocast(device_type=state.device.type, enabled=False): return (keys.float() @ query.float()) / math.sqrt(self.config.projection_size) class FADepthRouter(nn.Module): def __init__(self, hidden_size, rms_norm_eps=1e-6): super().__init__() self.query = nn.Parameter(torch.zeros(hidden_size)) self.rms_norm_eps = rms_norm_eps def forward(self, carry, sources): """Return m * sum(alpha_j * v_j), with m equal to the source count. Centering the weights around 1/m uses the existing carry as the sum of sources. This equivalent form avoids regrouping low-precision residual additions at initialization: query=0 gives exactly r=carry. Neither keys nor values are detached; all depth paths retain their gradients. """ with torch.autocast(device_type=carry.device.type, enabled=False): values = torch.stack(sources).float() keys = values * torch.rsqrt(values.square().mean(-1, keepdim=True) + self.rms_norm_eps) logits = torch.einsum("d,nbtd->nbt", self.query.float(), keys) alpha = logits.softmax(dim=0) count = len(sources) centered = (alpha - 1.0 / count) * count correction = torch.einsum("nbt,nbtd->btd", centered, values) return (carry.float() + correction).to(carry.dtype) def _attention(module, hidden, positional, previous, *, bidirectional=False): batch, length = hidden.shape[:2] shape = (batch, length, -1, module.head_dim) query, gate = module.q_proj(hidden).view(batch, length, -1, 2 * module.head_dim).chunk(2, dim=-1) query = module.q_norm(query).transpose(1, 2) key = module.k_norm(module.k_proj(hidden).view(shape)).transpose(1, 2) value = module.v_proj(hidden).view(shape).transpose(1, 2) query, key = apply_rotary_pos_emb(query, key, *positional) prefix_length = previous[0].shape[2] if previous else 0 if previous: key = torch.cat((previous[0], key), dim=2) value = torch.cat((previous[1], value), dim=2) # Keep compact GQA cache tensors while avoiding a quadratic math fallback # in SDPA implementations that do not fuse grouped-query attention. groups = query.shape[1] // key.shape[1] attention_query = query.to(value.dtype) attention_key = key.to(value.dtype) attention_value = value if groups > 1: attention_key = attention_key.repeat_interleave(groups, dim=1) attention_value = attention_value.repeat_interleave(groups, dim=1) padding = (-module.head_dim) % 8 if padding: attention_query, attention_key, attention_value = ( F.pad(tensor, (0, padding)) for tensor in (attention_query, attention_key, attention_value) ) mask = None if prefix_length and not bidirectional: q_positions = torch.arange(length, device=hidden.device)[:, None] k_positions = torch.arange(prefix_length + length, device=hidden.device)[None, :] allowed = k_positions <= prefix_length + q_positions width = math.ceil((prefix_length + length) / 8) * 8 mask = attention_query.new_zeros((length, width))[:, : prefix_length + length] mask.masked_fill_(~allowed, float("-inf")) with ( torch.autocast(device_type=hidden.device.type, enabled=False), sdpa_kernel([SDPBackend.FLASH_ATTENTION, SDPBackend.EFFICIENT_ATTENTION]), ): attended = F.scaled_dot_product_attention( attention_query, attention_key, attention_value, attn_mask=mask, is_causal=not prefix_length and not bidirectional, dropout_p=0.0, scale=module.scaling, )[..., : module.head_dim] attended = attended.transpose(1, 2).reshape(batch, length, -1) output = module.o_proj(attended * gate.reshape(batch, length, -1).sigmoid()) return output, (key, value) def _linear_context(module, hidden, previous, config): if previous is None: cache = Qwen3_5DynamicCache(config) cache.last_linear_layer = module.layer_idx output = module(hidden, cache_params=cache) return output, (cache.conv_states[module.layer_idx], cache.recurrent_states[module.layer_idx]) batch, length = hidden.shape[:2] raw = module.in_proj_qkv(hidden).transpose(1, 2) window = torch.cat((previous[0], raw), dim=-1) # The cached window holds kernel_size raw values. Its oldest value precedes # the first suffix output's receptive field; exclude that value from conv. if module.causal_conv1d_fn is not None: # Match the upstream fused convolution+SiLU rounding in BF16. Applying # the same causal kernel to the cached window preserves continuation # gradients and avoids an extra low-precision convolution round trip. mixed = module.causal_conv1d_fn( x=window, weight=module.conv1d.weight.squeeze(1), bias=module.conv1d.bias, activation=module.activation, )[:, :, -length:] else: mixed = F.silu( F.conv1d(window[:, :, 1:], module.conv1d.weight, module.conv1d.bias, groups=raw.shape[1]) ) mixed = mixed.transpose(1, 2) query, key, value = mixed.split((module.key_dim, module.key_dim, module.value_dim), dim=-1) query = query.reshape(batch, length, -1, module.head_k_dim) key = key.reshape(batch, length, -1, module.head_k_dim) value = value.reshape(batch, length, -1, module.head_v_dim) repeat = module.num_v_heads // module.num_k_heads if repeat > 1: query = query.repeat_interleave(repeat, dim=2) key = key.repeat_interleave(repeat, dim=2) beta = module.in_proj_b(hidden).sigmoid() decay = -module.A_log.float().exp() * F.softplus(module.in_proj_a(hidden).float() + module.dt_bias) mixed, recurrent = module.chunk_gated_delta_rule( query, key, value, g=decay, beta=beta, initial_state=previous[1], output_final_state=True, use_qk_l2norm_in_kernel=True, ) gate = module.in_proj_z(hidden).reshape(-1, module.head_v_dim) mixed = module.norm(mixed.reshape(-1, module.head_v_dim), gate).reshape(batch, length, -1) return module.out_proj(mixed), (window[:, :, -module.conv_kernel_size :].clone(), recurrent) def _continue_gdn(module, hidden, previous, config, *, terminal=False, scan=None): """Continue on the native 64-token chunk grid, preserving terminal scans.""" scan = scan or _linear_context chunk_size = 64 if torch.is_grad_enabled() or module.training: raise ValueError("Chunk-aware prefix state is inference-only") if previous is not None and not isinstance(previous, GDNState): raise ValueError("GS1 continuation requires chunk-aware native prefix state") if terminal: # The trained GS1 reference scans context first, then the terminal token # from that final state. Do not merge the terminal back into the chunk. if previous is None or hidden.shape[1] != 1: raise ValueError("A terminal segment requires a prefix and exactly one token") return scan(module, hidden, (previous.conv, previous.recurrent), config) boundary = ( (previous.boundary_conv, previous.boundary_recurrent) if previous is not None and previous.boundary_conv is not None else None ) tail = previous.tail if previous is not None else hidden[:, :0] combined = torch.cat((tail, hidden), dim=1) if tail.shape[1] else hidden result = scan(module, combined, boundary, config) output, final = result[:2] exported = result[2] if len(result) == 3 else None completed = combined.shape[1] // chunk_size * chunk_size if completed: if completed == combined.shape[1]: boundary = final elif exported is not None: boundary = exported else: boundary = scan(module, combined[:, :completed], boundary, config)[1] state = GDNState( *final, *(boundary if boundary is not None else (None, None)), combined[:, completed:].clone(), (previous.position if previous is not None else 0) + hidden.shape[1], ) return output[:, tail.shape[1] :], state class GS1PretrainedModel(Qwen3_5PreTrainedModel): """Shared initialization and strict native serialization for GS1 modules.""" config_class = GS1Config base_model_prefix = "" _no_split_modules = ["GS1DecoderLayer", "Qwen3_5VisionBlock"] _tied_weights_keys = {} def get_output_embeddings(self): return None def gradient_checkpointing_enable(self, gradient_checkpointing_kwargs=None): for module in self.modules(): if isinstance(module, GS1TextModel): module.gradient_checkpointing = True def gradient_checkpointing_disable(self): for module in self.modules(): if isinstance(module, GS1TextModel): module.gradient_checkpointing = False def save_pretrained(self, save_directory, *, state_dict=None, **kwargs): from .compiler import save_chat_template self.config.decision = self.decision_config.model_dump() if hasattr(self, "calibrator"): self.config.calibration = self.calibrator.config.model_dump() if hasattr(self, "decision_head"): self.config.head_compute_dtype = "bfloat16" if self.decision_head.bf16_compute else "float32" kwargs.pop("safe_serialization", None) save_chat_template(save_directory) state = self.state_dict() if state_dict is None else state_dict return super().save_pretrained( save_directory, state_dict={name: value.detach().contiguous() for name, value in state.items()}, safe_serialization=True, **kwargs, ) @classmethod def from_pretrained( cls, source, *, device="cpu", dtype=None, config=None, revision=None, local_files_only=False, cache_dir=None, token=None, force_download=False, **kwargs, ): for name in ( "_from_auto", "_from_pipeline", "_commit_hash", "trust_remote_code", "code_revision", "weights_only", "use_safetensors", "adapter_kwargs", ): kwargs.pop(name, None) if "torch_dtype" in kwargs: if dtype is not None: raise ValueError("Use only dtype or torch_dtype") dtype = kwargs.pop("torch_dtype") device_map = kwargs.pop("device_map", None) if device_map is not None: if isinstance(device_map, dict) and set(device_map) == {""}: device = device_map[""] elif isinstance(device_map, str) and ( device_map in {"cpu", "cuda"} or device_map.startswith("cuda:") ): device = device_map else: raise ValueError( "GS1 supports a single device; use .to(device) instead of automatic device mapping" ) if kwargs: raise TypeError(f"Unsupported GS1 loading options: {sorted(kwargs)}") directory = Path(source).expanduser() if directory.is_dir(): directory = directory.resolve() else: from huggingface_hub import snapshot_download directory = Path( snapshot_download( str(source), revision=revision, local_files_only=local_files_only, cache_dir=cache_dir, token=token, force_download=force_download, allow_patterns=["*.json", "*.safetensors", "*.jinja", "*.py", "*.model", "*.txt"], ) ) raw = json.loads((directory / "config.json").read_text()) if raw.get("model_type") != "gs1": raise ValueError("Expected native GS1 weights") config = config or GS1Config(**raw) config._name_or_path = str(source) # Keep nonpersistent rotary buffers real while constructing empty parameters. with init_empty_weights(include_buffers=False): model = cls(config) state = _read_weights(directory) expected = model.state_dict() if state.keys() != expected.keys(): raise ValueError( f"Incomplete GS1 weights: missing={sorted(expected.keys() - state.keys())}, " f"unexpected={sorted(state.keys() - expected.keys())}" ) for name, value in state.items(): if value.shape != expected[name].shape: raise ValueError(f"GS1 parameter shape mismatch: {name}") model.load_state_dict(state, strict=True, assign=True) model._processor_source = directory if dtype is not None and dtype != "auto": if isinstance(dtype, str): dtype = getattr(torch, dtype) model.to(dtype=dtype) return model.to(device).eval() class GS1TextModel(GS1PretrainedModel): """Native hybrid text decoder with GS1 depth routing and segmented attention.""" def __init__(self, config): super().__init__(config) text = config.text_config self.decision_config = DecisionConfig.model_validate(config.decision) self.initialization_seed = config.initialization_seed self.full_layers = tuple(i for i, kind in enumerate(text.layer_types) if kind == "full_attention") if text.linear_num_value_heads % text.linear_num_key_heads: raise ValueError("GDN value heads must be a multiple of key heads") if text.num_attention_heads % text.num_key_value_heads: raise ValueError("Attention heads must be a multiple of key/value heads") for selected in (self.decision_config.attnres_layers, self.decision_config.bidirectional_layers): if selected is not None and not set(selected) <= set(self.full_layers): raise ValueError("Selected layers must be full-attention decoder indices") self.embed_tokens = nn.Embedding(text.vocab_size, text.hidden_size, text.pad_token_id) self.layers = nn.ModuleList(GS1DecoderLayer(text, i) for i in range(text.num_hidden_layers)) self.norm = Qwen3_5RMSNorm(text.hidden_size, eps=text.rms_norm_eps) self.rotary_emb = Qwen3_5TextRotaryEmbedding(text) self.gradient_checkpointing = False self.post_init() receivers = self.decision_config.attnres_layers or self.full_layers self.attnres_routers = nn.ModuleDict( { str(i): FADepthRouter(text.hidden_size, self.decision_config.attnres_eps) for i in receivers if self.decision_config.attnres } ) def get_input_embeddings(self): return self.embed_tokens def set_input_embeddings(self, value): self.embed_tokens = value def is_bidirectional(self, index): selected = self.decision_config.bidirectional_layers return self.decision_config.bidirectional and (selected is None or index in selected) def _depth_input(self, index, carry, bank, partial): if str(index) not in self.attnres_routers: return carry sources = bank if partial is None else (*bank, partial) return self.attnres_routers[str(index)](carry, sources) @staticmethod def _advance_depth(bank, partial, mixed, ffn, full): increment = mixed + ffn partial = increment if partial is None else partial + increment return ((*bank, partial), None) if full else (bank, partial) def prepare_encode(self, hidden, position_ids, roles=None, *, terminal_query=False, attention_mask=None): """Validate immutable batch geometry before tensor execution or CUDA capture.""" batch, length = hidden.shape[:2] if not batch or not length: raise ValueError("GS1 expects nonempty branch rows") valid = ( torch.ones((batch, length), dtype=torch.bool, device=hidden.device) if attention_mask is None else attention_mask.bool() ) if valid.shape != (batch, length) or not valid[:, -1].all() or (valid[:, :-1] & ~valid[:, 1:]).any(): raise ValueError("GS1 batches require nonempty left-padded rows") if terminal_query and (valid.sum(-1) < 2).any(): raise ValueError("A decision requires context before its terminal token") if roles is None: roles = torch.zeros((batch, length), dtype=torch.long, device=hidden.device) elif roles.shape == (length,): roles = roles[None].expand(batch, -1) if roles.shape != (batch, length): raise ValueError("Token roles must match the branch rows") keep = valid.unsqueeze(-1) valid_lengths = tuple(valid.sum(-1).tolist()) hidden = hidden * keep masks = {} eye = torch.eye(length, dtype=torch.bool, device=hidden.device)[None] for bidirectional in {self.is_bidirectional(i) for i in self.full_layers}: if bidirectional: allowed = (roles[:, :, None] >= roles[:, None, :]) & (roles[:, None, :] != 3) allowed |= eye else: allowed = torch.ones((batch, length, length), dtype=torch.bool, device=hidden.device).tril_() allowed &= valid[:, None, :] allowed |= ~valid[:, :, None] & eye masks[bidirectional] = hidden.new_zeros((batch, 1, length, length)).masked_fill_( ~allowed[:, None], torch.finfo(hidden.dtype).min ) positional = self.rotary_emb(hidden, position_ids) cohorts = tuple( ( length, torch.tensor( [i for i, value in enumerate(valid_lengths) if value == length], device=hidden.device ), ) for length in sorted(set(valid_lengths)) ) return keep, valid_lengths, masks, positional, cohorts def encode(self, hidden, position_ids, roles=None, *, terminal_query=False, attention_mask=None): layout = self.prepare_encode( hidden, position_ids, roles, terminal_query=terminal_query, attention_mask=attention_mask ) return self.encode_prepared(hidden, layout, terminal_query=terminal_query) def encode_prepared(self, hidden, layout, *, terminal_query=False, execution_policy=None): """The unchanged decoder equations over already validated, static geometry.""" if execution_policy is not None: return execution_policy.encode(self, hidden, layout, terminal_query=terminal_query) keep, valid_lengths, masks, positional, cohorts = layout hidden = hidden * keep bank, partial = (hidden,), None def run(carry, layer, index, bank, partial): read = self._depth_input(index, carry, bank, partial) mask = masks.get(self.is_bidirectional(index)) output, mixed, ffn = layer(carry, positional, mask, read, terminal_query, valid_lengths, cohorts) # Prefix PADs start from zero GDN state and remain zero after every # layer. Valid tokens retain the ordinary convolution/time scan. output, mixed, ffn = output * keep, mixed * keep, ffn * keep if self.attnres_routers: bank, partial = self._advance_depth( bank, partial, mixed, ffn, layer.layer_type == "full_attention" ) return output, bank, partial for index, layer in enumerate(self.layers): if self.training and self.gradient_checkpointing: hidden, bank, partial = checkpoint( run, hidden, layer, index, bank, partial, use_reentrant=False ) else: hidden, bank, partial = run(hidden, layer, index, bank, partial) return self.norm(hidden) * keep def segment_forward( self, input_ids, start_position=0, prefix=None, *, terminal_query=False, execution_policy=None ): """Inference-only segment continuation, including GDN convolution/state caches.""" if torch.is_grad_enabled() or self.training: raise ValueError("Shared segment caches require eval mode without gradients") if execution_policy is not None: return execution_policy.segment( self, input_ids, start_position, prefix, terminal_query=terminal_query ) hidden = self.embed_tokens(input_ids) positions = torch.arange(start_position, start_position + hidden.shape[1], device=hidden.device)[ None ].expand(hidden.shape[0], -1) positional = self.rotary_emb(hidden, positions) bank, partial, states = (hidden,), None, [] for index, layer in enumerate(self.layers): read = self._depth_input(index, hidden, bank, partial) normalized = layer.input_layernorm(read) previous = prefix[index] if prefix is not None else None if layer.layer_type == "full_attention": mixed, state = _attention( layer.self_attn, normalized, positional, previous, bidirectional=self.is_bidirectional(index), ) else: mixed, state = _continue_gdn( layer.linear_attn, normalized, previous, self.config.text_config, terminal=terminal_query ) hidden = hidden + mixed ffn = layer.mlp(layer.post_attention_layernorm(hidden)) hidden = hidden + ffn if self.attnres_routers: bank, partial = self._advance_depth( bank, partial, mixed, ffn, layer.layer_type == "full_attention" ) states.append(state) return self.norm(hidden), tuple(states) def forward( self, input_ids=None, *, inputs_embeds=None, position_ids=None, roles=None, terminal_query=False, attention_mask=None, prepared_layout=None, execution_policy=None, ): if (input_ids is None) == (inputs_embeds is None): raise ValueError("Specify exactly one of input_ids or inputs_embeds") hidden = self.embed_tokens(input_ids) if inputs_embeds is None else inputs_embeds if position_ids is None: position_ids = torch.arange(hidden.shape[1], device=hidden.device)[None].expand( hidden.shape[0], -1 ) output = ( self.encode( hidden, position_ids, roles, terminal_query=terminal_query, attention_mask=attention_mask ) if prepared_layout is None else self.encode_prepared( hidden, prepared_layout, terminal_query=terminal_query, execution_policy=execution_policy ) ) return GS1TextOutput(last_hidden_state=output) class GS1Model(GS1PretrainedModel): """Complete GS1: native text, upstream vision and one decision head.""" def __init__(self, config): super().__init__(config) self.initialization_seed = config.initialization_seed self.language_model = GS1TextModel(config) self.decision_config = self.language_model.decision_config self.visual = Qwen3_5VisionModel._from_config(config.vision_config) self.post_init() self.decision_head = CrossAttentionDecisionHead( config.text_config.hidden_size, self.decision_config, self.initialization_seed ) self.decision_head.bf16_compute = config.head_compute_dtype == "bfloat16" from .schema import Calibrator self.calibrator = Calibrator(config.calibration) def get_input_embeddings(self): return self.language_model.get_input_embeddings() def set_input_embeddings(self, value): self.language_model.set_input_embeddings(value) def _vision_positions(self, input_ids, types, image_grid, video_grid=None): positions = torch.zeros(3, *input_ids.shape, device=input_ids.device, dtype=torch.long) grids = { 1: iter(image_grid) if image_grid is not None else None, 2: iter(video_grid) if video_grid is not None else None, } merge = self.config.vision_config.spatial_merge_size for row in range(input_ids.shape[0]): begin, offset = 0, 0 while begin < input_ids.shape[1]: kind = int(types[row, begin]) end = begin + 1 while end < input_ids.shape[1] and int(types[row, end]) == kind: end += 1 if kind == 0: piece = torch.arange(end - begin, device=input_ids.device)[None].expand(3, -1) + offset offset += end - begin else: t, h, w = map(int, next(grids[kind])) h, w = h // merge, w // merge piece = torch.stack( ( torch.full((t * h * w,), offset, device=input_ids.device), torch.arange(offset, offset + h, device=input_ids.device).repeat_interleave( w * t ), torch.arange(offset, offset + w, device=input_ids.device).repeat(h * t), ) ) offset += max(h, w) if piece.shape[1] != end - begin: raise ValueError("Vision grid and placeholder token count differ") positions[:, row, begin:end] = piece begin = end return positions def forward( self, input_ids, *, branch=None, state_length=None, branches=None, state_lengths=None, position_ids=None, pixel_values=None, image_grid_thw=None, mm_token_type_ids=None, attention_mask=None, readout_callback=None, **kwargs, ): hidden = self.get_input_embeddings()(input_ids) if pixel_values is not None: features = self.visual( pixel_values.to(self.visual.dtype), grid_thw=image_grid_thw, return_dict=True ).pooler_output.to(hidden.device, hidden.dtype) mask = (input_ids == self.config.image_token_id).unsqueeze(-1).expand_as(hidden) if hidden[mask].numel() != features.numel(): raise ValueError("Image features and placeholder tokens differ") hidden = hidden.masked_scatter(mask, features) if position_ids is None: if mm_token_type_ids is None: raise ValueError("Visual branches require modality token types") position_ids = self._vision_positions(input_ids, mm_token_type_ids, image_grid_thw) if position_ids is None: position_ids = torch.arange(hidden.shape[1], device=hidden.device)[None].expand( hidden.shape[0], -1 ) if branch is not None and branches is not None: raise ValueError("Specify one branch or batched branch metadata") rows = branches if branches is not None else [branch] if branch is not None else None lengths = state_lengths if branches is not None else [state_length] if branch is not None else None roles = None starts = [] if rows is not None: if len(rows) != input_ids.shape[0] or lengths is None or len(lengths) != len(rows): raise ValueError("Compiled branch metadata differs from batch rows") roles = torch.full_like(input_ids, -1) for row, (item, offset) in enumerate(zip(rows, lengths, strict=True)): total = offset + len(item.suffix) begin = input_ids.shape[1] - total if begin < 0 or (attention_mask is not None and int(attention_mask[row].sum()) != total): raise ValueError("Compiled branch length differs from model input") if branch is not None and begin: raise ValueError("Single compiled branches must be unpadded") starts.append(begin) roles[row, begin : begin + offset] = 0 roles[row, begin + offset :] = torch.tensor( [ 1 if r == "question_context" else 2 if r == "candidate_context" else 3 for r in item.token_roles ], device=hidden.device, ) hidden = self.language_model( inputs_embeds=hidden, position_ids=position_ids, roles=roles, terminal_query=rows is not None, attention_mask=attention_mask, ).last_hidden_state logits = None if rows is not None: values = [ (readout_callback or self.readout)(hidden[row, begin:], item, offset) for row, (begin, item, offset) in enumerate(zip(starts, rows, lengths, strict=True)) ] width = max(value.numel() for value in values) logits = ( values[0] if branch is not None else torch.stack( [F.pad(value, (0, width - value.numel()), value=-torch.inf) for value in values] ) ) return GS1Output(logits=logits, last_hidden_state=hidden) def readout(self, hidden, branch, offset, *, prepared_memory=None, state_hidden=None): def pool(start, span, end): content = hidden[offset + span[0] : offset + span[1]].mean(0) closing = hidden[offset + end] return ( (hidden[offset + start] + content + closing) / 3 if self.decision_config.bidirectional else (content + closing) / 2 ) question = pool(branch.question_start_position, branch.question_span, branch.question_end_position) options = torch.stack( [ pool(start, span, end) for start, span, end in zip( branch.option_starts, branch.option_spans, branch.option_ends, strict=True ) ] ) return self.decision_head( hidden[branch.state_evidence_start : offset] if state_hidden is None else state_hidden, question, hidden[offset + branch.decision_position], options, prepared_memory=prepared_memory, ) def probabilities(self, logits, candidate_mask=None, *, calibration_active=True): from .schema import probabilities return probabilities(logits, candidate_mask, self.calibrator if calibration_active else None) def decision_compiler(self, tokenizer=None, **overrides): """Build the native prompt compiler from standard pretrained assets.""" from .compiler import Compiler if tokenizer is None: from transformers import AutoTokenizer if not hasattr(self, "_decision_tokenizer"): self._decision_tokenizer = AutoTokenizer.from_pretrained( self._processor_source, local_files_only=True ) tokenizer = self._decision_tokenizer source = getattr(self, "_processor_source", None) limits = { key: self.config.inference_limits[key] for key in ("max_context", "max_questions", "max_request_tokens", "max_branches") } for key, value in overrides.items(): limits[key] = min(limits[key], value) compiler = Compiler(tokenizer, decision=self.decision_config, processor=source, **limits) if compiler.controls != getattr(self.config, "structural_token_ids", {}): raise ValueError("Structural token identity mismatch") return compiler @torch.inference_mode() def decide(self, state, questions, *, tokenizer=None, precision="bf16"): """Return typed decisions using the standalone eager reference path. This method needs only the files in this model's Hub repository. The separately installed gs1 runtime provides qualified CUDA graph execution. """ from .compiler import branch_logits from .schema import DecisionRequest, DecisionResponse, Usage, make_answer if precision not in {"fp32", "bf16"}: raise ValueError("precision must be fp32 or bf16") if self.training: raise ValueError("Call model.eval() before making decisions") compiler = self.decision_compiler(tokenizer) tokenizer = compiler.tokenizer name = getattr(self.config, "_name_or_path", "") or "GS1" request = DecisionRequest(state=state, questions=questions, model=name) compiled = compiler.compile(request) device = self.get_input_embeddings().weight.device answers = {} with torch.autocast(device.type, dtype=torch.bfloat16, enabled=precision == "bf16"): for branch in compiled.branches: logits = branch_logits( self, [(compiled.state, branch)], pad_token_id=tokenizer.pad_token_id or 0 )[0] values = self.probabilities(logits).p_final.tolist() answers[branch.question_id] = make_answer(request.questions[branch.question_id], values) response = DecisionResponse( model=name, answers=answers, usage=Usage(input_tokens=compiled.input_tokens, output_tokens=0) ) return response.model_dump(mode="json") # Native discovery resolves the complete text/vision/decision model. AutoConfig.register("gs1", GS1Config, exist_ok=True) AutoModel.register(GS1Config, GS1Model, exist_ok=True) GS1Config.register_for_auto_class("AutoConfig") GS1Model.register_for_auto_class("AutoModel")