95 lines
3.5 KiB
Rust
95 lines
3.5 KiB
Rust
// AETHER temporal head over CSI feature windows (ADR-096).
|
|
//
|
|
// Wraps `ruvllm_sparse_attention::SubquadraticSparseAttention` so AETHER
|
|
// callers in `wifi-densepose-train` and `wifi-densepose-signal` can swap
|
|
// dense MHA for sparse-GQA without touching the contrastive recipe.
|
|
//
|
|
// Status: scaffolding for ADR-096 §4.3. Sparse backend is functional;
|
|
// the dense back-compat backend is a follow-up (Phase 2 of the roadmap
|
|
// in #513). Streaming `step()` lands once the per-track KvCache lifecycle
|
|
// (ADR-096 §8.5) is finalized.
|
|
|
|
pub mod config;
|
|
pub mod error;
|
|
pub mod sparse;
|
|
pub mod weights;
|
|
|
|
pub use config::{TemporalBackendKind, TemporalHeadConfig};
|
|
pub use error::TemporalError;
|
|
pub use sparse::SparseGqaHead;
|
|
pub use weights::{
|
|
WeightBlob, WeightBlobHeader, WeightDtype, WEIGHT_BLOB_HEADER_LEN, WEIGHT_BLOB_MAGIC,
|
|
WEIGHT_BLOB_VERSION,
|
|
};
|
|
|
|
// Re-export the upstream Tensor3 + KvCache so callers don't need a
|
|
// direct `ruvllm_sparse_attention` dep.
|
|
pub use ruvllm_sparse_attention::{KvCache, Tensor3};
|
|
|
|
/// Thin facade so callers can pick a backend by name.
|
|
///
|
|
/// Today only `SparseGqa` is implemented; `Dense` is reserved per
|
|
/// ADR-096 §4.4 and returns `TemporalError::DenseBackendNotImplemented`
|
|
/// until the back-compat path lands.
|
|
pub enum AetherTemporalHead {
|
|
SparseGqa(SparseGqaHead),
|
|
Dense, // placeholder; ADR-096 §4.4 selection rule
|
|
}
|
|
|
|
impl AetherTemporalHead {
|
|
pub fn new(cfg: &TemporalHeadConfig) -> Result<Self, TemporalError> {
|
|
match cfg.backend {
|
|
TemporalBackendKind::SparseGqa => {
|
|
Ok(AetherTemporalHead::SparseGqa(SparseGqaHead::new(cfg)?))
|
|
}
|
|
TemporalBackendKind::Dense => Err(TemporalError::DenseBackendNotImplemented),
|
|
}
|
|
}
|
|
|
|
/// Window-level prefill. Returns the per-token attention output as
|
|
/// a Tensor3 of shape (window, q_heads, head_dim). Pooling to a
|
|
/// single embedding is the caller's responsibility — different
|
|
/// AETHER consumers use different pool ops (mean for re-ID,
|
|
/// last-token for streaming).
|
|
pub fn forward(
|
|
&self,
|
|
q: &Tensor3,
|
|
k: &Tensor3,
|
|
v: &Tensor3,
|
|
) -> Result<Tensor3, TemporalError> {
|
|
match self {
|
|
AetherTemporalHead::SparseGqa(h) => h.forward(q, k, v),
|
|
AetherTemporalHead::Dense => Err(TemporalError::DenseBackendNotImplemented),
|
|
}
|
|
}
|
|
|
|
/// Streaming decode (ADR-096 §3.2). Caller owns the `cache`; the
|
|
/// natural lifetime is per-tracked-person (one cache per
|
|
/// `PoseTrack`, dropped when the track evicts).
|
|
///
|
|
/// Returns the attention output for the single new token. Caller
|
|
/// is responsible for downstream pooling / classifier head.
|
|
pub fn step(
|
|
&self,
|
|
q_new: &Tensor3,
|
|
k_new: &Tensor3,
|
|
v_new: &Tensor3,
|
|
cache: &mut KvCache,
|
|
) -> Result<Tensor3, TemporalError> {
|
|
match self {
|
|
AetherTemporalHead::SparseGqa(h) => h.step(q_new, k_new, v_new, cache),
|
|
AetherTemporalHead::Dense => Err(TemporalError::DenseBackendNotImplemented),
|
|
}
|
|
}
|
|
|
|
/// Allocate a `KvCache` sized correctly for this head. Convenience
|
|
/// wrapper so AETHER's `pose_tracker.rs` doesn't need to import
|
|
/// the upstream crate.
|
|
pub fn make_cache(&self, capacity: usize) -> Result<KvCache, TemporalError> {
|
|
match self {
|
|
AetherTemporalHead::SparseGqa(h) => Ok(h.make_cache(capacity)),
|
|
AetherTemporalHead::Dense => Err(TemporalError::DenseBackendNotImplemented),
|
|
}
|
|
}
|
|
}
|