Commit Graph
4 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
Omar SobhandClaude Sonnet 4.6 33135194d4 docs: GPU perf batch 1 design spec (FA3, FP8, CUDA Graphs, SnapKV)
Research-driven design for 4 high-impact optimizations targeting
Blackwell SM_120 (RTX 5060 Ti):
- Item 1: CUDA Graphs wiring (90% exists, half-day task)
- Item 2: FP8 E4M3/E5M2 training (30-40% throughput, 50% memory)
- Item 3: FlashAttention-3 WGMMA+TMA+warp specialization (1.5-2x)
- Item 4: SnapKV + prefix caching (50-70% KV reduction)

Based on: arXiv:2407.08608 (FA3), arXiv:2511.05811 (MOSS FP8),
arXiv:2404.14469 (SnapKV), PyTorch 2025 state survey.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-26 19:40:26 +00:00
Omar SobhandClaude Sonnet 4.6 57e5252caa docs: add full repository review spec (2026-06-26)
Comprehensive layered architectural review covering all 109 crates,
306K LOC, 13,351 tests. Identifies 9 gaps (G0-G8) with the highest-
priority being 45 unimplemented! panics across rtx-backend-cuda/rocm/sycl
and the rtx-distributed workspace exclusion. Includes 17-item 4-phase
roadmap through 90 days.

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-26 12:48:03 +00:00