DSS pipeline: GPU kNN, LSH index, embedding cache

Three upgrades to the DynamicSimilaritySampler infrastructure:

1. **GPU kNN** (`dss_knn_device="cuda"`, default):
   Moves the per-seed similarity matmul to the GPU. At 25K train items
   this cuts per-epoch sampler generation from 17s to 1.6s — a 10.8x
   speedup. Negligible VRAM (100MB for the [N, 1024] embedding tensor).

2. **LSH index** (`src/datasets/lsh_index.py`, opt-in via `dss_use_lsh=True`):
   Random-projection cosine-LSH with H tables of B bits each. When enabled,
   the sampler narrows the candidate pool per seed via hash-bucket lookup
   before exact refinement. At 25K it's a wash (pool already fits in VRAM)
   but provides a scaling path for 100K+ where the N² similarity matrix
   would stop fitting. Default off.

3. **Embedding cache** (`src/datasets/embedding_cache.py`, `dss_cache_dir` config):
   Disk-backed cache for drone query embeddings, keyed by epoch. Skips
   re-embedding on --resume and lets ablations replay from a snapshot.
   Atomic writes via `.tmp` → `.replace`.

Measured on 25K train entries, 1024-dim random embeddings:
  CPU kNN:      17.44s
  GPU kNN:       1.62s  (10.8x)
  GPU + LSH:     1.42s  (LSH candidate pool 0.05% of N)

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

src/training/train_gtauav.py

@@ -39,6 +39,7 @@ from src.datasets.gtauav_dataset import (
     collate_sat_gallery,
 )
 from src.datasets.dynamic_similarity_sampler import DynamicSimilaritySampler
+from src.datasets.embedding_cache import EmbeddingCache
 from src.datasets.mutually_exclusive_sampler import MutuallyExclusiveSampler
 from src.losses.multi_infonce import InfoNCELoss
 from src.losses.weighted_infonce import WeightedInfoNCELoss
@@ -118,6 +119,11 @@ class TrainConfigGTAUAV:
     sampler_type: str = "mutex"  # "mutex" (no false negatives) or "dss" (DSS + mutex)
     dss_reembed_every: int = 1  # Re-embed train queries every N epochs for DSS.
     dss_warmup_epochs: int = 1  # Use mutex-only for the first N epochs (fresh model embeddings aren't useful)
+    dss_knn_device: str = "cuda"  # Device for similarity matmul in DSS sampler.
+    dss_use_lsh: bool = False  # Approximate kNN via LSH (opt-in; exact is fast at 25K).
+    dss_lsh_num_tables: int = 8
+    dss_lsh_num_bits: int = 14
+    dss_cache_dir: str | None = None  # Disk cache for embeddings; None = disabled.
     # Legacy alias kept for backward compatibility.
     use_mutex_sampler: bool = True
 
@@ -622,10 +628,15 @@ def train(cfg: TrainConfigGTAUAV) -> None:
     if effective_sampler_type == "dss":
         batch_sampler = DynamicSimilaritySampler(
             sat_cand_list, batch_size=cfg.batch_size, shuffle=True, seed=cfg.seed,
+            knn_device=cfg.dss_knn_device,
+            use_lsh=cfg.dss_use_lsh,
+            lsh_num_tables=cfg.dss_lsh_num_tables,
+            lsh_num_bits=cfg.dss_lsh_num_bits,
         )
         LOGGER.info(
-            "Sampler: DynamicSimilarity — embedding-ranked batches with mutex constraint "
-            "(warmup=%d epochs mutex-only, re-embed every %d epochs)",
+            "Sampler: DynamicSimilarity — kNN on %s%s, warmup=%d, re-embed every %d epochs",
+            cfg.dss_knn_device,
+            " + LSH" if cfg.dss_use_lsh else "",
             cfg.dss_warmup_epochs, cfg.dss_reembed_every,
         )
     elif effective_sampler_type == "mutex":
@@ -655,6 +666,11 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             pin_memory=True,
             drop_last=True,
         )
+
+    emb_cache: EmbeddingCache | None = None
+    if cfg.dss_cache_dir is not None:
+        emb_cache = EmbeddingCache(cfg.dss_cache_dir)
+        LOGGER.info("DSS embedding cache: %s", cfg.dss_cache_dir)
     test_loader = DataLoader(
         test_ds,
         batch_size=cfg.batch_size,
@@ -753,15 +769,23 @@ def train(cfg: TrainConfigGTAUAV) -> None:
             and epoch >= cfg.dss_warmup_epochs
             and (epoch - cfg.dss_warmup_epochs) % cfg.dss_reembed_every == 0
         ):
-            LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch)
-            t_embed = time.time()
-            query_embs = _embed_drone_queries(
-                model, train_ds, cfg.device,
-                batch_size=cfg.batch_size * cfg.grad_accum_steps,
-                num_workers=cfg.num_workers,
-            )
+            query_embs: torch.Tensor | None = None
+            if emb_cache is not None:
+                query_embs = emb_cache.load(epoch)
+            if query_embs is None:
+                LOGGER.info("DSS: re-embedding %d train queries (epoch=%d)", len(train_ds), epoch)
+                t_embed = time.time()
+                query_embs = _embed_drone_queries(
+                    model, train_ds, cfg.device,
+                    batch_size=cfg.batch_size * cfg.grad_accum_steps,
+                    num_workers=cfg.num_workers,
+                )
+                LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed)
+                if emb_cache is not None:
+                    emb_cache.save(epoch, query_embs)
+            t_sampler = time.time()
             batch_sampler.update_embeddings(query_embs)
-            LOGGER.info("DSS: re-embed took %.1fs", time.time() - t_embed)
+            LOGGER.info("DSS: sampler update_embeddings took %.2fs", time.time() - t_sampler)
 
         epoch_start = time.time()
         agg: dict[str, float] = {}