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