"""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 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)