"""KernelConfig-mappable drop-in forwards for Hugging Face Transformers ViT encoders. ``NeuronDINOv3ViTModel`` -> ``DINOv3ViTModel`` (all DINOv3 ViT sizes: S/S+/B/L/H+; 7B needs TP) ``NeuronVJEPA2Encoder`` -> ``VJEPA2Encoder`` (V-JEPA2 video encoder, up to 2048 tokens / clip) Both run the whole transformer stack on the NKI kernels (``encoder.EncoderRunner``) and keep the model's own patch embedding, RoPE position computation and final LayerNorm. Weights are folded once, on the first device forward, from the stock module's parameters (checkpoint loading is untouched) and cached on the module. Anything the kernels do not support (off-device input, training mode, > 2048 tokens, unsupported head size) falls back to the stock forward. Usage (graft, works on any Transformers version):: import types model.forward = types.MethodType(kernels.NeuronDINOv3ViTModel.forward, model) model = torch.compile(model, backend="neuron") ``torch.compile(backend="neuron")`` is part of the intended configuration; the classes set ``can_torch_compile = True`` so the ``kernels`` library keeps them under ``Mode.TORCH_COMPILE``. """ from __future__ import annotations import os as _osm import torch from .encoder import LNC, BlockWeights, EncoderRunner, pack_block, rope_tables _PACK = "_vit_enc_neuron_packed" _TABS = "_vit_enc_neuron_tabs" _MAX_TOKENS = 2048 def _rup(n: int, m: int) -> int: return ((n + m - 1) // m) * m def _supported(H: int, heads: int) -> bool: D = H // heads return H % 128 == 0 and D in (64, 128) and H == heads * D def _stack(hs: torch.Tensor, N: int): B, n, H = hs.shape x = hs.float() if N > n: x = torch.cat([x, x.new_zeros(B, N - n, H)], 1) return x.reshape(B * N, H).contiguous() # ------------------------------------------------------------------------------------------------ # DINOv3 # ------------------------------------------------------------------------------------------------ def _dinov3_layers(model): enc = getattr(model, "model", None) # newer Transformers: model.model.layer return enc.layer if enc is not None and hasattr(enc, "layer") else model.layer def _dinov3_block(lyr) -> BlockWeights: a = lyr.attention bw = BlockWeights( ln1_w=lyr.norm1.weight, ln1_b=lyr.norm1.bias, ln1_eps=lyr.norm1.eps, wq=a.q_proj.weight, bq=a.q_proj.bias, wk=a.k_proj.weight, bk=a.k_proj.bias, wv=a.v_proj.weight, bv=a.v_proj.bias, wo=a.o_proj.weight, bo=a.o_proj.bias, ln2_w=lyr.norm2.weight, ln2_b=lyr.norm2.bias, ln2_eps=lyr.norm2.eps, ls1=lyr.layer_scale1.lambda1, ls2=lyr.layer_scale2.lambda1, scale=float(a.scaling)) m = lyr.mlp if hasattr(m, "gate_proj"): bw.gate_w, bw.gate_b = m.gate_proj.weight, m.gate_proj.bias bw.up_w, bw.up_b = m.up_proj.weight, m.up_proj.bias bw.down_w, bw.down_b = m.down_proj.weight, m.down_proj.bias else: bw.fc1_w, bw.fc1_b = m.up_proj.weight, m.up_proj.bias bw.fc2_w, bw.fc2_b = m.down_proj.weight, m.down_proj.bias return bw def _dinov3_ok(model) -> bool: c = model.config if not _supported(c.hidden_size, c.num_attention_heads): return False if getattr(c, "use_gated_mlp", False): return getattr(c, "hidden_act", "silu") == "silu" return getattr(c, "hidden_act", "gelu") == "gelu" @torch._dynamo.disable def _dinov3_freeze(model): if model.training: raise RuntimeError("NeuronDINOv3ViTModel: these kernels have no autograd rule; call model.eval() first.") dev = model.norm.weight.device heads = model.config.num_attention_heads packed = [pack_block(_dinov3_block(l), heads, "half", dev) for l in _dinov3_layers(model)] setattr(model, _PACK, EncoderRunner(packed)) setattr(model, _TABS, {}) @torch._dynamo.disable def _dinov3_tab(model, pixel_values, n_prefix, N): key = (tuple(pixel_values.shape[-2:]), N) tabs = getattr(model, _TABS) if key not in tabs: cos, sin = model.rope_embeddings(pixel_values[:1].float()) # [n_patches, D] each c = model.config tabs[key] = rope_tables(cos, sin, n_prefix, N, c.num_attention_heads, c.hidden_size // c.num_attention_heads, "half", model.norm.weight.device) return tabs[key] class NeuronDINOv3ViTModel(torch.nn.Module): """Drop-in forward for ``transformers.DINOv3ViTModel`` (bf16 NKI encoder, fp32 residual stream). Returns ``last_hidden_state`` (final-normed, ``[B, 1 + registers + patches, H]``) and ``pooler_output`` (CLS), like the stock model. Any batch size; images up to 2048 tokens (e.g. 720x720 at patch 16). """ can_torch_compile = True has_backward = False def forward(self, pixel_values, bool_masked_pos=None, head_mask=None, **kwargs): from transformers.modeling_outputs import BaseModelOutputWithPooling pv = pixel_values.to(self.embeddings.patch_embeddings.weight.dtype) hs = self.embeddings(pv, bool_masked_pos=bool_masked_pos) # [B, n, H] B, n, H = hs.shape N = _rup(n, 128 * LNC) if (hs.device.type != "neuron" or self.training or head_mask is not None or N > _MAX_TOKENS or not _dinov3_ok(self)): import transformers.models.dinov3_vit.modeling_dinov3_vit as mv if head_mask is not None: kwargs["head_mask"] = head_mask return mv.DINOv3ViTModel.forward(self, pixel_values, bool_masked_pos=bool_masked_pos, **kwargs) if not hasattr(self, _PACK): _dinov3_freeze(self) n_prefix = 1 + self.config.num_register_tokens tab = _dinov3_tab(self, pixel_values, n_prefix, N) x = getattr(self, _PACK)(_stack(hs, N), tab, n, B) x = x.reshape(B, N, H)[:, :n] out = self.norm(x.to(self.norm.weight.dtype)) return BaseModelOutputWithPooling(last_hidden_state=out, pooler_output=out[:, 0, :]) # ------------------------------------------------------------------------------------------------ # V-JEPA2 # ------------------------------------------------------------------------------------------------ def _vjepa2_block(lyr) -> BlockWeights: a = lyr.attention return BlockWeights( ln1_w=lyr.norm1.weight, ln1_b=lyr.norm1.bias, ln1_eps=lyr.norm1.eps, wq=a.query.weight, bq=a.query.bias, wk=a.key.weight, bk=a.key.bias, wv=a.value.weight, bv=a.value.bias, wo=a.proj.weight, bo=a.proj.bias, ln2_w=lyr.norm2.weight, ln2_b=lyr.norm2.bias, ln2_eps=lyr.norm2.eps, fc1_w=lyr.mlp.fc1.weight, fc1_b=lyr.mlp.fc1.bias, fc2_w=lyr.mlp.fc2.weight, fc2_b=lyr.mlp.fc2.bias, scale=float(a.scaling)) @torch._dynamo.disable def _vjepa2_tab(encoder, N, n): """Signed-sin 3D RoPE tables from the model's own position ids (interleaved pairing).""" a = encoder.layer[0].attention H, D = a.all_head_size, a.attention_head_size dev = a.query.weight.device pos = a.get_position_ids(torch.zeros(1, n, H, device=dev), masks=None) to1d = lambda t: (t if t.dim() == 1 else t.reshape(-1, n)[0]).float().cpu() cos = torch.ones(n, D) sin = torch.zeros(n, D) s = 0 for sub, p in zip((a.d_dim, a.h_dim, a.w_dim), (to1d(pos[0]), to1d(pos[1]), to1d(pos[2]))): half = sub // 2 omega = 1.0 / (10000.0 ** (torch.arange(half, dtype=torch.float32) / float(half))) f = p[:, None] * omega c, sn = f.cos(), f.sin() cos[:, s:s + half] = c; cos[:, s + half:s + sub] = c sin[:, s:s + half] = sn; sin[:, s + half:s + sub] = sn s += sub g = torch.ones(D); g[0::2] = -1.0 return rope_tables(cos, sin * g, 0, N, H // D, D, "interleaved", dev) @torch._dynamo.disable def _vjepa2_freeze(encoder): if encoder.training: raise RuntimeError("NeuronVJEPA2Encoder: these kernels have no autograd rule; call model.eval() first.") a = encoder.layer[0].attention dev = a.query.weight.device heads = a.all_head_size // a.attention_head_size setattr(encoder, _PACK, EncoderRunner([pack_block(_vjepa2_block(l), heads, "interleaved", dev) for l in encoder.layer])) setattr(encoder, _TABS, {}) class NeuronVJEPA2Encoder(torch.nn.Module): """Drop-in forward for ``transformers.VJEPA2Encoder``. Any batch size; clips up to 2048 tokens (e.g. 16 frames at 256 px, tubelet 2, patch 16). """ can_torch_compile = True has_backward = False def forward(self, pixel_values_videos=None, head_mask=None, output_attentions=False, output_hidden_states=False, **kwargs): from transformers.modeling_outputs import BaseModelOutput hs = self.embeddings(pixel_values_videos) # [B, n, H] B, n, H = hs.shape N = _rup(n, 128 * LNC) a = self.layer[0].attention if (hs.device.type != "neuron" or self.training or head_mask is not None or output_attentions or output_hidden_states or N > _MAX_TOKENS or not _supported(H, a.all_head_size // a.attention_head_size)): import transformers.models.vjepa2.modeling_vjepa2 as mv return mv.VJEPA2Encoder.forward(self, pixel_values_videos, head_mask=head_mask, output_attentions=output_attentions, output_hidden_states=output_hidden_states, **kwargs) if not hasattr(self, _PACK): _vjepa2_freeze(self) tabs = getattr(self, _TABS) if N not in tabs: tabs[N] = _vjepa2_tab(self, N, n) x = getattr(self, _PACK)(_stack(hs, N), tabs[N], n, B) out = self.layernorm(x.reshape(B, N, H)[:, :n]) return BaseModelOutput(last_hidden_state=out) # ------------------------------------------------------------------------------------------------ # XLM-RoBERTa / RoBERTa / BERT (post-LayerNorm encoders, no RoPE) -- v1.8 # ------------------------------------------------------------------------------------------------ _POST = "_vit_enc_neuron_postln" def _bert_blocks(model): """BlockWeights per layer for a post-LN BERT-family encoder. Block i's "ln1" is the LayerNorm that produced its input (embeddings.LayerNorm for block 0, else layer i-1's output.LayerNorm) and "ln2" its attention-output LayerNorm, so the kernels' affine-free LN + folded affine reproduce the post-LN math exactly.""" out, prev = [], model.embeddings.LayerNorm for lyr in model.encoder.layer: a = lyr.attention out.append(BlockWeights( ln1_w=prev.weight, ln1_b=prev.bias, ln1_eps=prev.eps, wq=a.self.query.weight, bq=a.self.query.bias, wk=a.self.key.weight, bk=a.self.key.bias, wv=a.self.value.weight, bv=a.self.value.bias, wo=a.output.dense.weight, bo=a.output.dense.bias, ln2_w=a.output.LayerNorm.weight, ln2_b=a.output.LayerNorm.bias, ln2_eps=a.output.LayerNorm.eps, fc1_w=lyr.intermediate.dense.weight, fc1_b=lyr.intermediate.dense.bias, fc2_w=lyr.output.dense.weight, fc2_b=lyr.output.dense.bias)) prev = lyr.output.LayerNorm return out def _bert_ok(model, S): c = model.config return (_supported(c.hidden_size, c.num_attention_heads) and getattr(c, "hidden_act", "gelu") == "gelu" and getattr(c, "position_embedding_type", "absolute") == "absolute" and S <= _MAX_TOKENS and S % (128 * LNC) == 0 and LNC == 1) @torch._dynamo.disable def _bert_freeze(model): if model.training: raise RuntimeError("NeuronXLMRobertaModel: these kernels have no autograd rule; call model.eval() first.") c = model.config dev = model.embeddings.word_embeddings.weight.device heads, H = c.num_attention_heads, c.hidden_size D = H // heads runner = EncoderRunner([pack_block(bw, heads, "half", dev) for bw in _bert_blocks(model)]) f32 = lambda t: t.detach().float().to(dev) lns = [model.embeddings.LayerNorm] + [l.output.LayerNorm for l in model.encoder.layer] st = {"runner": runner, "D": D, "heads": heads, "H": H, "tabs": {}, "ln_in": [(f32(m.weight), f32(m.bias), float(m.eps)) for m in lns], "ln_attn": [(f32(l.attention.output.LayerNorm.weight), f32(l.attention.output.LayerNorm.bias), float(l.attention.output.LayerNorm.eps)) for l in model.encoder.layer]} setattr(model, _POST, st) @torch._dynamo.disable def _bert_tab(st, S, dev): if S not in st["tabs"]: D = st["D"] # no RoPE: identity tables (the kernel skips them, norope=1) st["tabs"][S] = rope_tables(torch.ones(S, D), torch.zeros(S, D), 0, S, st["heads"], D, "half", dev) return st["tabs"][S] @torch._dynamo.disable def _bert_masked(attention_mask) -> bool: """True if any position is masked out. Runs outside the compiled graph (one small host check per call).""" return not bool(torch.all(attention_mask == 1)) class NeuronXLMRobertaModel(torch.nn.Module): """Drop-in forward for ``XLMRobertaModel`` / ``RobertaModel`` / ``BertModel`` (post-LN, GELU, absolute positions). Graft on an instance (``model.forward = types.MethodType(NeuronXLMRobertaModel.forward, model)``); do not use it as a class-level KernelConfig replacement, because the fallback calls ``type(model).forward``. Returns ``last_hidden_state``. Kernel path: LNC=1, sequence length a multiple of 128 and <= 2048. v1.9: a padding ``attention_mask`` (any 0/1 pattern) stays on the kernel path as an additive key bias (``VIT_ENC_MASK_KERNEL=0`` restores the v1.8 fallback). Anything else (LNC=2, other lengths, off-device, training) falls back to the stock forward. """ can_torch_compile = True has_backward = False def forward(self, input_ids=None, attention_mask=None, token_type_ids=None, position_ids=None, **kwargs): import torch.nn.functional as F from transformers.modeling_outputs import BaseModelOutputWithPoolingAndCrossAttentions as Out from .encoder import _ab1ln from .nki_kernels.mlp import nki_mlp_k4_xp from .nki_kernels.mlp_res_ln import nki_mlp_res_ln_post from .nki_kernels.res_layernorm import nki_res_layernorm_fold S = input_ids.shape[1] if input_ids is not None else -1 # The mask check runs outside the compiled graph (one host sync per call); pass attention_mask=None to skip it. masked = attention_mask is not None and _bert_masked(attention_mask) if masked and _osm.environ.get("VIT_ENC_MASK_KERNEL", "1") == "0": kmask = None else: kmask = attention_mask if masked else None if (input_ids is None or input_ids.device.type != "neuron" or self.training or (masked and kmask is None) or position_ids is not None or kwargs.get("inputs_embeds") is not None or not _bert_ok(self, S)): # stock forward of the grafted instance's own class (XLMRobertaModel / RobertaModel / BertModel). The graft # sets an INSTANCE attribute (types.MethodType), so the class attribute is still the stock forward. return type(self).forward(self, input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, position_ids=position_ids, **kwargs) if not hasattr(self, _POST): _bert_freeze(self) st = getattr(self, _POST) R, H, D = st["runner"], st["H"], st["D"] B = input_ids.shape[0] # the kernels need the PRE-LayerNorm embedding sum (the embedding LN is block 0's "ln1"); same as embeddings() e = self.embeddings pos = e.position_ids[:, :S] if getattr(e, "padding_idx", None) is None else None if pos is None: m = input_ids.ne(e.padding_idx).int() pos = (torch.cumsum(m, dim=1) * m).long() + e.padding_idx tt = token_type_ids if token_type_ids is not None else torch.zeros_like(input_ids) x = (e.word_embeddings(input_ids) + e.position_embeddings(pos) + e.token_type_embeddings(tt)).float() x = x.reshape(B * S, H).contiguous() tab = _bert_tab(st, S, x.device) g0, b0, eps0 = st["ln_in"][0] _, z1 = nki_res_layernorm_fold(x, x.new_zeros(1, x.shape[0], H, dtype=torch.bfloat16), eps0) h = F.layer_norm(x, (H,), g0, b0, eps0) nL = len(R.L) from .encoder import _PACK2 pk2 = int(_PACK2 and B % 2 == 0 and 2 * S <= 512 and kmask is None) # v1.9: padding mask -> additive key bias [B, S] (0 attended, -1e9 masked), same softmax as the stock # extended mask; every query row (padding rows included) is computed, as in the stock forward. kb = None if kmask is None else ((1.0 - kmask.to(torch.float32)) * -1e9).contiguous() for li, p in enumerate(R.L): ga, ba, _ = st["ln_attn"][li] if kb is None: _, z2, h1 = _ab1ln.attention_block_l1_ln_post( z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], h, p["scale"], D, S, p["eps2"], ga, ba, B, R._res(_ab1ln, S), pk2, 1) else: _, z2, h1 = _ab1ln.attention_block_l1_ln_post_kb( z1, p["wqk"], p["bqk"], tab, p["wvt"], p["bv"], p["wot"], p["bo_full"], h, p["scale"], D, S, p["eps2"], ga, ba, kb, B, R._res(_ab1ln, S), 1) go, bo, eo = st["ln_in"][li + 1] if li + 1 < nL: _, z1, h = nki_mlp_res_ln_post(z2, p["w1t"], p["b1"], p["w2t"], p["b2"], h1, eo, go, bo) else: h = F.layer_norm(h1 + nki_mlp_k4_xp(z2, p["w1t"], p["b1"], p["w2t"], p["b2"]).float(), (H,), go, bo, eo) out = h.reshape(B, S, H) return Out(last_hidden_state=out, pooler_output=None)