Files
rustytorch/crates/training/rtx-flash-attention/Cargo.toml
T
Omar SobhandClaude Sonnet 4.6 082a50e3a0
CI / Build (ubuntu-latest) (push) Failing after 8s
Performance Benchmarks / Run Benchmarks (push) Successful in 8s
CI / Clippy Check (push) Failing after 8s
CI / Build CPU-Only (Explicit) (push) Failing after 8s
Documentation / Build API Documentation (push) Failing after 7s
CI / Format Check (push) Failing after 9s
GPU Tests / Check GPU Availability (push) Successful in 0s
Documentation / Build User Guide (push) Successful in 7s
GPU Tests / CUDA Tests (11.8) (push) Has been skipped
GPU Tests / CUDA Tests (12.1) (push) Has been skipped
CI / Build (macos-latest) (push) Failing after 14s
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
CI / Distributed Training Tests (push) Has been skipped
CI / CI Success (push) Failing after 0s
GPU Tests / Metal Tests (push) Has been skipped
feat(perf): GPU perf batch 1 — wire GPU execution paths for FP8, FA3, CUDA Graphs, SnapKV
FP8 GPU FFI (rtx-tensor)
- `fp8_cast.rs`: replaced `not_implemented` stubs with real cudarc 0.18.2 PTX launches;
  `cast_bf16_to_fp8_e4m3` and `cast_fp8_e4m3_to_bf16` now dispatch to NVCC-compiled
  `fp8_cast.ptx` via `LazyLock` module cache, matching the `inplace_ops` pattern
- `build.rs`: `create_dummy_ptx` now emits `fp8_cast.ptx` alongside `element_wise.ptx`
  so `include_str!` resolves cleanly when NVCC is absent

FlashAttention-3 typed kernel launch (rtx-flash-attention)
- `flash_v3_forward.rs`: `forward()` now takes typed `CudaSlice<bf16>` Q/K/V/O + `CudaSlice<f32>`
  LSE buffer; dispatches via `stream.launch_builder` with block_dim=(128,1,1),
  grid_dim=(ceil(seq_len/64), batch*heads, 1), shared_mem_bytes=0 (PTX metadata-resolved)
- `simple.rs`: added `has_flash_v3()` + `flash_attention_v3_forward_raw()` dispatch
- `Cargo.toml`: `half` added as optional cuda-gated dependency

CUDA Graphs stream threading (rtx-transformers)
- `training_loop.rs`: added `cuda_stream: Option<CudaStreamHandle>` field; `set_cuda_backend()`
  now creates a non-default capture stream; capture step calls real `begin_capture(stream)` +
  `end_capture(stream)`; added `set_cuda_stream()` override; replay unchanged (no stream needed)

SnapKV + prefix cache BatchScheduler wiring (rtx-inference)
- `scheduler.rs`: added `prefix_hit_pages: Option<Vec<PageId>>` + `evicted_positions: Vec<usize>`
  to `SchedulerRequest`; `BatchScheduler` gains `kv_cache` + `snapkv_eviction` fields;
  `submit_request` does non-blocking `try_lock` prefix lookup; added `notify_prefill_complete`
  (registers prefix + runs `select_evict_positions`), `set_kv_cache`, `set_snapkv_eviction`,
  `get_evicted_positions`, `get_prefix_hit_pages` — +5 new integration tests

Test results: 22 + 42 + 75 + 102 = 241 tests, 0 failures

Co-Authored-By: Claude Sonnet 4.6 <[email protected]>
2026-06-26 20:39:58 +00:00

74 lines
2.0 KiB
TOML

[package]
name = "rtx-flash-attention"
version = "1.0.0"
edition.workspace = true
rust-version = "1.92"
authors.workspace = true
license.workspace = true
repository.workspace = true
description = "Revolutionary Flash Attention implementation with Quantum/Neuromorphic/Edge variants"
[dependencies]
# Core RustyTorch dependencies
rtx-tensor = { workspace = true }
rtx-runtime = { workspace = true }
rtx-kernel = { workspace = true }
# GPU acceleration
cudarc = { workspace = true, optional = true }
half = { workspace = true, optional = true }
# Core utilities
anyhow = { workspace = true }
thiserror = { workspace = true }
tracing = { workspace = true }
serde = { workspace = true }
# Performance and memory management
parking_lot = { workspace = true }
rayon = { workspace = true }
dashmap = { workspace = true }
# Math and numerics
nalgebra = { workspace = true }
# Async support
async-trait = { workspace = true }
tokio = { workspace = true }
[target.'cfg(target_os = "macos")'.dependencies]
objc2 = { version = "0.6", optional = true, features = ["std"] }
objc2-metal = { version = "0.3", optional = true }
objc2-foundation = { version = "0.3", optional = true }
block2 = { version = "0.6", optional = true }
[dev-dependencies]
criterion = { workspace = true }
proptest = { workspace = true }
tokio-test = { workspace = true }
tempfile = { workspace = true }
[build-dependencies]
bindgen = "0.69"
cc = "1.0"
[[bench]]
name = "flash_attention_bench"
harness = false
[[example]]
name = "flash_attention_demo"
path = "examples/flash_attention_demo.rs"
required-features = ["disabled_tests"]
[features]
# Note: No default GPU features - select cuda or metal explicitly based on your platform
default = []
cuda = ["rtx-kernel/cuda", "rtx-tensor/cuda", "rtx-runtime/cuda", "dep:cudarc", "dep:half"]
metal = ["rtx-kernel/metal", "rtx-tensor/metal", "dep:objc2", "dep:objc2-metal", "dep:objc2-foundation", "dep:block2"]
variants = []
experimental = []
disabled_tests = []
[lints]
workspace = true