"""transformers remote code for the exp_train hybrid LM checkpoints (trust_remote_code=True). tok = AutoTokenizer.from_pretrained(path) model = AutoModelForCausalLM.from_pretrained(path, trust_remote_code=True, dtype=torch.bfloat16, device_map="cuda") out = model.generate(**tok("The capital of France is", return_tensors="pt").to("cuda"), max_new_tokens=5) Architecture (exp_train / a100stack): 32 layers, 24 Gated DeltaNet + 8 gated GQA softmax attention (every 4th), SwiGLU MLPs, gated residuals, Delta Block attention residuals ("repo" variant) and a hashed bigram/trigram memory (ngram.safetensors, canonical-token hashing, +-1 sign trick) injected before blocks 2 and 16. Parameter names are the training checkpoint's, so model.safetensors loads unchanged; the n-gram table is read from ngram.safetensors after loading. Matrices run in BF16, norms / gates / the residual stream in FP32 as in training. Inference only, batch size 1. Prefill uses FLA's chunked gated-delta-rule kernel; decoding uses the fused recurrent kernel, a preallocated KV cache, and a CUDA graph for the single-token step. """ import hashlib import json import math import torch import torch.nn as nn import torch.nn.functional as F from fla.ops.gated_delta_rule import chunk_gated_delta_rule, fused_recurrent_gated_delta_rule from transformers import GenerationMixin, PretrainedConfig, PreTrainedModel from transformers.modeling_outputs import CausalLMOutputWithPast class HybridLMConfig(PretrainedConfig): model_type = "exp_train_hybrid" def __init__(self, model_config=None, modifications=None, kv_cache_len=4096, use_cuda_graph=True, bos_token_id=2, eos_token_id=3, pad_token_id=0, tie_word_embeddings=False, **kwargs): self.model_config = model_config or {} self.modifications = modifications or {} self.kv_cache_len = kv_cache_len # KV-cache capacity (prompt + generated tokens) self.use_cuda_graph = use_cuda_graph # capture the single-token decode step in a CUDA graph mc = self.model_config if mc: self.vocab_size = mc["vocab_size"] self.hidden_size = mc["d_model"] self.num_hidden_layers = mc["n_layers"] self.num_attention_heads = mc["n_heads"] self.num_key_value_heads = mc["n_kv_heads"] self.head_dim = mc["head_dim"] self.intermediate_size = mc["ffn_dim"] super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, pad_token_id=pad_token_id, tie_word_embeddings=tie_word_embeddings, **kwargs) # ----------------------------------------------------------------------------- n-gram hashing (ngram.py) _MULTS = [0x9E3779B97F4A7C15, 0xC2B2AE3D27D4EB4F, 0x165667B19E3779F9, 0xD6E8FEB86659FD93, 0xFF51AFD7ED558CCD, 0xC4CEB9FE1A85EC53, 0x94D049BB133111EB, 0xBF58476D1CE4E5B9] _SIGN_MULTS = [0xA0761D6478BD642F, 0xE7037ED1A0B428DB, 0x8EBC6AF09C88C6E3, 0x589965CC75374CC3, 0x1D8E4E27C47D124F, 0x2D358DCCAA6C78A5, 0x9FB21C651E98DF25, 0xD3833E804F4C574B] def _i64(x): return x - (1 << 64) if x >= (1 << 63) else x def _shift(ids, j): # tok[t] = ids[t - j] return torch.cat([ids.new_zeros(j), ids[: ids.numel() - j]]) if j else ids def build_canon_map(tokenizer_file): """ngram_canon.py: token ids whose text normalizes to the same string share n-gram rows (hashing only).""" from tokenizers import Regex, Tokenizer, normalizers tokenizer = Tokenizer.from_file(str(tokenizer_file)) norm = normalizers.Sequence([ normalizers.NFKC(), normalizers.NFD(), normalizers.StripAccents(), normalizers.Lowercase(), normalizers.Replace(Regex(r"[ \t\r\n]+"), " "), normalizers.Replace(Regex(r"^ $"), ""), normalizers.Strip(), normalizers.Replace("", " ")]) vocab = tokenizer.get_vocab_size(with_added_tokens=True) special = {t["id"] for t in json.loads(tokenizer.to_str()).get("added_tokens", []) if t.get("special")} first, canon = {}, torch.arange(vocab, dtype=torch.int64) for t in range(vocab): if t in special: continue text = tokenizer.decode([t], skip_special_tokens=False) if "�" in text: continue key = norm.normalize_str(text) if key: canon[t] = first.setdefault(key, t) return canon def _rms(x, w, eps): return F.rms_norm(x.float(), (x.shape[-1],), w.float(), eps) def _select(sources, qw, eps): """Delta Block: sum_i softmax_i(q . RMSNorm_w(s_i)) s_i (qw = query * weight).""" logits = torch.stack([(s * qw).sum(-1) * torch.rsqrt(s.pow(2).mean(-1) + eps) for s in sources]) p = logits.softmax(0) out = p[0].unsqueeze(-1) * sources[0] for pi, s in zip(p[1:], sources[1:]): out = out + pi.unsqueeze(-1) * s return out # ----------------------------------------------------------------------------- modules (checkpoint names) def _linear(i, o): return nn.Linear(i, o, bias=False) class _Weight(nn.Module): def __init__(self, *shape): super().__init__() self.weight = nn.Parameter(torch.ones(*shape)) class _Gate(nn.Module): def __init__(self, d): super().__init__() self.a = nn.Parameter(torch.zeros(d)) class GatedAttention(nn.Module): def __init__(self, c): super().__init__() H, KV, hd, d = c["n_heads"], c["n_kv_heads"], c["head_dim"], c["d_model"] self.q_proj, self.gate_proj, self.o_proj = _linear(d, H * hd), _linear(d, H * hd), _linear(H * hd, d) self.k_proj, self.v_proj = _linear(d, KV * hd), _linear(d, KV * hd) self.q_norm, self.k_norm = _Weight(hd), _Weight(hd) class GatedDeltaNet(nn.Module): def __init__(self, c): super().__init__() H, dk, d = c["gdn_heads"], c["gdn_head_dim"], c["d_model"] dv, K = int(dk * c["gdn_expand_v"]), c["gdn_conv_size"] self.q_proj, self.k_proj = _linear(d, H * dk), _linear(d, H * dk) self.v_proj, self.g_proj, self.o_proj = _linear(d, H * dv), _linear(d, H * dv), _linear(H * dv, d) self.a_proj, self.b_proj = _linear(d, H), _linear(d, H) self.A_log, self.dt_bias = nn.Parameter(torch.zeros(H)), nn.Parameter(torch.zeros(H)) self.q_conv1d, self.k_conv1d, self.v_conv1d = _Weight(H * dk, 1, K), _Weight(H * dk, 1, K), _Weight(H * dv, 1, K) self.o_norm = _Weight(dv) class SwiGLU(nn.Module): def __init__(self, d, f): super().__init__() self.w_gate, self.w_up, self.w_down = _linear(d, f), _linear(d, f), _linear(f, d) class Block(nn.Module): def __init__(self, c, i, has_attn_sources): super().__init__() self.is_attn = (i + 1) % c["full_attn_every"] == 0 d = c["d_model"] self.norm1, self.norm2 = _Weight(d), _Weight(d) self.mixer = GatedAttention(c) if self.is_attn else GatedDeltaNet(c) self.mlp = SwiGLU(d, c["ffn_dim"]) self.res1, self.res2 = _Gate(d), _Gate(d) self.delta_mlp_query, self.delta_mlp_weight = nn.Parameter(torch.zeros(d)), nn.Parameter(torch.ones(d)) if has_attn_sources: self.delta_attn_query, self.delta_attn_weight = nn.Parameter(torch.zeros(d)), nn.Parameter(torch.ones(d)) class NgramMemory(nn.Module): def __init__(self, ng, d_model): super().__init__() T, D, n = len(ng["orders"]) * ng["heads_per_order"], ng["dim"], len(ng["inject_layers"]) self.proj = nn.ModuleList([_linear(T * D, d_model) for _ in range(n)]) self.key_norm = nn.ModuleList([_Weight(d_model) for _ in range(n)]) self.val_norm = nn.ModuleList([_Weight(d_model) for _ in range(n)]) self.register_buffer("sign_pool", torch.zeros(ng["sign_pool"], D), persistent=True) # table [T * rows, D] (BF16) and the canonical-id map are attached by from_pretrained (not in model.safetensors) class _State: """Static decode buffers (one active sequence). Fixed addresses so the decode step can be a CUDA graph.""" def __init__(self, m, max_len, device): self.max_len, self.device = max_len, device C = 2 * m.GH * m.GK + m.GH * m.GV self.conv = [None if a else torch.zeros(C, m.CONV - 1, device=device) for a in m.is_attn] self.S = [None if a else torch.zeros(1, m.GH, m.GK, m.GV, device=device) for a in m.is_attn] self.k = [torch.zeros(1, m.KV, max_len, m.HD, device=device, dtype=m.mm_dtype) if a else None for a in m.is_attn] self.v = [torch.zeros(1, m.KV, max_len, m.HD, device=device, dtype=m.mm_dtype) if a else None for a in m.is_attn] self.ctx = torch.zeros(3, dtype=torch.int64, device=device) # last 3 canonical ids (trigram window) self.pos = torch.zeros(1, dtype=torch.int64, device=device) # next position self.tok = torch.zeros(1, dtype=torch.int64, device=device) # decode-step input self.graph = self.logits = None self.owner = None def recurrent(self): return [t for t in self.conv + self.S + [self.ctx, self.pos] if t is not None] class HybridCache: """What generate() passes back as past_key_values. A new prefill takes over the model's static buffers.""" is_compileable = False def __init__(self, state): self.state, self.seen = state, 0 state.owner = self def get_seq_length(self, layer_idx=0): return self.seen # ----------------------------------------------------------------------------- model class HybridLMForCausalLM(PreTrainedModel, GenerationMixin): config_class = HybridLMConfig base_model_prefix = "hybrid" _no_split_modules = ["Block"] _is_stateful = True _supports_sdpa = True # kept in FP32 even with dtype=torch.bfloat16 (training's residual stream / norms / gates were FP32) _keep_in_fp32_modules_strict = ["norm", "res1", "res2", "delta_", "A_log", "dt_bias", "embed", "sign_pool", "conv1d"] def __init__(self, config): super().__init__(config) c, mods = config.model_config, config.modifications assert mods.get("delta_block") and mods.get("delta_variant", "repo") == "repo" assert mods.get("depth_attn", "off") == "off" self.D, self.EPS = c["d_model"], c["norm_eps"] self.H, self.KV, self.HD = c["n_heads"], c["n_kv_heads"], c["head_dim"] self.GH, self.GK = c["gdn_heads"], c["gdn_head_dim"] self.GV, self.CONV = int(self.GK * c["gdn_expand_v"]), c["gdn_conv_size"] n = c["n_layers"] self.per_block = math.ceil(n / mods.get("delta_blocks", 4)) self.is_attn = [(i + 1) % c["full_attn_every"] == 0 for i in range(n)] ng = c["ngram"] self.ng = ng self.rows = max(1024, ng["total_params"] // (len(ng["orders"]) * ng["heads_per_order"] * ng["dim"])) self.inject = {l: k for k, l in enumerate(ng["inject_layers"])} self.embed = nn.Embedding(c["vocab_size"], c["d_model"]) # every layer after the first sees at least one Delta Block source (the embedding) at its mixer input self.blocks = nn.ModuleList([Block(c, i, i > 0) for i in range(n)]) self.norm = _Weight(c["d_model"]) self.lm_head = _linear(c["d_model"], c["vocab_size"]) self.ngram = NgramMemory(ng, c["d_model"]) self._fused_ptr, self._state = None, None self.post_init() def _init_weights(self, module): pass # everything comes from the checkpoint def get_input_embeddings(self): return self.embed def get_output_embeddings(self): return self.lm_head @classmethod def _supports_default_dynamic_cache(cls): return False # HybridCache is created by forward() # ------------------------------------------------------------------ loading @classmethod def from_pretrained(cls, pretrained_model_name_or_path, *args, **kwargs): model = super().from_pretrained(pretrained_model_name_or_path, *args, **kwargs) from safetensors import safe_open from transformers.utils import cached_file hub = {k: kwargs[k] for k in ("revision", "token", "cache_dir", "local_files_only") if k in kwargs} device = model.embed.weight.device with safe_open(cached_file(pretrained_model_name_or_path, "ngram.safetensors", **hub), "pt", device="cpu") as f: table = f.get_tensor("table") expected = (len(model.ng["orders"]) * model.ng["heads_per_order"] * model.rows, model.ng["dim"]) assert tuple(table.shape) == expected, (table.shape, expected) canon = build_canon_map(cached_file(pretrained_model_name_or_path, "tokenizer.json", **hub)) sha = model.config.modifications.get("ngram_canonical_sha256") if model.ng.get("canonical_tokens"): assert sha is None or hashlib.sha256(canon.numpy().tobytes()).hexdigest() == sha, \ "canonical-token map differs from the one used in training" else: canon = torch.arange(model.config.vocab_size) model.ngram.register_buffer("table", table.to(device), persistent=False) model.ngram.register_buffer("canon", canon.to(device), persistent=False) model.eval() return model @property def mm_dtype(self): return self.lm_head.weight.dtype def _ready(self): """Fuse projections into single GEMMs (the original parameters become views into the fused weight).""" if self._fused_ptr == self.blocks[0].mixer.q_proj.weight.data_ptr(): return with torch.no_grad(): for blk in self.blocks: m = blk.mixer names = (("q_proj", "k_proj", "v_proj", "gate_proj") if blk.is_attn else ("q_proj", "k_proj", "v_proj", "a_proj", "b_proj", "g_proj")) blk._in_w = self._fuse([getattr(m, n) for n in names]) blk._mlp_w = self._fuse([blk.mlp.w_gate, blk.mlp.w_up]) if not blk.is_attn: blk._conv_w = torch.cat([m.q_conv1d.weight[:, 0], m.k_conv1d.weight[:, 0], m.v_conv1d.weight[:, 0]]).float() blk._neg_A = -m.A_log.float().exp() blk._g1, blk._g2 = 2 * torch.sigmoid(blk.res1.a.float()), 2 * torch.sigmoid(blk.res2.a.float()) blk._dmq = (blk.delta_mlp_query * blk.delta_mlp_weight).float() if hasattr(blk, "delta_attn_query"): blk._daq = (blk.delta_attn_query * blk.delta_attn_weight).float() self._fused_ptr = self.blocks[0].mixer.q_proj.weight.data_ptr() self._state = None @staticmethod def _fuse(linears): w = torch.cat([l.weight.data for l in linears]) o = 0 for l in linears: n = l.weight.shape[0] l.weight.data = w[o:o + n] o += n return w def _get_state(self): dev = self.embed.weight.device if self._state is None or self._state.device != dev: self._state = _State(self, self.config.kv_cache_len, dev) return self._state # ------------------------------------------------------------------ pieces def _ngram_memory(self, cids, pos): ng, rows, signs, valid = self.ng, [], [], [] t = 0 for order in ng["orders"]: for head in range(ng["heads_per_order"]): h = torch.zeros_like(cids) m = _i64(_MULTS[(head * 3 + order) % 8]) for j in range(order): h = h * m + (_shift(cids, j) + 1 + j * 7919) h = h ^ (h >> 29); h = h * _i64(0xBF58476D1CE4E5B9); h = h ^ (h >> 32) rows.append(torch.remainder(h, self.rows) + t * self.rows) if ng.get("sign_pool"): s = torch.full_like(cids, _i64(0x632BE59BD9B4E019)) m = _i64(_SIGN_MULTS[(head * 3 + order) % 8]) for j in range(order): s = s * m + (_shift(cids, j) + 1 + j * 104729) s = s ^ (s >> 31); s = s * _i64(0x94D049BB133111EB); s = s ^ (s >> 29) signs.append(s & (ng["sign_pool"] - 1)) valid.append(pos >= order - 1) t += 1 emb = self.ngram.table[torch.stack(rows, 1)].float() # [N, T, dim] mult = torch.stack(valid, 1).unsqueeze(-1).float() if signs: mult = self.ngram.sign_pool[torch.stack(signs, 1)] * mult return (emb * mult).flatten(1) def _rope(self, x, pos): d = x.shape[-1] inv = 1.0 / (self.config.model_config["rope_theta"] ** (torch.arange(0, d, 2, device=x.device, dtype=torch.float32) / d)) f = pos.float()[:, None] * inv[None] cos, sin = f.cos()[:, None], f.sin()[:, None] x1, x2 = x[..., : d // 2], x[..., d // 2:] return torch.cat([x1 * cos - x2 * sin, x2 * cos + x1 * sin], -1) def _gdn(self, blk, i, xn, st, decode): m, N, H, dk, dv = blk.mixer, xn.shape[0], self.GH, self.GK, self.GV qkv, a, b, gate = (xn @ blk._in_w.T).split((2 * H * dk + H * dv, H, H, H * dv), -1) w = blk._conv_w if decode: # one causal depthwise-conv step from the cached last K-1 inputs win = torch.cat([st.conv[i], qkv.T.float()], 1) y = (win * w).sum(-1, keepdim=True) st.conv[i].copy_(win[:, 1:]) else: win = F.pad(qkv.T.float(), (self.CONV - 1, 0)) y = F.conv1d(win[None], w[:, None], groups=w.shape[0])[0] if st is not None: st.conv[i].copy_(win[:, -(self.CONV - 1):]) q, k, v = F.silu(y).T.to(xn.dtype).split((H * dk, H * dk, H * dv), -1) g = (blk._neg_A * F.softplus(a.float() + m.dt_bias.float()))[None] # log decay beta = b.float().sigmoid()[None] kern = fused_recurrent_gated_delta_rule if decode else chunk_gated_delta_rule o, S = kern(q=q.reshape(1, N, H, dk), k=k.reshape(1, N, H, dk), v=v.reshape(1, N, H, dv), g=g, beta=beta, scale=dk ** -0.5, initial_state=st.S[i] if decode else None, output_final_state=st is not None, use_qk_l2norm_in_kernel=True) if st is not None: st.S[i].copy_(S) o = _rms(o[0], m.o_norm.weight, self.EPS) * F.silu(gate.float().view(N, H, dv)) return o.reshape(N, -1).to(xn.dtype) @ m.o_proj.weight.T def _attention(self, blk, i, xn, pos, st, decode): m, N, H, KV, HD = blk.mixer, xn.shape[0], self.H, self.KV, self.HD q, k, v, gate = (xn @ blk._in_w.T).split((H * HD, KV * HD, KV * HD, H * HD), -1) q = self._rope(_rms(q.view(N, H, HD), m.q_norm.weight, self.EPS), pos).to(xn.dtype) k = self._rope(_rms(k.view(N, KV, HD), m.k_norm.weight, self.EPS), pos).to(xn.dtype) v = v.view(N, KV, HD) if decode: st.k[i].index_copy_(2, pos, k.transpose(0, 1)[None]) st.v[i].index_copy_(2, pos, v.transpose(0, 1)[None]) mask = (torch.arange(st.max_len, device=xn.device) <= pos)[None, None, None] o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], st.k[i], st.v[i], attn_mask=mask, enable_gqa=True) else: if st is not None: st.k[i][:, :, :N] = k.transpose(0, 1)[None] st.v[i][:, :, :N] = v.transpose(0, 1)[None] o = F.scaled_dot_product_attention(q.transpose(0, 1)[None], k.transpose(0, 1)[None], v.transpose(0, 1)[None], is_causal=True, enable_gqa=True) o = o[0].transpose(0, 1) * torch.sigmoid(gate.view(N, H, HD)) return o.reshape(N, -1) @ m.o_proj.weight.T def _run(self, ids, st, decode, keep): """ids [N] (prefill from position 0) or [1] (decode at st.pos). Returns logits [keep or N, V] (FP32).""" N, dt, eps = ids.shape[0], self.mm_dtype, self.EPS cids = self.ngram.canon[ids] if decode: pos = st.pos st.ctx.copy_(torch.cat([st.ctx[1:], cids])) mem = self._ngram_memory(st.ctx, torch.cat([pos - 2, pos - 1, pos]))[-1:] else: pos = torch.arange(N, device=ids.device) mem = self._ngram_memory(cids, pos) if st is not None: st.ctx.copy_(F.pad(cids, (3, 0))[-3:]) st.pos.fill_(N) x = self.embed.weight[ids].float() # residual stream in FP32 sources = [] for i, blk in enumerate(self.blocks): if i in self.inject: # n-gram memory, context-gated add k = self.inject[i] e = (mem.to(dt) @ self.ngram.proj[k].weight.T).float() gate = torch.sigmoid((_rms(x, self.ngram.key_norm[k].weight, eps) * _rms(e, self.ngram.val_norm[k].weight, eps)).sum(-1, keepdim=True) / math.sqrt(self.D)) x = x + gate * e attn_sources = list(sources) if i % self.per_block == 0: # Delta Block source, "repo": s_k = x_k - s_{k-1} sources.append(x if not sources else x - sources[-1]) h = x + _select(attn_sources, blk._daq, eps) if attn_sources else x xn = _rms(h, blk.norm1.weight, eps).to(dt) y = self._attention(blk, i, xn, pos, st, decode) if blk.is_attn else self._gdn(blk, i, xn, st, decode) x = x + blk._g1 * y.float() h = x + _select(sources, blk._dmq, eps) gt, up = (_rms(h, blk.norm2.weight, eps).to(dt) @ blk._mlp_w.T).chunk(2, -1) x = x + blk._g2 * ((F.silu(gt) * up) @ blk.mlp.w_down.weight.T).float() if keep: x = x[-keep:] logits = (_rms(x, self.norm.weight, eps).to(dt) @ self.lm_head.weight.T).float() if decode: st.pos += 1 return logits def _decode(self, st, token): st.tok.copy_(token) if not (self.config.use_cuda_graph and st.device.type == "cuda"): return self._run(st.tok, st, True, 0) if st.graph is None: snap = [t.clone() for t in st.recurrent()] s = torch.cuda.Stream() s.wait_stream(torch.cuda.current_stream()) with torch.cuda.stream(s): # warm up (Triton autotune) outside the graph for _ in range(2): self._run(st.tok, st, True, 0) torch.cuda.current_stream().wait_stream(s) graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): st.logits = self._run(st.tok, st, True, 0) for dst, src in zip(st.recurrent(), snap): dst.copy_(src) st.graph = graph st.graph.replay() return st.logits.clone() # ------------------------------------------------------------------ HF API def forward(self, input_ids=None, attention_mask=None, past_key_values=None, use_cache=None, labels=None, logits_to_keep=0, return_dict=True, **kwargs): if input_ids is None or input_ids.shape[0] != 1: raise ValueError("this model implementation supports input_ids of batch size 1") if attention_mask is not None and not bool(attention_mask.all()): raise ValueError("padding is not supported (batch size 1)") self._ready() use_cache = self.config.use_cache if use_cache is None else use_cache ids = input_ids[0] keep = logits_to_keep if isinstance(logits_to_keep, int) else 0 cache = past_key_values if isinstance(past_key_values, HybridCache) else None with torch.no_grad(): if cache is not None and cache.seen > 0: # continue a cached sequence token by token assert cache.state.owner is cache, "this cache was superseded by a newer prefill" if cache.seen + ids.numel() > cache.state.max_len: raise ValueError(f"sequence exceeds kv_cache_len={cache.state.max_len}") logits = torch.cat([self._decode(cache.state, t) for t in ids.split(1)]) cache.seen += ids.numel() else: cache = None if use_cache: if ids.numel() > self.config.kv_cache_len: raise ValueError(f"prompt exceeds kv_cache_len={self.config.kv_cache_len}") cache = HybridCache(self._get_state()) cache.seen = ids.numel() logits = self._run(ids, cache.state if cache else None, False, 0 if labels is not None else keep) loss = None if labels is not None: loss = F.cross_entropy(logits[:-1], labels[0, 1:], ignore_index=-100) return CausalLMOutputWithPast(loss=loss, logits=logits[None], past_key_values=cache) def prepare_inputs_for_generation(self, input_ids, past_key_values=None, attention_mask=None, use_cache=True, **kwargs): if isinstance(past_key_values, HybridCache) and past_key_values.seen > 0: input_ids = input_ids[:, past_key_values.seen:] return dict(input_ids=input_ids, past_key_values=past_key_values, use_cache=use_cache, logits_to_keep=1) HybridLMConfig.register_for_auto_class() HybridLMForCausalLM.register_for_auto_class("AutoModelForCausalLM")