"""Streaming FP32 text-only Qwen3.5 reference, original BF16 weights. Mathematical semantics checked against Transformers Qwen3.5, commit fd3cf0140de03f7619ac95b5cfed05b711ff9e22 (Apache-2.0, Qwen Team/HuggingFace). No cached decode, padding, images, generation, fused kernels or autograd. Full-sequence gated delta recurrence replaces the algebraically equivalent chunk algorithm; both are independently compared in remote tests. """ import ctypes import json import time from pathlib import Path import numpy as np from pilot.tensors import load_tensors, decode PREFIX = 'model.language_model.' def sigmoid(x): return np.exp(-np.logaddexp(np.float32(0), -x)) def silu(x): return x * sigmoid(x) def norm(x, weight, eps=1e-6, centered=True): return (x / np.sqrt(np.mean(x*x, axis=-1, keepdims=True) + eps) * (1 + weight if centered else weight)).astype(np.float32) def rope(x, theta=10000000., fraction=.25): # Text-only: all three MRoPE position axes coincide. dim = int(x.shape[-1] * fraction) inv = np.float32(1) / np.float32(theta) ** (np.arange(0, dim, 2, dtype=np.float32) / dim) freq = np.arange(x.shape[1], dtype=np.float32)[:, None] * inv[None, :] angle = np.concatenate([freq, freq], axis=-1)[None, :, None, :] a = x[..., :dim] rotated = np.concatenate([-a[..., dim//2:], a[..., :dim//2]], axis=-1) return np.concatenate([a*np.cos(angle) + rotated*np.sin(angle), x[..., dim:]], axis=-1) def delta_rule(q, k, v, g, beta): q = q / np.sqrt(np.sum(q*q, axis=-1, keepdims=True) + 1e-6) k = k / np.sqrt(np.sum(k*k, axis=-1, keepdims=True) + 1e-6) q = q / np.float32(np.sqrt(q.shape[-1])) state = np.zeros((q.shape[0], q.shape[2], q.shape[3], v.shape[3]), dtype=np.float32) out = np.empty_like(v) for t in range(q.shape[1]): state *= np.exp(g[:, t, :, None, None]) predicted = np.sum(state * k[:, t, :, :, None], axis=-2) delta = (v[:, t] - predicted) * beta[:, t, :, None] state += k[:, t, :, :, None] * delta[:, :, None, :] out[:, t] = np.sum(state * q[:, t, :, :, None], axis=-2) return out, state def causal_conv(x, kernel): out = np.zeros_like(x) width = kernel.shape[-1] for j in range(width): lag = width - 1 - j if lag < x.shape[1]: out[:, lag:] += x[:, :x.shape[1]-lag] * kernel[:, j] return silu(out) def attention(q, k, v): repeats = q.shape[2] // k.shape[2] k, v = np.repeat(k, repeats, axis=2), np.repeat(v, repeats, axis=2) scores = np.einsum('bthd,bshd->bhts', q, k) / np.float32(np.sqrt(q.shape[-1])) mask = np.triu(np.ones(scores.shape[-2:], dtype=bool), 1) scores = np.where(mask, -np.inf, scores) exp = np.exp(scores - scores.max(axis=-1, keepdims=True)) p = exp / exp.sum(axis=-1, keepdims=True) return np.einsum('bhts,bshd->bthd', p, v), p class StreamingLM: def __init__(self, weights, config, seconds=7200, enforce_memory=True): self.tensors = load_tensors(Path(weights)) self.config = config.get('text_config', config) self.deadline = time.monotonic() + seconds self.enforce_memory = enforce_memory self.eps = self.config['rms_norm_eps'] self.norm_replacement = None self.capture_qk = False self.qk_inputs = {} self.swap_observation = None self.swap_pressure_windows = 0 self.malloc_trim = getattr(ctypes.CDLL(None), 'malloc_trim', None) def guard(self): if time.monotonic() > self.deadline: raise TimeoutError('CPU computation deadline reached') if not self.enforce_memory: return status = dict(line.split(':', 1) for line in Path('/proc/self/status').read_text().splitlines() if ':' in line) rss_kib = int(status['VmRSS'].split()[0]) available = int(next(line.split()[1] for line in Path('/proc/meminfo').read_text().splitlines() if line.startswith('MemAvailable:'))) now = time.monotonic() if self.swap_observation is None or now-self.swap_observation[0] >= 5: counters = dict(line.split() for line in Path('/proc/vmstat').read_text().splitlines()) count = int(counters['pswpin'])+int(counters['pswpout']) if self.swap_observation is not None: elapsed = now-self.swap_observation[0] rate = (count-self.swap_observation[1])*4096/elapsed self.swap_pressure_windows = self.swap_pressure_windows+1 if rate > 4*1024**2 else 0 own_swap_kib = int(status.get('VmSwap', '0 kB').split()[0]) if self.swap_pressure_windows >= 3 and (available < 384*1024 or own_swap_kib > 128*1024): raise MemoryError(f'Sustained host swap pressure plus constrained memory: available={available} KiB, process_swap={own_swap_kib} KiB') self.swap_observation = (now,count) if rss_kib > 512*1024 or available < 256*1024: raise MemoryError(f'Soft memory stop: RSS {rss_kib} KiB, available {available} KiB') def small(self, name): t = self.tensors[PREFIX + name] if t.numel > 1048576: raise ValueError('Use blockwise linear for large tensors') x = t.read().reshape(t.shape) t.evict() return x def linear(self, x, name, rows=256): t = self.tensors[PREFIX + name + '.weight'] outdim, indim = t.shape flat = x.reshape(-1, indim) result = np.empty((flat.shape[0], outdim), dtype=np.float32) for start in range(0, outdim, rows): self.guard() end = min(start + rows, outdim) w = t.read(slice(start*indim, end*indim)).reshape(end-start, indim) result[:, start:end] = flat @ w.T del w t.evict() if self.malloc_trim is not None: self.malloc_trim(0) # Release only this process's freed temporary heap pages. return result.reshape(*x.shape[:-1], outdim) def embedding(self, ids): t = self.tensors[PREFIX + 'embed_tokens.weight'] raw = t.raw.reshape(t.shape) result = decode(raw[ids], t.dtype) t.evict() return result def normalization(self, x, name): if getattr(self, 'norm_replacement', None) and name == self.norm_replacement[0]: name = self.norm_replacement[1] return norm(x, self.small(name + '.weight'), self.eps) def mlp(self, h, layer): p = f'layers.{layer}.mlp.' gate = silu(self.linear(h, p+'gate_proj')) up = self.linear(h, p+'up_proj') return self.linear(gate*up, p+'down_proj') def full_attention(self, h, layer, shared_p=None): p = f'layers.{layer}.self_attn.' c = self.config heads, kvheads, dim = c['num_attention_heads'], c['num_key_value_heads'], c['head_dim'] # Q and output gate are interleaved per head, not globally split. mixed = self.linear(h, p+'q_proj').reshape(*h.shape[:-1], heads, 2*dim) q, gate = mixed[..., :dim], mixed[..., dim:] k = self.linear(h, p+'k_proj').reshape(*h.shape[:-1], kvheads, dim) v = self.linear(h, p+'v_proj').reshape(*h.shape[:-1], kvheads, dim) if getattr(self, 'capture_qk', False): self.qk_inputs[layer] = (q.copy(), k.copy()) q = self.normalization(q, p+'q_norm') k = self.normalization(k, p+'k_norm') rp = c['rope_parameters'] q, k = (rope(x, rp['rope_theta'], rp['partial_rotary_factor']) for x in (q,k)) out, probs = attention(q,k,v) if shared_p is not None: out = np.einsum('bhts,bshd->bthd', shared_p, np.repeat(v, heads//kvheads, axis=2)) out = (out*sigmoid(gate)).reshape(*h.shape[:-1], -1) return self.linear(out, p+'o_proj'), probs def linear_attention(self, h, layer): p = f'layers.{layer}.linear_attn.' c = self.config kh, vh = c['linear_num_key_heads'], c['linear_num_value_heads'] kd, vd = c['linear_key_head_dim'], c['linear_value_head_dim'] qkv = self.linear(h, p+'in_proj_qkv') qkv = causal_conv(qkv, self.small(p+'conv1d.weight')[:,0,:]) q, k, v = np.split(qkv, [kh*kd, 2*kh*kd], axis=-1) q = q.reshape(*h.shape[:-1], kh, kd) k = k.reshape(*h.shape[:-1], kh, kd) v = v.reshape(*h.shape[:-1], vh, vd) if vh != kh: q, k = np.repeat(q, vh//kh, axis=2), np.repeat(k, vh//kh, axis=2) z = self.linear(h, p+'in_proj_z').reshape(*h.shape[:-1], vh, vd) beta = sigmoid(self.linear(h, p+'in_proj_b')) a = self.linear(h, p+'in_proj_a') g = -np.exp(self.small(p+'A_log')) * np.logaddexp(np.float32(0), a+self.small(p+'dt_bias')) out, _ = delta_rule(q,k,v,g,beta) out = norm(out, self.small(p+'norm.weight'), self.eps, centered=False)*silu(z) return self.linear(out.reshape(*h.shape[:-1], -1), p+'out_proj') def forward(self, ids, capture=None, replacement=None, progress=None): h = self.embedding(ids) saved = {} capture = set(capture or []) for layer, kind in enumerate(self.config['layer_types']): self.guard() raw = h n = self.normalization(h, f'layers.{layer}.input_layernorm') if ('attn', layer) in capture: saved[f'attn_{layer}'] = n.copy() donor = replacement[2] if replacement and replacement[:2] == ('attn',layer) else layer if kind == 'full_attention': out, _ = self.full_attention(n,donor) else: out = self.linear_attention(n,donor) h = raw + out n = self.normalization(h, f'layers.{layer}.post_attention_layernorm') if ('mlp',layer) in capture: saved[f'mlp_{layer}'] = n.copy() donor = replacement[2] if replacement and replacement[:2] == ('mlp',layer) else layer h = h + self.mlp(n,donor) if not np.isfinite(h).all(): raise FloatingPointError(f'Nonfinite hidden states at layer {layer}') if progress: progress(layer) return self.normalization(h,'norm'), saved def nll(self, h, ids, rows=512): # Tied output embedding; online logsumexp avoids full vocabulary logits. x = h[:,:-1].reshape(-1,h.shape[-1]) targets = ids[:,1:].reshape(-1) t = self.tensors[PREFIX+'embed_tokens.weight'] mx = np.full(x.shape[0],-np.inf,dtype=np.float64) total = np.zeros(x.shape[0],dtype=np.float64) correct = np.empty(x.shape[0],dtype=np.float64) for start in range(0,t.shape[0],rows): self.guard() end=min(start+rows,t.shape[0]) w=t.read(slice(start*t.shape[1],end*t.shape[1])).reshape(end-start,t.shape[1]) scores=x@w.T m=np.maximum(mx,scores.max(axis=-1)) total=total*np.exp(mx-m)+np.exp(scores-m[:,None]).sum(axis=-1) mx=m which=(targets>=start)&(targets