diff --git a/Cargo.lock b/Cargo.lock index 6a41fdaf..0c541dd2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3800,6 +3800,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "openinfer-higgs-audio" +version = "0.1.0" +dependencies = [ + "anyhow", + "cudarc", + "half", + "openinfer-core", + "openinfer-kernels", + "openinfer-kv-cache", + "serde", + "serde_json", + "thiserror 2.0.18", +] + [[package]] name = "openinfer-kernels" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index 33876140..db3fb6c2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -17,6 +17,7 @@ members = [ "openinfer-kimi-k2", "openinfer-qwen3-4b", "openinfer-qwen35-4b", + "openinfer-higgs-audio", "openinfer-kv-cache", "openinfer-kv-offload", # ---- openinfer-comm (EP all-to-all) ---- @@ -129,6 +130,7 @@ openinfer-kernels = { path = "openinfer-kernels" } openinfer-kimi-k2 = { path = "openinfer-kimi-k2" } openinfer-qwen3-4b = { path = "openinfer-qwen3-4b" } openinfer-qwen35-4b = { path = "openinfer-qwen35-4b" } +openinfer-higgs-audio = { path = "openinfer-higgs-audio" } openinfer-deepseek-v2-lite = { path = "openinfer-deepseek-v2-lite" } openinfer-vllm-frontend = { path = "openinfer-vllm-frontend" } openinfer-vllm-support = { path = "openinfer-vllm-support" } diff --git a/openinfer-higgs-audio/Cargo.toml b/openinfer-higgs-audio/Cargo.toml new file mode 100644 index 00000000..5dcbcdc7 --- /dev/null +++ b/openinfer-higgs-audio/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "openinfer-higgs-audio" +license = "Apache-2.0" +version = "0.1.0" +edition = "2024" +autobenches = false +autobins = false +autotests = false + +[dependencies] +# Pure Rust / CPU — always compiled (Mac can build) +anyhow = { workspace = true } +half = { workspace = true } +serde = { workspace = true } +serde_json = { workspace = true } +thiserror = { workspace = true } +# GPU stack — only under higgs-audio feature +openinfer-core = { workspace = true, optional = true } +openinfer-kernels = { workspace = true, optional = true } +openinfer-kv-cache = { workspace = true, optional = true } +cudarc = { workspace = true, optional = true } + +[features] +default = [] +higgs-audio = [ + "dep:openinfer-core", + "dep:openinfer-kernels", + "dep:openinfer-kv-cache", + "dep:cudarc", +] + +[lints] +workspace = true + +[dev-dependencies] +safetensors = { workspace = true } + +[[test]] +name = "backbone_parity" +required-features = ["higgs-audio"] diff --git a/openinfer-higgs-audio/src/backbone.rs b/openinfer-higgs-audio/src/backbone.rs new file mode 100644 index 00000000..7a91e63f --- /dev/null +++ b/openinfer-higgs-audio/src/backbone.rs @@ -0,0 +1,552 @@ +//! Higgs backbone forward — mirrors `openinfer-qwen3-4b` weight loading and +//! prefill path, adapted for `body.*` weight names and `text_config` from +//! Higgs config.json. +//! +//! Compiled only under `#[cfg(feature = "higgs-audio")]`. + +use anyhow::{Context, Result}; +use cudarc::driver::CudaSlice; +use half::bf16; +use log::{debug, info}; +use std::time::Instant; + +use crate::config::HiggsConfig; +use openinfer_core::kv_pool::KvLayout; +use openinfer_core::ops; +use openinfer_core::ops::PrefillPagedPlan; +use openinfer_core::tensor::{DeviceContext, DeviceMatrix, DeviceVec, HiddenStates}; +use openinfer_core::weight_loader::{ + deserialize_shards, load_shard_info, load_tensor_1d, load_tensor_2d, mmap_shards, + precompute_rope, +}; + +// ── Data structures (identical to qwen3, minus LoRA / TP / CUDA graph) ────── + +/// Attention layer weights. +/// +/// QKV stored as a single concatenated matrix `[q_dim + 2*kv_dim, hidden_size]`. +/// Individual projections accessed via row offsets (zero extra memory). +pub(super) struct Attention { + pub(super) qkv_proj: DeviceMatrix, + pub(super) o_proj: DeviceMatrix, + pub(super) q_norm: DeviceVec, + pub(super) k_norm: DeviceVec, + pub(super) q_dim: usize, + pub(super) kv_dim: usize, +} + +/// MLP layer weights. +/// +/// Gate+Up stored as a single concatenated matrix `[2*intermediate_size, hidden_size]`. +pub(super) struct MLP { + pub(super) gate_up_proj: DeviceMatrix, + pub(super) down_proj: DeviceMatrix, +} + +/// Transformer block. +pub(super) struct TransformerBlock { + pub(super) input_layernorm: DeviceVec, + pub(super) attention: Attention, + pub(super) post_attention_layernorm: DeviceVec, + pub(super) mlp: MLP, +} + +/// Pre-allocated scratch buffers for one prefill forward pass. +/// Created once per forward, eliminating per-layer allocation overhead. +pub(super) struct PrefillBuffers { + pub(super) hidden_out: HiddenStates, // hidden_dim × seq_len + pub(super) normed: HiddenStates, // hidden_dim × seq_len (reused for normed2) + pub(super) q_batch: HiddenStates, // q_dim × seq_len + pub(super) k_batch: HiddenStates, // kv_dim × seq_len + pub(super) v_batch: HiddenStates, // kv_dim × seq_len + pub(super) o_buf: HiddenStates, // hidden_dim × seq_len (reused for mlp_out) + pub(super) gate_out: HiddenStates, // inter_dim × seq_len + pub(super) up_out: HiddenStates, // inter_dim × seq_len + pub(super) act_out: HiddenStates, // inter_dim × seq_len + pub(super) attn_output: HiddenStates, // q_dim × seq_len +} + +impl PrefillBuffers { + fn new( + ctx: &DeviceContext, + hidden_dim: usize, + q_dim: usize, + kv_dim: usize, + inter_dim: usize, + seq_len: usize, + ) -> Result { + Ok(Self { + hidden_out: HiddenStates::zeros(ctx, hidden_dim, seq_len)?, + normed: HiddenStates::zeros(ctx, hidden_dim, seq_len)?, + q_batch: HiddenStates::zeros(ctx, q_dim, seq_len)?, + k_batch: HiddenStates::zeros(ctx, kv_dim, seq_len)?, + v_batch: HiddenStates::zeros(ctx, kv_dim, seq_len)?, + o_buf: HiddenStates::zeros(ctx, hidden_dim, seq_len)?, + gate_out: HiddenStates::zeros(ctx, inter_dim, seq_len)?, + up_out: HiddenStates::zeros(ctx, inter_dim, seq_len)?, + act_out: HiddenStates::zeros(ctx, inter_dim, seq_len)?, + attn_output: HiddenStates::zeros(ctx, q_dim, seq_len)?, + }) + } +} + +// ── HiggsBackbone ────────────────────────────────────────────────────────── + +/// Higgs backbone model — weights and config only. +/// +/// Holds GPU-resident weights for embed_tokens, 36 transformer layers, +/// final RMSNorm, and precomputed RoPE cache. +pub struct HiggsBackbone { + pub(super) ctx: DeviceContext, + pub(super) config: crate::config::TextConfig, + pub(super) embed_tokens: DeviceMatrix, + pub(super) layers: Vec, + pub(super) norm: DeviceVec, + pub(super) cos_cache: DeviceVec, + pub(super) sin_cache: DeviceVec, +} + +// SAFETY: Each backbone instance is pinned to a single CUDA device and is only +// driven from one thread at a time. +unsafe impl Send for HiggsBackbone {} +unsafe impl Sync for HiggsBackbone {} + +impl HiggsBackbone { + /// Load Higgs backbone weights from safetensors. + /// + /// Reads Higgs `config.json` → extracts `text_config` for architecture + /// params. Loads all `body.*` tensors (36 layers × 11 weights) and both + /// tied embeddings. Precomputes RoPE cache on GPU. + pub fn from_safetensors(model_path: &str, device_ordinal: usize) -> Result { + info!("Loading Higgs backbone from: {}", model_path); + debug!("Initializing GPU device {}", device_ordinal); + let ctx = DeviceContext::new_with_device(device_ordinal)?; + + // Load Higgs config → extract text_config + let model_path = std::path::Path::new(model_path); + let higgs_config = + HiggsConfig::from_path(model_path).context("failed to parse Higgs config.json")?; + let config = higgs_config.text_config; + + // Validate architecture facts before loading weights + anyhow::ensure!( + config.num_hidden_layers == 36, + "expected 36 hidden layers, got {}", + config.num_hidden_layers + ); + anyhow::ensure!( + config.hidden_size == 2560, + "expected hidden_size=2560, got {}", + config.hidden_size + ); + + // Load safetensors + let model_path_str = model_path + .to_str() + .ok_or_else(|| anyhow::anyhow!("model_path is not valid UTF-8"))?; + let (shard_paths, weight_map) = load_shard_info(model_path_str)?; + debug!("Loading {} safetensor shard(s)", shard_paths.len()); + let mmaps = mmap_shards(&shard_paths)?; + let shards = deserialize_shards(&mmaps)?; + + let t_gpu = Instant::now(); + + // Load tied embeddings (used for both embed_tokens and lm_head) + debug!("Loading embeddings to GPU"); + let embed_tokens = load_tensor_2d( + &ctx, + &shards, + &weight_map, + "tied.embedding.text_embedding.weight", + ) + .with_context(|| { + "failed to load tied.embedding.text_embedding.weight — \ + is this a Higgs checkpoint?" + })?; + + // Load 36 layers using `body.layers.{i}.{rest}` naming + let num_layers = config.num_hidden_layers; + debug!("Loading {} layers to GPU", num_layers); + let mut layers = Vec::with_capacity(num_layers); + + for i in 0..num_layers { + let prefix = format!("body.layers.{}", i); + + let q_proj = load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.self_attn.q_proj.weight", prefix), + )?; + let k_proj = load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.self_attn.k_proj.weight", prefix), + )?; + let v_proj = load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.self_attn.v_proj.weight", prefix), + )?; + let q_dim = q_proj.rows; + let kv_dim = k_proj.rows; + let qkv_proj = DeviceMatrix::vstack(&ctx, &[&q_proj, &k_proj, &v_proj])?; + drop(q_proj); + drop(k_proj); + drop(v_proj); + + let gate_proj = load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.mlp.gate_proj.weight", prefix), + )?; + let up_proj = load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.mlp.up_proj.weight", prefix), + )?; + let gate_up_proj = DeviceMatrix::vstack(&ctx, &[&gate_proj, &up_proj])?; + drop(gate_proj); + drop(up_proj); + + let block = TransformerBlock { + input_layernorm: load_tensor_1d( + &ctx, + &shards, + &weight_map, + &format!("{}.input_layernorm.weight", prefix), + )?, + attention: Attention { + qkv_proj, + o_proj: load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.self_attn.o_proj.weight", prefix), + )?, + q_norm: load_tensor_1d( + &ctx, + &shards, + &weight_map, + &format!("{}.self_attn.q_norm.weight", prefix), + )?, + k_norm: load_tensor_1d( + &ctx, + &shards, + &weight_map, + &format!("{}.self_attn.k_norm.weight", prefix), + )?, + q_dim, + kv_dim, + }, + post_attention_layernorm: load_tensor_1d( + &ctx, + &shards, + &weight_map, + &format!("{}.post_attention_layernorm.weight", prefix), + )?, + mlp: MLP { + gate_up_proj, + down_proj: load_tensor_2d( + &ctx, + &shards, + &weight_map, + &format!("{}.mlp.down_proj.weight", prefix), + )?, + }, + }; + layers.push(block); + } + + // Load final norm + let norm = load_tensor_1d(&ctx, &shards, &weight_map, "body.norm.weight") + .with_context(|| "failed to load body.norm.weight")?; + + // Precompute RoPE cache + debug!("Precomputing RoPE cache on GPU"); + let rope_theta = config.rope_theta(); + let (cos_cache, sin_cache) = precompute_rope( + &ctx, + config.head_dim, + config.max_position_embeddings, + rope_theta, + )?; + + ctx.sync()?; + info!( + "Higgs backbone loaded in {:.0}ms", + t_gpu.elapsed().as_secs_f64() * 1e3 + ); + + Ok(Self { + ctx, + config, + embed_tokens, + layers, + norm, + cos_cache, + sin_cache, + }) + } + + /// Return the config used by this backbone (from `text_config`). + pub fn config(&self) -> &crate::config::TextConfig { + &self.config + } + + /// Return the device context for external coordination (KV allocation etc.). + pub fn device_ctx(&self) -> &DeviceContext { + &self.ctx + } + + // ── Forward pass ──────────────────────────────────────────────────── + + /// Embed a batch of token IDs into hidden states `[hidden_dim, seq_len]`. + pub(super) fn get_embeddings_batch(&self, token_ids: &[u32]) -> Result { + let seq_len = token_ids.len(); + let hidden_dim = self.config.hidden_size; + + let token_ids_gpu = self + .ctx + .stream + .clone_htod(token_ids) + .map_err(|e| anyhow::anyhow!("H2D copy failed: {}", e))?; + + let mut out = HiddenStates::zeros(&self.ctx, hidden_dim, seq_len)?; + ops::embedding_batch(&self.ctx, &self.embed_tokens, &token_ids_gpu, &mut out)?; + + Ok(out) + } + + /// Process a single transformer layer (per-token prefill). + /// + /// Steps: RMSNorm → QKV projections → paged attention (norm+RoPE+attend) → + /// O projection → fused residual+RMSNorm → MLP (gate/up → silu_mul → down) → + /// residual add. + #[allow(clippy::too_many_arguments)] + fn forward_layer( + &self, + layer_idx: usize, + layer: &TransformerBlock, + hidden: &mut HiddenStates, + kv_buffer: &CudaSlice, + layout: &KvLayout, + plan: &PrefillPagedPlan, + bufs: &mut PrefillBuffers, + ) -> Result<()> { + let num_heads = self.config.num_attention_heads; + let num_kv_heads = self.config.num_key_value_heads; + let head_dim = self.config.head_dim; + + // 1. RMSNorm → bufs.normed + ops::rms_norm_batch_into( + &self.ctx, + hidden, + &layer.input_layernorm, + self.config.rms_norm_eps, + &mut bufs.normed, + ); + + // 2. QKV projections from fused qkv_proj + let q_dim = layer.attention.q_dim; + let kv_dim = layer.attention.kv_dim; + ops::gemm_rows_into( + &self.ctx, + &layer.attention.qkv_proj, + 0, + q_dim, + &bufs.normed, + &mut bufs.q_batch, + ); + ops::gemm_rows_into( + &self.ctx, + &layer.attention.qkv_proj, + q_dim, + kv_dim, + &bufs.normed, + &mut bufs.k_batch, + ); + ops::gemm_rows_into( + &self.ctx, + &layer.attention.qkv_proj, + q_dim + kv_dim, + kv_dim, + &bufs.normed, + &mut bufs.v_batch, + ); + + // 3. Paged prefill: norm+RoPE → append K/V to paged → batch attention + ops::prefill_attention_paged_into( + &self.ctx, + &mut bufs.q_batch, + &mut bufs.k_batch, + &bufs.v_batch, + &layer.attention.q_norm, + &layer.attention.k_norm, + &self.cos_cache, + &self.sin_cache, + kv_buffer, + layout, + layer_idx, + plan, + &mut bufs.attn_output, + num_heads, + num_kv_heads, + head_dim, + self.config.rms_norm_eps, + )?; + + // 4. O projection → bufs.o_buf + ops::gemm_into( + &self.ctx, + &layer.attention.o_proj, + &bufs.attn_output, + &mut bufs.o_buf, + ); + + // 5+6. Residual add + MLP RMSNorm (fused) + openinfer_kernels::ops::fused_add_rms_norm_batch_into( + &self.ctx, + hidden, + &bufs.o_buf, + &layer.post_attention_layernorm, + self.config.rms_norm_eps, + &mut bufs.normed, + )?; + + // 7. MLP: split gate/up GEMMs → silu_mul → down + let inter_dim = self.config.intermediate_size; + ops::gemm_rows_into( + &self.ctx, + &layer.mlp.gate_up_proj, + 0, + inter_dim, + &bufs.normed, + &mut bufs.gate_out, + ); + ops::gemm_rows_into( + &self.ctx, + &layer.mlp.gate_up_proj, + inter_dim, + inter_dim, + &bufs.normed, + &mut bufs.up_out, + ); + ops::silu_mul_batch_into(&self.ctx, &bufs.gate_out, &bufs.up_out, &mut bufs.act_out)?; + ops::gemm_into( + &self.ctx, + &layer.mlp.down_proj, + &bufs.act_out, + &mut bufs.o_buf, + ); + + // 8. Residual add: hidden + mlp_out → bufs.hidden_out + ops::add_batch_into(&self.ctx, hidden, &bufs.o_buf, &mut bufs.hidden_out)?; + std::mem::swap(hidden, &mut bufs.hidden_out); + + Ok(()) + } + + /// Process all layers for a single prefill sequence. + fn process_all_layers( + &self, + mut hidden: HiddenStates, + layout: &KvLayout, + kv_buffer: &CudaSlice, + plan: &PrefillPagedPlan, + ) -> Result { + let total_tokens = hidden.seq_len; + let inter_dim = self.config.intermediate_size; + let q_dim = self.config.num_attention_heads * self.config.head_dim; + let kv_dim = self.config.num_key_value_heads * self.config.head_dim; + + let mut bufs = PrefillBuffers::new( + &self.ctx, + self.config.hidden_size, + q_dim, + kv_dim, + inter_dim, + total_tokens, + )?; + + for (layer_idx, layer) in self.layers.iter().enumerate() { + self.forward_layer( + layer_idx, + layer, + &mut hidden, + kv_buffer, + layout, + plan, + &mut bufs, + )?; + } + + Ok(hidden) + } + + /// Compute logits for all positions in hidden states: final RMSNorm + lm_head. + /// + /// Returns `HiddenStates` with shape `[vocab_size, seq_len]`. + pub fn compute_all_position_logits(&self, hidden: &HiddenStates) -> Result { + let mut normed = HiddenStates::zeros(&self.ctx, hidden.hidden_dim, hidden.seq_len)?; + ops::rms_norm_batch_into( + &self.ctx, + hidden, + &self.norm, + self.config.rms_norm_eps, + &mut normed, + ); + ops::gemm(&self.ctx, &self.embed_tokens, &normed) + } + + /// Extract last-token logits: gather the final column of `hidden`, + /// apply final RMSNorm and lm_head. Returns `[vocab_size, 1]`. + pub fn last_token_logits(&self, hidden: &HiddenStates) -> Result { + let last_idx = (hidden.seq_len - 1) as i32; + let indices_d = self + .ctx + .stream + .clone_htod(&[last_idx]) + .map_err(|e| anyhow::anyhow!("H2D copy failed: {}", e))?; + + let mut gathered = HiddenStates::zeros(&self.ctx, hidden.hidden_dim, 1)?; + ops::gather_hidden_tokens_into(&self.ctx, hidden, &indices_d, 1, &mut gathered)?; + + let mut normed = HiddenStates::zeros(&self.ctx, hidden.hidden_dim, 1)?; + ops::rms_norm_batch_into( + &self.ctx, + &gathered, + &self.norm, + self.config.rms_norm_eps, + &mut normed, + ); + + ops::gemm(&self.ctx, &self.embed_tokens, &normed) + } + + /// Full prefill forward for a single sequence. + /// + /// Takes a `kv_buffer` (allocated by caller with enough pages for this + /// sequence), a `KvLayout`, and a `PrefillPagedPlan`. Returns + /// `(hidden_states, logits)` where `hidden_states` has shape + /// `[hidden_dim, seq_len]` and `logits` has shape `[vocab_size, 1]` + /// (last-token only). + /// + /// For "echo" mode (all-position logits), call + /// `compute_all_position_logits` on the returned hidden states before + /// calling `last_token_logits`. + pub fn forward( + &self, + input_ids: &[u32], + kv_buffer: &CudaSlice, + layout: &KvLayout, + plan: &PrefillPagedPlan, + ) -> Result { + let hidden = self.get_embeddings_batch(input_ids)?; + self.process_all_layers(hidden, layout, kv_buffer, plan) + } +} diff --git a/openinfer-higgs-audio/src/config.rs b/openinfer-higgs-audio/src/config.rs new file mode 100644 index 00000000..9e1f8972 --- /dev/null +++ b/openinfer-higgs-audio/src/config.rs @@ -0,0 +1,134 @@ +use anyhow::Result; +use serde::Deserialize; +use std::fs; +use std::path::Path; + +/// Nested `rope_parameters` in text_config. +#[derive(Clone, Debug, Deserialize)] +pub struct RopeParameters { + pub rope_theta: f32, + #[serde(default)] + pub rope_type: String, +} + +/// `text_config` sub-object inside Higgs config.json. +#[derive(Clone, Debug, Deserialize)] +pub struct TextConfig { + pub hidden_size: usize, + pub intermediate_size: usize, + pub num_hidden_layers: usize, + pub num_attention_heads: usize, + pub num_key_value_heads: usize, + pub head_dim: usize, + pub vocab_size: usize, + pub rms_norm_eps: f32, + pub hidden_act: String, + pub tie_word_embeddings: bool, + #[serde(default = "default_max_position_embeddings")] + pub max_position_embeddings: usize, + pub bos_token_id: u32, + pub eos_token_id: u32, + #[serde(default)] + pub rope_parameters: Option, +} + +fn default_max_position_embeddings() -> usize { + 32768 +} + +impl TextConfig { + /// Resolve `rope_theta` from nested `rope_parameters` if present, + /// otherwise fall back to a direct `rope_theta` field (not present in Higgs + /// but supported for robustness). + pub fn rope_theta(&self) -> f32 { + if let Some(ref rp) = self.rope_parameters { + rp.rope_theta + } else { + 1_000_000.0 + } + } +} + +/// `audio_encoder_config` sub-object. +#[derive(Clone, Debug, Deserialize)] +pub struct AudioEncoderConfig { + pub num_codebooks: usize, + pub vocab_size: usize, + pub out_dim: usize, + pub use_delay_pattern: bool, + pub tie_word_embeddings: bool, + pub model_type: String, + pub encoder_type: String, +} + +/// Top-level Higgs config. +#[derive(Clone, Debug, Deserialize)] +pub struct HiggsConfig { + pub text_config: TextConfig, + pub audio_encoder_config: AudioEncoderConfig, + pub audio_token_id: i64, + pub model_type: String, +} + +impl HiggsConfig { + /// Load config from a directory containing `config.json`. + pub fn from_path(model_path: &Path) -> Result { + let config_path = model_path.join("config.json"); + let content = fs::read_to_string(&config_path)?; + let config: HiggsConfig = serde_json::from_str(&content)?; + Ok(config) + } +} + +#[cfg(test)] +mod tests { + use super::*; + use std::path::PathBuf; + + fn test_model_path() -> PathBuf { + let env_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .map(PathBuf::from) + .unwrap_or_else(|_| PathBuf::from("docs/private/higgs-audio-v3-tts-4b")); + // Resolve relative paths against the workspace root (CARGO_MANIFEST_DIR/..). + if env_path.is_absolute() { + env_path + } else { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("workspace root") + .join(&env_path) + } + } + + #[test] + fn parse_higgs_config_and_assert_facts() { + let config = HiggsConfig::from_path(&test_model_path()).expect("failed to load config"); + + // text_config facts + assert_eq!(config.text_config.hidden_size, 2560); + assert_eq!(config.text_config.intermediate_size, 9728); + assert_eq!(config.text_config.num_hidden_layers, 36); + assert_eq!(config.text_config.num_attention_heads, 32); + assert_eq!(config.text_config.num_key_value_heads, 8); + assert_eq!(config.text_config.head_dim, 128); + assert_eq!(config.text_config.vocab_size, 151936); + assert!((config.text_config.rms_norm_eps - 1e-6).abs() < 1e-10); + assert!(config.text_config.tie_word_embeddings); + assert_eq!(config.text_config.max_position_embeddings, 32768); + assert_eq!(config.text_config.hidden_act, "silu"); + + // rope_theta from nested rope_parameters + assert_eq!(config.text_config.rope_theta(), 1_000_000.0); + + // audio_encoder_config facts + assert_eq!(config.audio_encoder_config.num_codebooks, 8); + assert_eq!(config.audio_encoder_config.vocab_size, 1026); + assert_eq!(config.audio_encoder_config.out_dim, 2560); + assert!(config.audio_encoder_config.use_delay_pattern); + assert!(config.audio_encoder_config.tie_word_embeddings); + + // top-level + assert_eq!(config.audio_token_id, -100); + assert_eq!(config.model_type, "higgs_multimodal_qwen3"); + } +} diff --git a/openinfer-higgs-audio/src/lib.rs b/openinfer-higgs-audio/src/lib.rs new file mode 100644 index 00000000..bf414d3f --- /dev/null +++ b/openinfer-higgs-audio/src/lib.rs @@ -0,0 +1,5 @@ +pub mod config; +pub mod weights; // weight name mapping (pure logic) + +#[cfg(feature = "higgs-audio")] +pub mod backbone; // GPU forward, only compiled under feature diff --git a/openinfer-higgs-audio/src/weights.rs b/openinfer-higgs-audio/src/weights.rs new file mode 100644 index 00000000..0a8bd6cf --- /dev/null +++ b/openinfer-higgs-audio/src/weights.rs @@ -0,0 +1,157 @@ +/// Slot that a Higgs checkpoint tensor name maps to in the backbone. +#[derive(Clone, Debug, PartialEq, Eq)] +pub enum BackboneSlot { + /// `tied.embedding.text_embedding.weight` — used as both embed_tokens and lm_head (tied). + EmbedTokens, + /// `tied.head.text_head.weight` — lm_head (tied with embed_tokens, value should match). + LmHead, + /// `body.layers.{i}.{rest}` — transformer block at layer index i. + Layer { layer_idx: usize, rest: String }, + /// `body.norm.weight` — final RMSNorm. + FinalNorm, +} + +/// Map a Higgs checkpoint tensor name to its backbone slot. +/// Returns `None` for non-backbone tensors (audio/codec/modality). +pub fn map_backbone(name: &str) -> Option { + // tied.embedding.text_embedding.weight + if name == "tied.embedding.text_embedding.weight" { + return Some(BackboneSlot::EmbedTokens); + } + // tied.head.text_head.weight + if name == "tied.head.text_head.weight" { + return Some(BackboneSlot::LmHead); + } + // body.norm.weight + if name == "body.norm.weight" { + return Some(BackboneSlot::FinalNorm); + } + // body.layers.{i}.{rest} + if let Some(rest) = name.strip_prefix("body.layers.") { + // rest looks like "0.self_attn.q_proj.weight" or "0.input_layernorm.weight" + if let Some(dot_pos) = rest.find('.') { + let layer_str = &rest[..dot_pos]; + if let Ok(layer_idx) = layer_str.parse::() { + let rest = rest[dot_pos + 1..].to_string(); + return Some(BackboneSlot::Layer { layer_idx, rest }); + } + } + } + // Everything else: audio/codec/modality — skip. + None +} + +#[cfg(test)] +mod tests { + use super::*; + use std::collections::HashSet; + use std::path::PathBuf; + + fn test_model_path() -> PathBuf { + let env_path = std::env::var("OPENINFER_TEST_MODEL_PATH") + .map(PathBuf::from) + .unwrap_or_else(|_| PathBuf::from("docs/private/higgs-audio-v3-tts-4b")); + if env_path.is_absolute() { + env_path + } else { + PathBuf::from(env!("CARGO_MANIFEST_DIR")) + .parent() + .expect("workspace root") + .join(&env_path) + } + } + + fn load_index_tensor_names(model_path: &PathBuf) -> Vec { + let index_path = model_path.join("model.safetensors.index.json"); + let content = std::fs::read_to_string(&index_path).expect("failed to read index"); + let index: serde_json::Value = + serde_json::from_str(&content).expect("failed to parse index"); + index["weight_map"] + .as_object() + .expect("weight_map not an object") + .keys() + .cloned() + .collect() + } + + #[test] + fn backbone_tensors_map_to_exactly_398_unique_slots() { + let names = load_index_tensor_names(&test_model_path()); + + let mut backbone_count = 0usize; + let mut seen_slots = HashSet::new(); + let mut embed_seen = false; + let mut norm_seen = false; + + for name in &names { + if let Some(slot) = map_backbone(name) { + backbone_count += 1; + match &slot { + BackboneSlot::EmbedTokens => embed_seen = true, + BackboneSlot::LmHead => {} // may or may not exist + BackboneSlot::FinalNorm => norm_seen = true, + BackboneSlot::Layer { layer_idx, rest } => { + seen_slots.insert((*layer_idx, rest.clone())); + } + } + } + } + + // 36 layers × 11 tensors per layer = 396 + embed + norm = 398 + // (lm_head is tied with embed_tokens, no separate tied.head.text_head.weight + // in this checkpoint) + assert_eq!(backbone_count, 398, "expected exactly 398 backbone tensors"); + assert!(embed_seen, "embed_tokens not found"); + assert!(norm_seen, "final norm not found"); + + // Check all 36 layers present with all 11 components + let expected_components: &[&str] = &[ + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.v_proj.weight", + "self_attn.o_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + "input_layernorm.weight", + "post_attention_layernorm.weight", + "mlp.gate_proj.weight", + "mlp.up_proj.weight", + "mlp.down_proj.weight", + ]; + + for layer_idx in 0..36 { + for comp in expected_components { + assert!( + seen_slots.contains(&(layer_idx, comp.to_string())), + "missing layer {layer_idx} component {comp}" + ); + } + } + + assert_eq!(seen_slots.len(), 396, "expected 396 unique layer slots"); + } + + #[test] + fn audio_codec_tensors_all_skip() { + let names = load_index_tensor_names(&test_model_path()); + + for name in &names { + if name.starts_with("body.") + || name == "tied.embedding.text_embedding.weight" + || name == "tied.head.text_head.weight" + { + // These are backbone — should map to Some + assert!( + map_backbone(name).is_some(), + "backbone tensor {name} returned None" + ); + } else { + // Everything else (audio/codec/modality) should be None + assert!( + map_backbone(name).is_none(), + "non-backbone tensor {name} should return None" + ); + } + } + } +} diff --git a/openinfer-higgs-audio/tests/backbone_parity.rs b/openinfer-higgs-audio/tests/backbone_parity.rs new file mode 100644 index 00000000..a04cc075 --- /dev/null +++ b/openinfer-higgs-audio/tests/backbone_parity.rs @@ -0,0 +1,298 @@ +//! Higgs backbone parity gate — validates that our backbone forward produces +//! logits within the bf16 noise floor of a pre-computed golden (from the +//! actual Higgs checkpoint via chat-template forward). +//! +//! The golden (`test_data/higgs/backbone_golden.safetensors`) was produced by +//! converting `backbone_golden.pt` to safetensors: +//! ```python +//! import torch +//! from safetensors.torch import save_file +//! g = torch.load("test_data/higgs/backbone_golden.pt", map_location="cpu", weights_only=False) +//! save_file({"input_ids": g["input_ids"].to(torch.int64).contiguous(), +//! "logits": g["logits"].to(torch.float32).contiguous(), +//! "hidden": g["hidden_states"].to(torch.float32).contiguous()}, +//! "test_data/higgs/backbone_golden.safetensors") +//! ``` +//! +//! Assertions: +//! * top-64 logprobs comparison: regret (reference has clear winner, we pick +//! wrong), mean delta, p99 delta. +//! * mean ≤ ~0.06 nat, p99 ≤ ~0.20 nat — tolerances calibrated from the +//! bf16 noise floor (same pattern as qwen3 golden gate). +//! * absolute max delta is printed but NOT asserted (coverage-unstable). +//! +//! Requires: CUDA GPU, Higgs checkpoint weights, and backbone_golden.safetensors. +//! Skipped cleanly when the model or golden is absent. + +use half::bf16; +use openinfer_core::kv_pool::KvLayout; +use openinfer_core::ops::PrefillPagedPlan; +use openinfer_higgs_audio::backbone::HiggsBackbone; +use safetensors::SafeTensors; +use std::path::Path; + +const GOLDEN_PATH: &str = concat!( + env!("CARGO_MANIFEST_DIR"), + "/../test_data/higgs/backbone_golden.safetensors" +); + +/// Number of top logprobs to compare (matching qwen3's LOGPROBS). +const LOGPROBS: usize = 64; +/// Max acceptable regret: how far below golden's argmax our pick may sit. +const MARGIN_TOL: f32 = 0.20; +/// Mean delta bound (systematic drift trips this). +const MEAN_TOL: f32 = 0.06; +/// P99 delta bound (spread inflation trips this). +const P99_TOL: f32 = 0.20; +/// Head depth for the logprob tolerance check. +const HEAD_K: usize = 8; + +fn model_path_or_skip() -> Option { + match std::env::var("OPENINFER_TEST_MODEL_PATH") { + Ok(path) => Some(path), + Err(_) => { + eprintln!( + "skipping backbone parity — set OPENINFER_TEST_MODEL_PATH \ + to Higgs checkpoint directory" + ); + None + } + } +} + +fn load_golden() -> Option<(Vec, Vec, Vec)> { + let path = Path::new(GOLDEN_PATH); + if !path.exists() { + eprintln!( + "skipping backbone parity — golden file not found at {}. \ + Convert backbone_golden.pt to safetensors first.", + GOLDEN_PATH + ); + return None; + } + let data = std::fs::read(path).expect("failed to read golden safetensors"); + let tensors = SafeTensors::deserialize(&data).expect("failed to deserialize golden"); + + let ids = tensors.tensor("input_ids").expect("missing input_ids"); + let logits_view = tensors.tensor("logits").expect("missing logits"); + let hidden_view = tensors.tensor("hidden").expect("missing hidden"); + + let input_ids: Vec = ids + .data() + .chunks_exact(8) + .map(|b| i64::from_le_bytes(b.try_into().unwrap()) as u32) + .collect(); + + let logits: Vec = logits_view + .data() + .chunks_exact(4) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(); + + let hidden: Vec = hidden_view + .data() + .chunks_exact(4) + .map(|b| f32::from_le_bytes(b.try_into().unwrap())) + .collect(); + + Some((input_ids, logits, hidden)) +} + +/// Top-64 logprobs regret / mean / p99 suite (same pattern as qwen3 golden gate). +fn check_logprobs(our_logprobs: &[f32], ref_logprobs: &[f32]) { + assert_eq!( + our_logprobs.len(), + ref_logprobs.len(), + "vocab size mismatch" + ); + let vocab = our_logprobs.len(); + + // Find top-64 indices in our sorted order for consistent coverage + let mut our_sorted: Vec<(f32, usize)> = our_logprobs + .iter() + .copied() + .enumerate() + .map(|(i, v)| (v, i)) + .collect(); + our_sorted.sort_unstable_by(|a, b| b.0.partial_cmp(&a.0).unwrap()); + let top_indices: Vec = our_sorted.iter().take(LOGPROBS).map(|(_, i)| *i).collect(); + + // Regret: for each top token, if golden has a clear winner (margin > MARGIN_TOL), + // ensure we pick the same token. + let golden_best_idx = ref_logprobs + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) + .map(|(i, _)| i) + .unwrap(); + let golden_best = ref_logprobs[golden_best_idx]; + let golden_second = ref_logprobs + .iter() + .enumerate() + .filter(|(i, _)| *i != golden_best_idx) + .map(|(_, v)| *v) + .max_by(|a, b| a.partial_cmp(b).unwrap()) + .unwrap_or(f32::NEG_INFINITY); + let margin = golden_best - golden_second; + + if margin > MARGIN_TOL { + // Golden has a clear winner — we must pick the same token + let our_best_idx = our_logprobs + .iter() + .enumerate() + .max_by(|(_, a), (_, b)| a.partial_cmp(b).unwrap()) + .map(|(i, _)| i) + .unwrap(); + if our_best_idx != golden_best_idx { + // Compute regret: difference between golden's best and what golden + // assigns to our choice + let regret = golden_best - ref_logprobs[our_best_idx]; + assert!( + regret <= MARGIN_TOL, + "regret {:.4} > {:.4}: golden best token={golden_best_idx} (logprob={golden_best:.4}), \ + our best token={our_best_idx} (golden assigns it logprob={:.4})", + regret, + MARGIN_TOL, + ref_logprobs[our_best_idx] + ); + } + } + + // Delta stats on head tokens + let mut deltas: Vec = top_indices + .iter() + .take(HEAD_K) + .map(|&i| (our_logprobs[i] - ref_logprobs[i]).abs()) + .collect(); + deltas.sort_unstable_by(|a, b| a.partial_cmp(b).unwrap()); + + let n = deltas.len(); + let mean_delta: f32 = deltas.iter().sum::() / n as f32; + let p99_idx = ((n - 1) as f64 * 0.99).ceil() as usize; + let p99_delta = deltas[p99_idx.min(n - 1)]; + let max_delta = deltas[n - 1]; + + eprintln!( + "delta stats (top-{HEAD_K} of top-{LOGPROBS}): \ + mean={mean_delta:.6}, p99={p99_delta:.6}, max={max_delta:.6}" + ); + + assert!( + mean_delta <= MEAN_TOL, + "mean delta {mean_delta:.6} > {MEAN_TOL} (systematic drift)" + ); + assert!( + p99_delta <= P99_TOL, + "p99 delta {p99_delta:.6} > {P99_TOL} (spread inflation)" + ); + // max_delta is reported but NOT asserted (coverage-unstable) +} + +#[test] +fn backbone_parity_against_golden() { + let model_path = match model_path_or_skip() { + Some(p) => p, + None => return, + }; + let (input_ids, golden_logits, _golden_hidden) = match load_golden() { + Some(g) => g, + None => return, + }; + + let seq_len = input_ids.len(); + assert!(seq_len > 0, "golden input_ids is empty"); + eprintln!("golden input_ids length: {seq_len}"); + + // Load the Higgs backbone + let backbone = + HiggsBackbone::from_safetensors(&model_path, 0).expect("failed to load Higgs backbone"); + + let config = backbone.config(); + assert_eq!(config.hidden_size, 2560); + assert_eq!(config.num_hidden_layers, 36); + assert_eq!(config.num_attention_heads, 32); + assert_eq!(config.num_key_value_heads, 8); + assert_eq!(config.head_dim, 128); + + // Set up KV buffer: one page big enough for the entire sequence + let page_size = seq_len.next_power_of_two().max(16); + let layout = KvLayout::new( + config.num_hidden_layers, + config.num_key_value_heads, + config.head_dim, + page_size, + ); + eprintln!( + "KV layout: page_size={page_size}, page_stride={} elements ({:.1} MB)", + layout.page_stride, + layout.page_stride as f64 * 2.0 / 1e6 + ); + + let ctx = backbone.device_ctx(); + let kv_buffer: cudarc::driver::CudaSlice = ctx + .stream + .alloc_zeros(layout.page_stride) + .expect("failed to allocate KV buffer"); + + // Create prefill plan for single sequence + let plan = PrefillPagedPlan::from_raw_batch_with_cta_tile_q( + ctx, + &[vec![0i32]], + &[seq_len], + &[0usize], + &[seq_len], + config.num_attention_heads, + config.num_key_value_heads, + config.head_dim, + 0, // use default CTA tile size + ) + .expect("failed to create prefill plan"); + + eprintln!("running backbone forward..."); + let hidden = backbone + .forward(&input_ids, &kv_buffer, &layout, &plan) + .expect("backbone forward failed"); + + assert_eq!(hidden.hidden_dim, config.hidden_size); + assert_eq!(hidden.seq_len, seq_len); + + // Compute last-token logits + let logits = backbone + .last_token_logits(&hidden) + .expect("last_token_logits failed"); + + assert_eq!(logits.hidden_dim, config.vocab_size); + assert_eq!(logits.seq_len, 1); + + // Copy logits to host and convert bf16 → f32 + let n_logits = logits.data.len(); + let mut logits_bf16 = vec![bf16::ZERO; n_logits]; + ctx.stream + .memcpy_dtoh(&logits.data, &mut logits_bf16) + .expect("failed to copy logits to host"); + ctx.stream.synchronize().expect("sync failed"); + let logits_host: Vec = logits_bf16.iter().map(|&x| f32::from(x)).collect(); + + // Compute log_softmax for both our logits and golden logits + let our_logprobs = log_softmax(&logits_host); + let golden_logprobs = log_softmax(&golden_logits); + + assert_eq!( + our_logprobs.len(), + golden_logprobs.len(), + "vocab size mismatch: ours={}, golden={}", + our_logprobs.len(), + golden_logprobs.len() + ); + + eprintln!("comparing logprobs (vocab={})...", our_logprobs.len()); + check_logprobs(&our_logprobs, &golden_logprobs); +} + +/// Compute log_softmax in f32. +fn log_softmax(logits: &[f32]) -> Vec { + let max_val = logits.iter().fold(f32::NEG_INFINITY, |a, &b| a.max(b)); + let sum_exp: f32 = logits.iter().map(|&x| (x - max_val).exp()).sum(); + let log_sum = sum_exp.ln(); + logits.iter().map(|&x| x - max_val - log_sum).collect() +}