Skip to content

Native encoder ~500x slower than HF port: fp16 SDPA dtype crash + un-fused attention (frozen-feature extraction) #175

Description

@jotron-ops

Two issues in the native V-JEPA 2 / 2.1 encoder path that make frozen-feature extraction ~500× slower than the HuggingFace port.

Setup: RTX 3070 Ti 8 GB, torch 2.6.0+cu124, extracting frozen features from vjepa2_1_vit_large_384 and vjepa2_vit_large via torch.hub.load("facebookresearch/vjepa2", ...).

1. fp16 forward crashes (mixed dtype in SDPA). Casting the encoder + input to .half() throws inside attention:

RuntimeError: Expected query, key, and value to have the same dtype,
but got query.dtype: float key.dtype: float and value.dtype: c10::Half instead.

Some tensor (a norm/pos-embed buffer or a projection) stays fp32 while value is fp16, so scaled_dot_product_attention rejects the mix. .half() on the whole module isn't enough.

2. Attention is not memory-efficient / fused at the token counts these models actually use. At the native vjepa2_1_vit_large_384 setting (384px, 64 frames → ~18k tokens) a single-clip forward exceeds 250 s on this GPU; even at 256px/16 frames (~2k tokens) it is far slower than expected. By contrast the HF port facebook/vjepa2-vitl-fpc64-256 with attn_implementation="sdpa" extracts the same clips at ~0.8 s/clip. The gap is consistent with the native path not routing attention through a fused/flash kernel.

Workaround we used: run the native models in fp32 (avoids #1) at 256px / 16 frames (makes #2 tractable, ~0.5 s/clip) — but that forces a resolution/frame-budget compromise on 384-native checkpoints.

Suggestions: (a) make the encoder fp16-safe (cast pos-embed/norm buffers, or document that fp16 is unsupported); (b) route attention through F.scaled_dot_product_attention so the native path gets flash/mem-efficient kernels like the HF port. Happy to share a minimal repro script.

Repo @ main, 2026-07-23. (Related: #174, the committed localhost:8300 base-URL that also blocks hub loads.)

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions