| """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() |
|
|
|
|
| |
| |
| |
| def _dinov3_layers(model): |
| enc = getattr(model, "model", None) |
| 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()) |
| 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 = 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, :]) |
|
|
|
|
| |
| |
| |
| 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 = 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) |
|
|
