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