Kernels:

---

aws-neuron

/

vjepa2-neuron-kernels

Like 2

Follow

AWS Inferentia and Trainium 190

"""``NeuronVJEPA2Encoder`` -- the KernelConfig-mappable layer for the V-JEPA2 encoder.

USAGE (HF Transformers ``KernelConfig``)::

    from transformers import AutoModel, KernelConfig
    kernel_config = KernelConfig({
        "VJEPA2Encoder": "OWNER/vjepa2-neuron-kernels:NeuronVJEPA2Encoder",  # <- set to the published repo id
    })
    model = AutoModel.from_pretrained("facebook/vjepa2-vitl-fpc64-256",
                                      dtype=torch.float32,
                                      kernel_config=kernel_config,
                                      device_map="neuron")
    model.encoder.forward = model.encoder.forward  # (kernelize does the graft)

WHY THE ENCODER CLASS AND NOT A PER-LAYER MAPPING. The K1/K3 kernel fuses the *previous*
sublayer's residual add with the *next* sublayer's LayerNorm, so the fusion spans transformer-
layer boundaries: the state carried between layers is the pair (fp32 residual stream, bf16
branch partials), not a single summed hidden_states. Replacing the whole encoder forward keeps
that chain -- and the instruction stream it was verified bitwise against -- intact. A per-layer
mapping would force a residual re-add at every boundary and change the stream.

MECHANICS. The kernels library grafts this class's ``forward`` onto the existing
``VJEPA2Encoder`` instance, so ``self`` inside ``forward`` is the *stock* encoder with its
stock parameters (checkpoint loading is untouched). On the first device forward, the per-layer
packed weights (LayerNorm affines folded into QKV/V and MLP-fc1, out-proj bias halved for the
LNC=2 partials, rope tables) are computed once from those stock parameters and cached on the
layer objects. Off-device or non-fp32 inputs fall back to the stock implementation.

``can_torch_compile = True`` is REQUIRED here: the published performance (23.94 ms/forward,
44% MFU) is of the ``torch.compile(backend="neuron")`` graph, and without this flag the kernels
library silently swaps the layer back to stock under Mode.TORCH_COMPILE (observed on other
Neuron kernel repos as a 25-127x silent regression).
"""

from __future__ import annotations

import torch

from .nki_kernels import (
    attention_block,
    build_rope_tables,
    mlp,
    pack_out_weights,
    pack_tables,
    pack_v_weights,
    pack_weights,
    res_layernorm_fold,
)

_PACKED_ATTR = "_vjepa2_nki_packed"
_TAB_ATTR = "_vjepa2_nki_tab"


def _pack_layer(lyr, tab):
    """Fold one transformer block's LayerNorm affines into the kernel weight packings.

    Exact fp32 fold on CPU (the algebra of the verified reference implementation):
      K1 gamma1/beta1 -> QKV and V packs;  out-proj bias HALVED (two cores' partials sum);
      K3 gamma2/beta2 -> MLP fc1;  fc2 unfolded.
    """
    attn = lyr.attention
    dev = attn.query.weight.device
    H = attn.all_head_size
    D = attn.attention_head_size
    heads = H // D

    g1 = lyr.norm1.weight.detach().float().cpu()
    t1 = lyr.norm1.bias.detach().float().cpu()
    g2 = lyr.norm2.weight.detach().float().cpu()
    t2 = lyr.norm2.bias.detach().float().cpu()
    wq = attn.query.weight.detach().float().cpu(); bq = attn.query.bias.detach().float().cpu()
    wk = attn.key.weight.detach().float().cpu();   bk = attn.key.bias.detach().float().cpu()
    wv = attn.value.weight.detach().float().cpu(); bv = attn.value.bias.detach().float().cpu()

    wqk_f, bqk_f = pack_weights(wq * g1[None, :], wk * g1[None, :],
                                bq + wq @ t1, bk + wk @ t1, heads, D)
    wvt_f, bv_f = pack_v_weights(wv * g1[None, :], bv + wv @ t1)
    wot, bo = pack_out_weights(attn.proj.weight, attn.proj.bias)

    w1 = lyr.mlp.fc1.weight.detach().float().cpu(); b1 = lyr.mlp.fc1.bias.detach().float().cpu()
    packed = {
        "D": D, "heads": heads,
        "k1_eps": float(lyr.norm1.eps), "k3_eps": float(lyr.norm2.eps),
        "scale": float(attn.scaling),
        "wqk_f": wqk_f.to(dev), "bqk_f": bqk_f.to(dev), "tab": tab,
        "wvt_f": wvt_f.to(dev), "bv_f": bv_f.to(dev), "wot": wot.to(dev), "bo": bo.to(dev),
        "w1t_f": (w1 * g2[None, :]).t().contiguous().to(torch.bfloat16).to(dev),
        "b1_f": (b1 + w1 @ t2).contiguous().to(dev),
        "w2t": lyr.mlp.fc2.weight.detach().t().contiguous().to(torch.bfloat16).to(dev),
        "b2": lyr.mlp.fc2.bias.detach().float().contiguous().to(dev),
    }
    return packed


def _rope_tables(encoder, tokens):
    attn0 = encoder.layer[0].attention
    H = attn0.all_head_size
    D = attn0.attention_head_size
    dev = attn0.query.weight.device
    pos = attn0.get_position_ids(torch.zeros(1, tokens, H, device=dev), masks=None)

    def to1d(t):
        return t if t.dim() == 1 else t.reshape(-1, tokens)[0]

    cos, sinsigned = build_rope_tables(
        (attn0.d_dim, attn0.h_dim, attn0.w_dim),
        (to1d(pos[0]), to1d(pos[1]), to1d(pos[2])),
        D, torch.float32, dev)
    return pack_tables(cos, sinsigned, H // D, D)


@torch._dynamo.disable
def _freeze(encoder, tokens):
    """Pack every layer once; cached on the layer objects. Refuses in training mode.

    ``@torch._dynamo.disable`` is load-bearing: without it, a user who wraps the model in
    ``torch.compile`` before the first forward gets the weight-fold arithmetic traced/executed
    by the compiler instead of eagerly, and the packed weights round DIFFERENTLY (measured:
    output norm 4487.45 vs the verified 4486.19). The fold must run eagerly, once.
    """
    if encoder.training:
        raise RuntimeError(
            "NeuronVJEPA2Encoder: encoder.training is True; these kernels have no autograd "
            "rule. Call model.eval() first.")
    tab = _rope_tables(encoder, tokens)
    setattr(encoder, _TAB_ATTR, tab)
    for lyr in encoder.layer:
        setattr(lyr, _PACKED_ATTR, _pack_layer(lyr, tab))


class NeuronVJEPA2Encoder(torch.nn.Module):
    """Drop-in forward for ``VJEPA2Encoder``: the 4-kernel NKI layer loop.

    Supported: trn2, 16 frames / 2048 tokens / batch 1, fp32 model. Anything else falls back
    to the stock implementation (off-device, wrong geometry) or raises (training mode during
    the first-freeze).
    """

    can_torch_compile = True
    has_backward = False

    def forward(self, pixel_values_videos=None, **kwargs):
        hs = self.embeddings(pixel_values_videos)
        N = hs.shape[1]
        # Fallback: off-device, or an unsupported geometry (only N=2048 is verified; the vendor
        # attention kernel cannot hold 8192 keys in SBUF -- see the card's Known limitations).
        if hs.device.type != "neuron" or N != 2048:
            import transformers.models.vjepa2.modeling_vjepa2 as mv
            return mv.VJEPA2Encoder.forward(self, pixel_values_videos, **kwargs)

        if not hasattr(self.layer[0], _PACKED_ATTR):
            _freeze(self, N)

        x = hs[0]                                        # [N, H] fp32 residual stream
        H = x.shape[1]
        branch = torch.zeros(1, N, H, dtype=torch.bfloat16, device=x.device)

        for lyr in self.layer:
            p = getattr(lyr, _PACKED_ATTR)
            x, z1 = res_layernorm_fold(x, branch, p["k1_eps"])
            attn_partials = attention_block(
                z1, p["wqk_f"], p["bqk_f"], p["tab"], p["wvt_f"], p["bv_f"],
                p["wot"], p["bo"], p["scale"], p["D"])
            x, z2 = res_layernorm_fold(x, attn_partials, p["k3_eps"])
            branch = mlp(z2, p["w1t_f"], p["b1_f"], p["w2t"], p["b2"]).unsqueeze(0)

        x = x + branch[0].to(torch.float32)
        out = self.layernorm(x.unsqueeze(0))
        import transformers.models.vjepa2.modeling_vjepa2 as mv
        return mv.BaseModelOutput(last_hidden_state=out)