| """Block-causal 'cached' variant (Stage 4). |
| |
| Sequence = [pass-A context incl. [CALL]x4 [SEP]] [ARGS] slot block [SEP]. |
| Mask rule: context tokens see only context; slot-block tokens see everything. |
| Consequences: |
| * training needs ONE forward (decision/length heads read context states that can't see the slots); |
| * at inference the context K/V of every layer is computed once (that IS pass A) and each refinement |
| pass runs only the ~10-40 slot tokens through the network. |
| """ |
| from __future__ import annotations |
|
|
| import torch |
| import torch.nn.functional as F |
| from transformers.models.modernbert.modeling_modernbert import apply_rotary_pos_emb |
|
|
|
|
| def block_masks(attn: torch.Tensor, ctx_len: torch.Tensor, window: int) -> dict: |
| """attn [B,L] (1 = real token), ctx_len [B]. Returns HF ModernBERT mask dict (bool, True = attend).""" |
| B, L = attn.shape |
| idx = torch.arange(L, device=attn.device) |
| valid_k = attn.bool()[:, None, None, :] |
| q_slot = (idx[None, :] >= ctx_len[:, None])[:, None, :, None] |
| k_ctx = (idx[None, :] < ctx_len[:, None])[:, None, None, :] |
| eye = torch.eye(L, dtype=torch.bool, device=attn.device)[None, None] |
| full = (valid_k & (q_slot | k_ctx)) | eye |
| win = ((idx[:, None] - idx[None, :]).abs() <= window)[None, None] |
| return {"full_attention": full, "sliding_attention": full & win} |
|
|
|
|
| def _layer_qkv(layer, x, cos, sin): |
| attn = layer.attn |
| B, L, _ = x.shape |
| qkv = attn.Wqkv(layer.attn_norm(x)).view(B, L, 3, -1, attn.head_dim) |
| q, k, v = (t.transpose(1, 2) for t in qkv.unbind(dim=2)) |
| q, k = apply_rotary_pos_emb(q, k, cos, sin, unsqueeze_dim=1) |
| return q, k, v |
|
|
|
|
| def _finish(layer, x, out): |
| B, L = x.shape[:2] |
| x = x + layer.attn.Wo(out.transpose(1, 2).reshape(B, L, -1)) |
| return x + layer.mlp(layer.mlp_norm(x)) |
|
|
|
|
| @torch.no_grad() |
| def encode_context(enc, ids: torch.Tensor): |
| """enc = ModernBertModel. ids [1,Lc]. Returns (final hidden [1,Lc,d], per-layer (k, v) cache).""" |
| cfg = enc.config |
| x = enc.embeddings(input_ids=ids) |
| L = ids.shape[1] |
| pos = torch.arange(L, device=ids.device)[None] |
| pe = {lt: enc.rotary_emb(x, pos, lt) for lt in set(cfg.layer_types)} |
| idx = torch.arange(L, device=ids.device) |
| win = ((idx[:, None] - idx[None, :]).abs() <= cfg.sliding_window)[None, None] |
| cache = [] |
| for layer in enc.layers: |
| cos, sin = pe[layer.attention_type] |
| q, k, v = _layer_qkv(layer, x, cos, sin) |
| mask = win if layer.attention_type == "sliding_attention" else None |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5) |
| x = _finish(layer, x, out) |
| cache.append((k, v)) |
| return enc.final_norm(x), cache |
|
|
|
|
| @torch.no_grad() |
| def encode_slots(enc, slot_ids: torch.Tensor, start: int, cache): |
| """Run only the slot-block tokens (absolute positions start..start+S-1) over the cached context K/V.""" |
| cfg = enc.config |
| x = enc.embeddings(input_ids=slot_ids) |
| S = slot_ids.shape[1] |
| Lc = cache[0][0].shape[2] |
| qpos = torch.arange(start, start + S, device=slot_ids.device) |
| kpos = torch.cat([torch.arange(Lc, device=slot_ids.device), qpos]) |
| pe = {lt: enc.rotary_emb(x, qpos[None], lt) for lt in set(cfg.layer_types)} |
| win = ((qpos[:, None] - kpos[None, :]).abs() <= cfg.sliding_window)[None, None] |
| for layer, (ck, cv) in zip(enc.layers, cache): |
| cos, sin = pe[layer.attention_type] |
| q, k, v = _layer_qkv(layer, x, cos, sin) |
| k, v = torch.cat([ck, k], dim=2), torch.cat([cv, v], dim=2) |
| mask = win if layer.attention_type == "sliding_attention" else None |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5) |
| x = _finish(layer, x, out) |
| return enc.final_norm(x) |
|
|
|
|
| |
| |
| |
|
|
| def _ctx_core(enc, ids, valid): |
| """ids [1,P] padded; valid = tensor([real length]). -> hidden [1,P,d], cache.""" |
| cfg = enc.config |
| x = enc.embeddings(input_ids=ids) |
| P = ids.shape[1] |
| idx = torch.arange(P, device=ids.device) |
| pos = idx[None] |
| pe = {lt: enc.rotary_emb(x, pos, lt) for lt in set(cfg.layer_types)} |
| keyok = (idx < valid)[None, None, None, :] |
| eye = torch.eye(P, dtype=torch.bool, device=ids.device)[None, None] |
| full = keyok | eye |
| win = full & ((idx[:, None] - idx[None, :]).abs() <= cfg.sliding_window)[None, None] |
| cache = [] |
| for layer in enc.layers: |
| cos, sin = pe[layer.attention_type] |
| q, k, v = _layer_qkv(layer, x, cos, sin) |
| mask = win if layer.attention_type == "sliding_attention" else full |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5) |
| x = _finish(layer, x, out) |
| cache.append((k, v)) |
| return enc.final_norm(x), cache |
|
|
|
|
| def _slot_core(enc, slot_ids, start, s_valid, c_valid, cache): |
| """slot_ids [1,Sp] padded; start/s_valid/c_valid tensors; cache from _ctx_core (padded context).""" |
| cfg = enc.config |
| x = enc.embeddings(input_ids=slot_ids) |
| Sp = slot_ids.shape[1] |
| Lc = cache[0][0].shape[2] |
| ar = torch.arange(Sp, device=slot_ids.device) |
| qpos = start + ar |
| kc = torch.arange(Lc, device=slot_ids.device) |
| kpos = torch.cat([kc, qpos]) |
| keyok = torch.cat([kc < c_valid, ar < s_valid])[None, None, None, :] |
| pe = {lt: enc.rotary_emb(x, qpos[None], lt) for lt in set(cfg.layer_types)} |
| full = keyok | torch.cat([torch.zeros(Sp, Lc, dtype=torch.bool, device=slot_ids.device), |
| torch.eye(Sp, dtype=torch.bool, device=slot_ids.device)], dim=1)[None, None] |
| win = full & ((qpos[:, None] - kpos[None, :]).abs() <= cfg.sliding_window)[None, None] |
| for layer, (ck, cv) in zip(enc.layers, cache): |
| cos, sin = pe[layer.attention_type] |
| q, k, v = _layer_qkv(layer, x, cos, sin) |
| k, v = torch.cat([ck, k], dim=2), torch.cat([cv, v], dim=2) |
| mask = win if layer.attention_type == "sliding_attention" else full |
| out = F.scaled_dot_product_attention(q, k, v, attn_mask=mask, scale=layer.attn.head_dim ** -0.5) |
| x = _finish(layer, x, out) |
| return enc.final_norm(x) |
|
|
|
|
| def _bucket(n: int, b: int) -> int: |
| return -(-n // b) * b |
|
|
|
|
| def make_static(enc, ctx_bucket: int = 256, slot_bucket: int = 128, compile: bool = True): |
| """-> (encode_context, encode_slots) drop-ins that pad to buckets (and run compiled, CUDA-graphed cores). |
| Buckets are coarse on purpose: every (context bucket, slot bucket) pair is a separate compiled graph.""" |
| ctx_core, slot_core = _ctx_core, _slot_core |
| if compile: |
| import torch._dynamo as dynamo |
| for name in ("cache_size_limit", "recompile_limit"): |
| if hasattr(dynamo.config, name): |
| setattr(dynamo.config, name, 64) |
| ctx_core = torch.compile(_ctx_core, mode="reduce-overhead", dynamic=False, fullgraph=False) |
| slot_core = torch.compile(_slot_core, mode="reduce-overhead", dynamic=False, fullgraph=False) |
| pad_id = enc.config.pad_token_id or 0 |
| state = {} |
|
|
| @torch.no_grad() |
| def encode_context_s(enc_, ids): |
| L = ids.shape[1] |
| P = _bucket(L, ctx_bucket) |
| ids_p = F.pad(ids, (0, P - L), value=pad_id) |
| h, cache = ctx_core(enc, ids_p, torch.tensor(L, device=ids.device)) |
| h, cache = h.clone(), [(k.clone(), v.clone()) for k, v in cache] |
| state["c_valid"] = L |
| return h[:, :L], cache |
|
|
| @torch.no_grad() |
| def encode_slots_s(enc_, slot_ids, start, cache): |
| S = slot_ids.shape[1] |
| Sp = _bucket(S, slot_bucket) |
| ids_p = F.pad(slot_ids, (0, Sp - S), value=pad_id) |
| dev = slot_ids.device |
| out = slot_core(enc, ids_p, torch.tensor(start, device=dev), torch.tensor(S, device=dev), |
| torch.tensor(state["c_valid"], device=dev), cache) |
| return out[:, :S].clone() |
|
|
| return encode_context_s, encode_slots_s |
|
|
