Extends jepa_train with distributed launcher, jepa_data with advanced
sampling and preprocessing, jepa_gpu with full CUDA kernel wiring,
jepa_distributed/runner/metrics/vit with additional training stages.
Adds jepa_integration module and project-local cargo config.
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
Implements variable-length (varlen) FlashAttention that processes
mixed-length batches without padding waste:
- New CUDA kernel flash_attention_varlen_forward with BLOCK_Q=64 /
BLOCK_K=64 tiling; grid=(ceil(max_seqlen_q/64), num_heads, 1).
Each block uses a linear scan over cu_seqlens_q to identify its
owning sequence and exits early when past sequence end.
- New Rust module flash_varlen_forward: always-compiled CPU simulation
(varlen_attention_cpu) for testing + #[cfg(cuda)] FlashVarlenKernel.
- SdpaBackend::VarLen variant added to backend_selector.
- 8 new CPU-only tests; total test count: 50.
Co-Authored-By: Claude Sonnet 4.6 <[email protected]>