| """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() |
|
|
|
|
| |
| |
| |
| 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) |
|
|
|
|
| |
| |
| |
| _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"] |
| 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 |
| |
| 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)): |
| |
| |
| 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] |
| |
| 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) |
| |
| |
| 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) |
|
|
