Use Flash-Attention 3 (prebuilt via HF kernels) as the attention backend

Load kernels-community/flash-attn3 (auto-resolves the prebuilt kernel for
ZeroGPU's Blackwell arch) and expose it as `flash_attn_interface`, which is what
DVLT's attention imports. Hardcode set_attn_backend("fa3") — dropped the
ATTN_BACKEND env + _select_attn_backend selection logic. Falls back to SDPA only
when the kernel is unavailable (e.g. local CPU).

Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

app.py

@@ -4,6 +4,7 @@
 # dependencies = [
 #     "gradio>=5.49,<6",
 #     "spaces",
+#     "kernels",
 #     "trimesh>=4.4",
 #     "torch>=2.5.1",
 #     "torchvision>=0.20.1",
@@ -23,6 +24,7 @@ Hugging Face ZeroGPU Space (deps in requirements.txt; torch from the image).
 from __future__ import annotations
 
 import os
+import sys
 import tempfile
 import time
 from dataclasses import dataclass, field
@@ -35,10 +37,20 @@ import torch
 from accelerate import Accelerator
 from PIL import Image
 
+# Flash-Attention 3, prebuilt for ZeroGPU's Blackwell GPUs and exposed under the
+# module name DVLT imports (`from flash_attn_interface import flash_attn_func`).
+try:
+    from kernels import get_kernel
+
+    sys.modules["flash_attn_interface"] = get_kernel("kernels-community/flash-attn3")
+except Exception as exc:  # local / no compatible kernel -> DVLT falls back to SDPA
+    print(f"[dvlt] FA3 kernel unavailable ({exc}); using default attention.")
+
 from dvlt.common.constants import DataField, PredictionField
 from dvlt.common.geometry import depth_to_world_coords_points
 from dvlt.common.pose import to4x4
 from dvlt.model.dvlt.model import DVLT, _slice_expand_flatten
+from dvlt.model_components import set_attn_backend
 from dvlt.util.preprocess import preprocess_images
 from dvlt.viz.depth import overlay_depth_map
 from dvlt.viz.glb import pointcloud_to_glb
@@ -52,7 +64,6 @@ PATCH_SIZE = 14
 DEFAULT_STEPS = 12  # K — the model's default inference step count
 MAX_STEPS = 24
 MAX_FRAMES = 16  # cap views so global attention stays within the ZeroGPU budget
-ATTN_BACKEND = os.environ.get("DVLT_ATTN", "auto")  # auto | flash | fa3 (prefers fa3 if installed)
 
 VIDEO_FPS_DEFAULT = 2.0
 DECODE_EVERY_DEFAULT = 3
@@ -71,24 +82,13 @@ _ACCEL = Accelerator(mixed_precision="bf16" if torch.cuda.is_available() else "n
 _MODEL: DVLT | None = None
 
 
-def _select_attn_backend():
-    """Pick the fastest available attention backend (fa3 > flash > auto)."""
-    from dvlt.model_components import set_attn_backend
-
-    order = {"fa3": ["fa3", "flash", "auto"], "flash": ["flash", "auto"]}.get(ATTN_BACKEND, ["auto"])
-    for backend in order:
-        try:
-            set_attn_backend(backend)
-            print(f"[dvlt] attention backend: {backend}", flush=True)
-            return
-        except Exception:
-            continue
-
-
 def load_model() -> DVLT:
     global _MODEL
     if _MODEL is None:
-        _select_attn_backend()
+        try:
+            set_attn_backend("fa3")
+        except Exception as exc:  # noqa: BLE001
+            print(f"[dvlt] fa3 unavailable ({exc}); using default attention.")
         model = DVLT(img_size=IMG_SIZE, depth_head_type="conv")
         model.load_pretrained(CHECKPOINT, strict=True)
         model.setup_test(_ACCEL)

requirements.txt

@@ -8,6 +8,7 @@
 
 gradio>=5.49,<6
 spaces
+kernels
 trimesh>=4.4
 
 dvlt @ https://huggingface.co/spaces/blanchon/dvlt/resolve/main/packages/dvlt/wheels/dvlt-0.0.1-py3-none-any.whl