mirror of
https://github.com/ruvnet/RuView
synced 2026-08-01 19:01:42 +00:00
feat(temporal): scaffold wifi-densepose-temporal crate (ADR-096 Phase 1-3, #513)
Implements Phases 1-3 of the ADR-096 roadmap: Phase 1: workspace integration - Add `ruvllm_sparse_attention` as a path-vendored workspace dep against `vendor/ruvector/crates/ruvllm_sparse_attention`, default-features=false, features=["fp16"]. Mirrors the no_std posture ADR-095 will need on the firmware side so both consumers share a single feature set. - Register `wifi-densepose-temporal` as workspace member. Phase 2: AETHER temporal head - `AetherTemporalHead` facade dispatches to a `SparseGqa` backend wrapping `SubquadraticSparseAttention`. Selection rule from ADR-096 §4.4 enforced at forward(): MHA branch when q_heads == kv_heads, GQA branch otherwise. - `Dense` backend reserved (returns typed `DenseBackendNotImplemented`) so config-time validation fails loudly instead of at forward(). - `TemporalHeadConfig::default_aether()` matches the AETHER training default per ADR-096 §3.1 (window=32, block=16, q=4, kv=1 → MQA). - Token 0 always wired as a global anchor — preserves AETHER's contrastive "session-start reference" role per ADR-024. Phase 3: smoke tests (5/5 passing) - forward at AETHER default config, both MHA and GQA dispatch paths, rejected dense backend, rejected non-divisible GQA ratio, and the long-window roadmap target (N=1000, the 10s @ 100Hz case from ADR-096 §3.1 — proves the kernel runs at lengths where dense MHA costs 10⁶ edge ops vs sparse 10⁴). Streaming `step()` deferred — KvCache lifecycle ties to PoseTrack per ADR-096 §8.5 and lands when the firmware-side ABI does (Phase 4+). Co-Authored-By: claude-flow <ruv@ruv.net>
This commit is contained in:
@@ -0,0 +1,19 @@
|
||||
[package]
|
||||
name = "wifi-densepose-temporal"
|
||||
version = "0.1.0"
|
||||
edition = "2021"
|
||||
license = "MIT"
|
||||
description = "AETHER temporal head for WiFi-DensePose — sparse-GQA attention over CSI feature windows (ADR-096)"
|
||||
repository = "https://github.com/ruvnet/RuView"
|
||||
|
||||
[dependencies]
|
||||
ruvllm_sparse_attention = { workspace = true }
|
||||
thiserror = "1"
|
||||
|
||||
[dev-dependencies]
|
||||
approx = "0.5"
|
||||
|
||||
[features]
|
||||
default = []
|
||||
# Enable FP16 KV cache path (mirrors the firmware-side ADR-095 build).
|
||||
fp16 = []
|
||||
@@ -0,0 +1,70 @@
|
||||
use crate::TemporalError;
|
||||
|
||||
/// Backend choice per ADR-096 §4.4.
|
||||
///
|
||||
/// * `Dense` — back-compat path against `ruvector-attention`. Reserved;
|
||||
/// not yet implemented in this crate (returns a typed error so callers
|
||||
/// can fail loudly during config validation rather than at forward()).
|
||||
/// * `SparseGqa` — `ruvllm_sparse_attention` `forward_gqa` for prefill,
|
||||
/// `decode_step` against `KvCache` for streaming inference.
|
||||
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
|
||||
pub enum TemporalBackendKind {
|
||||
Dense,
|
||||
SparseGqa,
|
||||
}
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct TemporalHeadConfig {
|
||||
pub backend: TemporalBackendKind,
|
||||
|
||||
/// Number of query heads. For pure MHA, equals `kv_heads`.
|
||||
pub q_heads: usize,
|
||||
/// Number of key/value heads. Must divide `q_heads`. GQA group size
|
||||
/// is `q_heads / kv_heads`.
|
||||
pub kv_heads: usize,
|
||||
/// Per-head feature dimension.
|
||||
pub head_dim: usize,
|
||||
|
||||
/// Local attention window radius (sparse pattern primitive #1, ADR-096 §3).
|
||||
pub window: usize,
|
||||
/// Landmark block size (sparse pattern primitive #3).
|
||||
pub block_size: usize,
|
||||
/// Whether the attention is causal. AETHER temporal aggregation is
|
||||
/// causal (cannot peek at future CSI frames during streaming re-ID).
|
||||
pub causal: bool,
|
||||
}
|
||||
|
||||
impl TemporalHeadConfig {
|
||||
/// Default config sized for the AETHER training default
|
||||
/// (`window_frames = 100`) but with the sparse machinery wired up
|
||||
/// so the long-window roadmap (10 s / 1000 frames) only requires
|
||||
/// changing `window` at the call site, not re-architecting.
|
||||
pub fn default_aether() -> Self {
|
||||
Self {
|
||||
backend: TemporalBackendKind::SparseGqa,
|
||||
q_heads: 4,
|
||||
kv_heads: 1, // MQA — collapses to one shared K/V across query heads
|
||||
head_dim: 32,
|
||||
window: 32,
|
||||
block_size: 16,
|
||||
causal: true,
|
||||
}
|
||||
}
|
||||
|
||||
pub fn validate(&self) -> Result<(), TemporalError> {
|
||||
if self.q_heads == 0 || self.kv_heads == 0 || self.head_dim == 0 {
|
||||
return Err(TemporalError::InvalidConfig(
|
||||
"q_heads, kv_heads, head_dim must all be > 0",
|
||||
));
|
||||
}
|
||||
if self.q_heads % self.kv_heads != 0 {
|
||||
return Err(TemporalError::InvalidConfig(
|
||||
"q_heads must be divisible by kv_heads (GQA constraint)",
|
||||
));
|
||||
}
|
||||
if self.block_size == 0 {
|
||||
return Err(TemporalError::InvalidConfig("block_size must be > 0"));
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum TemporalError {
|
||||
#[error("temporal head config invalid: {0}")]
|
||||
InvalidConfig(&'static str),
|
||||
|
||||
#[error("dense MHA backend not implemented yet (ADR-096 §4.4 follow-up)")]
|
||||
DenseBackendNotImplemented,
|
||||
|
||||
#[error("sparse attention kernel error: {0}")]
|
||||
Kernel(String),
|
||||
}
|
||||
|
||||
impl From<ruvllm_sparse_attention::AttentionError> for TemporalError {
|
||||
fn from(e: ruvllm_sparse_attention::AttentionError) -> Self {
|
||||
TemporalError::Kernel(format!("{e}"))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,60 @@
|
||||
// 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 use config::{TemporalBackendKind, TemporalHeadConfig};
|
||||
pub use error::TemporalError;
|
||||
pub use sparse::SparseGqaHead;
|
||||
|
||||
// Re-export the upstream Tensor3 so callers don't need a direct
|
||||
// `ruvllm_sparse_attention` dep.
|
||||
pub use ruvllm_sparse_attention::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),
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,68 @@
|
||||
use ruvllm_sparse_attention::{
|
||||
AttentionBackend, SparseAttentionConfig, SubquadraticSparseAttention, Tensor3,
|
||||
};
|
||||
|
||||
use crate::{TemporalError, TemporalHeadConfig};
|
||||
|
||||
/// AETHER temporal head implemented with `ruvllm_sparse_attention`.
|
||||
///
|
||||
/// The selection rule from ADR-096 §4.4 is enforced at `forward()`
|
||||
/// time: when `q_heads == kv_heads` we use `forward()` (plain MHA
|
||||
/// over the sparse pattern); when they differ we use `forward_gqa()`.
|
||||
/// The streaming `step()` path is staged behind a follow-up — KvCache
|
||||
/// lifecycle ties to `PoseTrack` per ADR-096 §8.5 and lives on the
|
||||
/// caller, not here.
|
||||
pub struct SparseGqaHead {
|
||||
cfg: TemporalHeadConfig,
|
||||
attn: SubquadraticSparseAttention,
|
||||
}
|
||||
|
||||
impl SparseGqaHead {
|
||||
pub fn new(cfg: &TemporalHeadConfig) -> Result<Self, TemporalError> {
|
||||
cfg.validate()?;
|
||||
|
||||
let attn_cfg = SparseAttentionConfig {
|
||||
window: cfg.window,
|
||||
block_size: cfg.block_size,
|
||||
global_tokens: alloc_first_token(),
|
||||
causal: cfg.causal,
|
||||
use_log_stride: true,
|
||||
use_landmarks: true,
|
||||
sort_candidates: false,
|
||||
};
|
||||
|
||||
let attn = SubquadraticSparseAttention::new(attn_cfg)?;
|
||||
Ok(Self {
|
||||
cfg: cfg.clone(),
|
||||
attn,
|
||||
})
|
||||
}
|
||||
|
||||
pub fn cfg(&self) -> &TemporalHeadConfig {
|
||||
&self.cfg
|
||||
}
|
||||
|
||||
pub fn forward(
|
||||
&self,
|
||||
q: &Tensor3,
|
||||
k: &Tensor3,
|
||||
v: &Tensor3,
|
||||
) -> Result<Tensor3, TemporalError> {
|
||||
// ADR-096 §4.4: dispatch by GQA shape.
|
||||
if self.cfg.q_heads == self.cfg.kv_heads {
|
||||
// Pure MHA — sparse `forward` is the right path.
|
||||
Ok(self.attn.forward(q, k, v)?)
|
||||
} else {
|
||||
// GQA / MQA — kv_heads < q_heads, group share factor = q/kv.
|
||||
Ok(self.attn.forward_gqa(q, k, v)?)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Always treat token 0 as a global anchor — AETHER's contrastive
|
||||
/// recipe (ADR-024) gives the first token a special role as the
|
||||
/// "session start" reference embedding, and global tokens in the
|
||||
/// sparse pattern preserve full visibility for that one position.
|
||||
fn alloc_first_token() -> Vec<usize> {
|
||||
vec![0]
|
||||
}
|
||||
@@ -0,0 +1,113 @@
|
||||
//! Smoke tests for the AETHER sparse-GQA temporal head (ADR-096 §5 gate is
|
||||
//! a separate accuracy benchmark; this file just proves the wiring works).
|
||||
|
||||
use wifi_densepose_temporal::{
|
||||
AetherTemporalHead, TemporalBackendKind, TemporalHeadConfig, TemporalError, Tensor3,
|
||||
};
|
||||
|
||||
fn make_qkv(seq: usize, q_heads: usize, kv_heads: usize, dim: usize) -> (Tensor3, Tensor3, Tensor3) {
|
||||
// Deterministic synthetic CSI-like activations so the test is
|
||||
// reproducible across machines without bringing in `rand`.
|
||||
let mut q = Tensor3::zeros(seq, q_heads, dim);
|
||||
for s in 0..seq {
|
||||
for h in 0..q_heads {
|
||||
for d in 0..dim {
|
||||
let v = ((s * 31 + h * 7 + d) as f32).sin() * 0.1;
|
||||
q.set(s, h, d, v);
|
||||
}
|
||||
}
|
||||
}
|
||||
let mut k = Tensor3::zeros(seq, kv_heads, dim);
|
||||
let mut v = Tensor3::zeros(seq, kv_heads, dim);
|
||||
for s in 0..seq {
|
||||
for h in 0..kv_heads {
|
||||
for d in 0..dim {
|
||||
let kv = (((s * 17 + h * 3 + d) as f32).cos()) * 0.1;
|
||||
k.set(s, h, d, kv);
|
||||
v.set(s, h, d, kv * 0.5);
|
||||
}
|
||||
}
|
||||
}
|
||||
(q, k, v)
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sparse_gqa_forward_runs_at_aether_default() {
|
||||
let cfg = TemporalHeadConfig::default_aether();
|
||||
let head = AetherTemporalHead::new(&cfg).expect("construct");
|
||||
|
||||
let (q, k, vt) = make_qkv(64, cfg.q_heads, cfg.kv_heads, cfg.head_dim);
|
||||
let out = head.forward(&q, &k, &vt).expect("forward");
|
||||
let (oseq, oh, od) = out.shape();
|
||||
assert_eq!(oseq, 64);
|
||||
assert_eq!(oh, cfg.q_heads);
|
||||
assert_eq!(od, cfg.head_dim);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn sparse_mha_path_runs_when_qkv_heads_match() {
|
||||
// q_heads == kv_heads forces the `forward` (non-GQA) branch.
|
||||
let cfg = TemporalHeadConfig {
|
||||
backend: TemporalBackendKind::SparseGqa,
|
||||
q_heads: 2,
|
||||
kv_heads: 2,
|
||||
head_dim: 16,
|
||||
window: 8,
|
||||
block_size: 4,
|
||||
causal: true,
|
||||
};
|
||||
let head = AetherTemporalHead::new(&cfg).expect("construct");
|
||||
let (q, k, vt) = make_qkv(32, 2, 2, 16);
|
||||
let out = head.forward(&q, &k, &vt).expect("forward");
|
||||
assert_eq!(out.shape(), (32, 2, 16));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn dense_backend_returns_typed_error() {
|
||||
let cfg = TemporalHeadConfig {
|
||||
backend: TemporalBackendKind::Dense,
|
||||
q_heads: 4,
|
||||
kv_heads: 1,
|
||||
head_dim: 32,
|
||||
window: 32,
|
||||
block_size: 16,
|
||||
causal: true,
|
||||
};
|
||||
let err = AetherTemporalHead::new(&cfg).err().expect("dense rejected");
|
||||
matches!(err, TemporalError::DenseBackendNotImplemented);
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn invalid_gqa_ratio_rejected_at_construction() {
|
||||
let cfg = TemporalHeadConfig {
|
||||
backend: TemporalBackendKind::SparseGqa,
|
||||
q_heads: 5,
|
||||
kv_heads: 2, // 5 % 2 != 0
|
||||
head_dim: 16,
|
||||
window: 8,
|
||||
block_size: 4,
|
||||
causal: true,
|
||||
};
|
||||
let err = AetherTemporalHead::new(&cfg).err().expect("rejected");
|
||||
matches!(err, TemporalError::InvalidConfig(_));
|
||||
}
|
||||
|
||||
#[test]
|
||||
fn long_window_at_aether_roadmap_target() {
|
||||
// ADR-096 §3.1 roadmap target: 10 s @ 100 Hz = 1000 frames. Verify
|
||||
// the kernel actually runs at this length so the long-window claim
|
||||
// is more than aspirational.
|
||||
let cfg = TemporalHeadConfig {
|
||||
backend: TemporalBackendKind::SparseGqa,
|
||||
q_heads: 4,
|
||||
kv_heads: 1,
|
||||
head_dim: 16,
|
||||
window: 64,
|
||||
block_size: 32,
|
||||
causal: true,
|
||||
};
|
||||
let head = AetherTemporalHead::new(&cfg).expect("construct");
|
||||
let (q, k, vt) = make_qkv(1000, 4, 1, 16);
|
||||
let out = head.forward(&q, &k, &vt).expect("forward at N=1000");
|
||||
assert_eq!(out.shape(), (1000, 4, 16));
|
||||
}
|
||||
Reference in New Issue
Block a user