File size: 8,220 Bytes
6fd3740 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 | """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)
# ---------- static-shape variants for CUDA graphs / torch.compile (v3 S6) ----------
# Lengths are padded up to a bucket and the real lengths are passed as tensors, so one compiled graph serves every
# prompt in a bucket. Padded keys are masked out; outputs at padded positions are garbage and get sliced off.
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] # outlive the CUDA-graph buffers
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
|