forgo.cloud
Sign in
Repo workspace

forkjoin-ai/gnosis

Pneuma Rust Mesh — Cross-File Contracts

distributed-inference/PNEUMA_CONTRACTS.md
forkjoin-ai/gnosis

Pneuma Rust Mesh — Cross-File Contracts

Shared interfaces every agent targets. Each agent implements one file against these signatures; the reconciler wires them together. Stubs are fine for cross-file deps — don't try to satisfy them yourself, just call the documented function.

File ownership

File Owner What it exports
src/math_pneuma.rs Agent A conv1d, conv_transpose1d, cross_attention
src/model_whisper_encoder.rs Agent B WhisperEncoder, WhisperConfig use
src/model_whisper_decoder.rs Agent C WhisperDecoder, KV cache for self+cross
src/loader_whisper.rs Agent D WhisperConfig, parse_whisper_config, name resolution
src/wasm_bindings_whisper.rs Agent E WasmWhisperPipeline (JS-callable)
open-source/bitwise/scripts/encode-whisper-knot.ts Agent F safetensors → .knot for Whisper

Math kernel contracts (Agent A delivers, B/C call)

/// 1D convolution. Whisper encoder uses two stacked Conv1d on mel input:
///   layer 1: in_ch=80 (mels), out_ch=hidden, kernel=3, stride=1, pad=1
///   layer 2: in_ch=hidden,    out_ch=hidden, kernel=3, stride=2, pad=1
/// Layout: x is (in_ch, in_len) row-major; w is (out_ch, in_ch, k) row-major;
///         out is (out_ch, out_len) row-major.
pub fn conv1d(
    out: &mut [f32],
    x: &[f32],
    w: &[f32],
    bias: Option<&[f32]>,
    in_ch: usize, out_ch: usize, in_len: usize,
    k: usize, stride: usize, pad: usize,
);

/// Output length = (in_len + 2*pad - k) / stride + 1
pub fn conv1d_out_len(in_len: usize, k: usize, stride: usize, pad: usize) -> usize {
    (in_len + 2 * pad - k) / stride + 1
}

/// 1D transposed convolution (vocoder upsampling, TTS Track B).
/// Layout: w is (in_ch, out_ch, k); out length = (in_len - 1) * stride + k - 2*pad.
pub fn conv_transpose1d(
    out: &mut [f32],
    x: &[f32],
    w: &[f32],
    bias: Option<&[f32]>,
    in_ch: usize, out_ch: usize, in_len: usize,
    k: usize, stride: usize, pad: usize,
);

/// Cross-attention helper. Q from decoder; K, V cached from encoder output.
///   q:    (seq_q,  num_heads, head_dim) — the current decoder query
///   k:    (seq_kv, num_heads, head_dim) — encoder K (cached, no growth)
///   v:    (seq_kv, num_heads, head_dim) — encoder V (cached, no growth)
///   out:  (seq_q,  num_heads, head_dim) — attention output
/// No causal mask (decoder always sees full encoder).
pub fn cross_attention(
    out: &mut [f32],
    q: &[f32],
    k: &[f32],
    v: &[f32],
    seq_q: usize, seq_kv: usize,
    num_heads: usize, head_dim: usize,
);

layer_norm and gelu_inplace are already in math.rs — Agent A does not duplicate them.

WhisperConfig contract (Agent D delivers, B/C/E read)

#[derive(Debug, Clone)]
pub struct WhisperConfig {
    // audio input
    pub n_mels: usize,        // 80 for all whisper variants
    pub n_audio_ctx: usize,   // 1500 (encoded audio sequence length after conv stride 2)
    // encoder
    pub encoder_layers: usize,      // 4 (tiny), 6 (base), 12 (small)
    pub encoder_hidden_dim: usize,  // 384, 512, 768
    pub encoder_heads: usize,       // 6, 8, 12
    pub encoder_intermediate_size: usize,  // 4 * hidden_dim
    // decoder
    pub decoder_layers: usize,
    pub decoder_hidden_dim: usize,
    pub decoder_heads: usize,
    pub decoder_intermediate_size: usize,
    pub n_text_ctx: usize,    // 448 max decode length
    pub vocab_size: usize,    // 51865 (multilingual) or 51864 (en-only)
    pub eps: f32,             // 1e-5
}

pub fn parse_whisper_config(meta: &serde_json::Value) -> Result<WhisperConfig, String>;

Tensor name conventions (Agent D delivers, B/C call)

.knot metadata tensors block uses these keys for Whisper:

encoder_conv1.weight              shape (hidden, n_mels, 3)
encoder_conv1.bias                shape (hidden,)
encoder_conv2.weight              shape (hidden, hidden, 3)
encoder_conv2.bias                shape (hidden,)
encoder_pos_embed.weight          shape (n_audio_ctx, hidden)
encoder_blk_{i}_attn_q.weight     shape (hidden, hidden)
encoder_blk_{i}_attn_q.bias       shape (hidden,)
encoder_blk_{i}_attn_k.weight     shape (hidden, hidden)
encoder_blk_{i}_attn_v.weight     shape (hidden, hidden)
encoder_blk_{i}_attn_v.bias       shape (hidden,)
encoder_blk_{i}_attn_out.weight   shape (hidden, hidden)
encoder_blk_{i}_attn_out.bias     shape (hidden,)
encoder_blk_{i}_ln1.weight        shape (hidden,)        # pre-attn LN
encoder_blk_{i}_ln1.bias          shape (hidden,)
encoder_blk_{i}_mlp_fc1.weight    shape (4*hidden, hidden)
encoder_blk_{i}_mlp_fc1.bias      shape (4*hidden,)
encoder_blk_{i}_mlp_fc2.weight    shape (hidden, 4*hidden)
encoder_blk_{i}_mlp_fc2.bias      shape (hidden,)
encoder_blk_{i}_ln2.weight        shape (hidden,)        # pre-mlp LN
encoder_blk_{i}_ln2.bias          shape (hidden,)
encoder_ln_final.weight           shape (hidden,)
encoder_ln_final.bias             shape (hidden,)

decoder_token_embed.weight        shape (vocab, hidden)
decoder_pos_embed.weight          shape (n_text_ctx, hidden)
decoder_blk_{i}_self_q.weight     same pattern as encoder
decoder_blk_{i}_self_k.weight
decoder_blk_{i}_self_v.weight
decoder_blk_{i}_self_out.weight
decoder_blk_{i}_self_ln.weight
decoder_blk_{i}_cross_q.weight    cross-attn Q from decoder
decoder_blk_{i}_cross_k.weight    cross-attn K from encoder
decoder_blk_{i}_cross_v.weight
decoder_blk_{i}_cross_out.weight
decoder_blk_{i}_cross_ln.weight
decoder_blk_{i}_mlp_fc1.weight
decoder_blk_{i}_mlp_fc2.weight
decoder_blk_{i}_mlp_ln.weight
decoder_ln_final.weight
# LM head shares weights with token_embed (Whisper uses tied embeddings).

Helper:

pub fn whisper_tensor_name(stage: &str, layer: Option<usize>, role: &str) -> String;
// e.g. whisper_tensor_name("encoder", Some(0), "attn_q.weight") → "encoder_blk_0_attn_q.weight"
//      whisper_tensor_name("encoder", None,    "ln_final.bias") → "encoder_ln_final.bias"

WhisperEncoder contract (Agent B)

pub struct WhisperEncoder {
    config: WhisperConfig,
    backend: Arc<dyn WeightBackend>,  // existing trait from src/loader.rs
    // any internal buffers
}

impl WhisperEncoder {
    pub fn new(config: WhisperConfig, backend: Arc<dyn WeightBackend>) -> Self;

    /// mel: row-major (n_mels, n_frames=3000)
    /// returns: row-major (n_audio_ctx=1500, hidden_dim) — audio context tensor
    pub fn forward(&mut self, mel: &[f32]) -> Vec<f32>;
}

Forward pass:

  1. conv1d(mel) → gelu → conv1d(stride 2) → gelu → shape (hidden, 1500)
  2. transpose to (1500, hidden), add positional embedding
  3. for each encoder layer:
    • residual = x; x = ln1(x); x = self_attention(x); x = residual + x
    • residual = x; x = ln2(x); x = mlp_fc1(x); gelu(x); x = mlp_fc2(x); x = residual + x
  4. x = ln_final(x); return.

WhisperDecoder contract (Agent C)

pub struct WhisperDecoder {
    config: WhisperConfig,
    backend: Arc<dyn WeightBackend>,
    // KV caches: per-layer self-attention K,V grow as we decode
    // cross-attention K,V cached once per encoded audio (constant size)
}

impl WhisperDecoder {
    pub fn new(config: WhisperConfig, backend: Arc<dyn WeightBackend>) -> Self;

    /// Pre-compute cross-attention K,V from a fresh encoder output.
    /// Call once per utterance before any decode_step.
    pub fn prepare_audio_context(&mut self, encoder_output: &[f32]);

    /// StructuralErrorgle autoregressive step.
    /// token: previous token id (or BOS for first step)
    /// pos: position in the sequence (0 = BOS)
    /// returns: vocab_size logits for the next token
    pub fn decode_step(&mut self, token: u32, pos: usize) -> Vec<f32>;

    /// Reset self-attn KV cache between utterances.
    pub fn reset(&mut self);
}

WasmWhisperPipeline contract (Agent E)

#[wasm_bindgen]
pub struct WasmWhisperPipeline { /* ... */ }

#[wasm_bindgen]
impl WasmWhisperPipeline {
    /// Pass the same backend handle used for LLM weights.
    #[wasm_bindgen(constructor)]
    pub fn new(backend: WasmHostBackend, config_json: String) -> Self;

    /// mel: Float32Array of length 80 * 3000
    /// returns: Float32Array of length 1500 * hidden_dim
    pub fn encode(&mut self, mel: &[f32]) -> Vec<f32>;

    /// returns: Vec<u32> of generated tokens (greedy decoding for v1)
    pub fn generate(&mut self, encoder_output: &[f32], max_tokens: u32, eos_token: u32) -> Vec<u32>;
}

NAPI wrapper mirrors the WASM API for Node use.

Weight converter contract (Agent F)

// open-source/bitwise/scripts/encode-whisper-knot.ts

/** Reads HF safetensors model + tokenizer, emits .knot file. */
export async function encodeWhisperKnot(
  hfModelDir: string,        // local dir with model.safetensors + config.json
  outPath: string,           // where to write whisper-tiny.knot
  options?: { fp16?: boolean; quantize?: 'none' | 'q4k' },
): Promise<void>;

// CLI:
// bun encode-whisper-knot.ts --in /path/to/openai-whisper-tiny --out whisper-tiny.knot

Mapping: HF tensor name model.encoder.layers.0.self_attn.q_proj.weight maps to knot tensor encoder_blk_0_attn_q.weight. Use the table in "Tensor name conventions" above.

Validation gates (per file)

Each agent's deliverable must pass:

  • cargo check (Rust files) or bun run typecheck (TS file)
  • All public functions documented with /// doc comments
  • Existing patterns honored (NEON SIMD where the file's neighbors use it, no new crates)
  • No todo!() or unimplemented!() in returned code — stubs OK if marked // AGENT-STUB: depends on Agent X's deliverable, replace once landed

What NOT to do

  • Do not modify src/math.rs (use src/math_pneuma.rs only — math.rs already has layer_norm and gelu_inplace added by the orchestrator)
  • Do not edit Cargo.toml (no new crates needed; if you think you need one, raise it in your final report instead)
  • Do not edit any file owned by another agent
  • Do not run cargo build (the workspace has heavy dependencies; just cargo check --no-default-features if you need a syntax check)