Commit Graph
2 Commits
Author SHA1 Message Date
Omar SobhandClaude Sonnet 4.6 e1b4061c23 feat(jepa): extended GPU training, data pipeline, integration, and cargo config
CI / Format Check (push) Failing after 12s
CI / Build (macos-latest) (push) Failing after 12s
CI / Build (ubuntu-latest) (push) Failing after 19s
CI / Distributed Training Tests (push) Has been skipped
CI / Clippy Check (push) Failing after 19s
CI / Test (macos-latest) (push) Has been skipped
CI / Test (ubuntu-latest) (push) Has been skipped
CI / Python Bindings (maturin) (macos-latest) (push) Has been skipped
CI / Python Bindings (maturin) (ubuntu-latest) (push) Has been skipped
CI / WASM Build + Size Check (push) Has been skipped
Documentation / Build User Guide (push) Successful in 8s
CI / Build CPU-Only (Explicit) (push) Failing after 16s
Documentation / Build API Documentation (push) Failing after 13s
CI / CI Success (push) Failing after 0s
Performance Benchmarks / Run Benchmarks (push) Successful in 43s
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]>
2026-06-29 21:35:34 +00:00
Omar SobhandClaude Sonnet 4.6 80d7c7fb6c feat(flash-attention): add varlen packed-sequence support
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]>
2026-06-27 01:09:39 +00:00