Kernels:

---

aws-neuron

/

vit-encoder-neuron-kernels

Like 3

Follow

AWS Inferentia and Trainium 194

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