feat(higgs-audio): add checkpoint loading and validation foundation

Cargo.toml

@@ -15,6 +15,7 @@ members = [
   "pegainfer-kimi-k2",
   "pegainfer-qwen3",
   "pegainfer-qwen35",
+  "pegainfer-higgs-audio",
   "pegainfer-sample",
   "pegainfer-kv-cache",
   "pegainfer-kv-offload",
@@ -128,6 +129,7 @@ pegainfer-kv-offload = { path = "pegainfer-kv-offload" }
 pegainfer-kv-store = { path = "pegainfer-kv-store" }
 pegainfer-qwen3 = { path = "pegainfer-qwen3" }
 pegainfer-qwen35 = { path = "pegainfer-qwen35" }
+pegainfer-higgs-audio = { path = "pegainfer-higgs-audio" }
 pegainfer-sample = { path = "pegainfer-sample" }
 opentelemetry = { version = "0.31.0", features = ["trace", "logs"] }
 opentelemetry-appender-tracing = "0.31.0"

docs/models/higgs-audio/a-layer-validation.md

@@ -0,0 +1,92 @@
+# Higgs-Audio A-Layer Validation
+
+This note records the intended review scope for the Higgs-Audio A-layer bring-up.
+It is deliberately narrower than full audio generation: the goal is to load the
+Higgs checkpoint/config, run the Qwen3 text backbone prefill, expose traceable
+hidden-state boundaries, and verify the one-step audio-head contract against
+offline golden tensors.
+
+## Scope
+
+Included:
+
+- Higgs-Audio crate wiring and one-step CLI tools.
+- Qwen3 config-view loading backed by tensor-name aliases, without copying a
+  renamed checkpoint payload.
+- Single-GPU Qwen3 diagnostic prefill hooks for final hidden, per-layer hidden,
+  and selected layer stage dumps.
+- One-step audio projection parity tools and a small offline fixture.
+
+Not included:
+
+- Full delay-pattern decode state machine.
+- Multi-codebook autoregressive audio decode.
+- Codec/vocoder integration or waveform output.
+- Tensor-parallel diagnostic trace support.
+
+## Validation Model
+
+The A-layer validation is golden-trace-driven:
+
+1. Generate or load a reference prompt and hidden/logit tensors from the Python
+   or sglang-omni reference stack.
+2. Run the PegaInfer Higgs one-step path against the same checkpoint/config.
+3. Compare semantic outputs first (`audio_argmax`, cosine similarity, top-k
+   overlap), then use layer and stage dumps to localize any drift.
+
+Strict elementwise parity is useful as a diagnostic, but it is not the only
+acceptance signal. The current 4090 evidence showed the semantic gate passing
+while small BF16/F32 elementwise drift remained:
+
+The RMSNorm rounding ablation showed that the HF/Qwen3-style fused-add-RMSNorm
+variant improves strict trace parity slightly, but is not required for the
+current one-step semantic gate. To keep this foundation PR focused, the shared
+`pegainfer-kernels` rounding change is left out of scope and can be discussed
+separately as a Qwen3 numeric-parity change if needed.
+
+Evidence source: old-4090 semantic comparison logs from the Higgs-Audio
+trace-driven bring-up run; the auto path and retained-session path reported the
+same semantic metrics.
+
+- `audio_argmax.ids`: exact.
+- `hidden_cosine`: `0.999994874`.
+- `logits_cosine`: `0.999998987`.
+- `top64_min_overlap`: `58`.
+- `top64_mean_overlap`: `61`.
+- `final_hidden.bf16 mean_abs`: about `0.006735`.
+- `audio_logits.f32 mean_abs`: about `0.040235`.
+
+## Local Checks
+
+Checks that do not require a Linux CUDA runtime:
+
+```bash
+cargo fmt --check -p pegainfer-core -p pegainfer-qwen3 -p pegainfer-higgs-audio
+cargo check -p pegainfer-higgs-audio --bins
+python3 -m py_compile \
+  tools/accuracy/analyze_higgs_one_step_actual.py \
+  tools/accuracy/analyze_higgs_projection_drift.py \
+  tools/accuracy/analyze_higgs_qk_norm_drift.py \
+  tools/accuracy/analyze_higgs_residual_drift.py \
+  tools/accuracy/analyze_higgs_rmsnorm_drift.py \
+  tools/accuracy/compare_higgs_layer_hidden.py \
+  tools/accuracy/compare_higgs_stage_dump.py \
+  tools/accuracy/compare_higgs_trace_dump.py \
+  tools/accuracy/dump_higgs_layer0_stages_golden.py \
+  tools/accuracy/dump_higgs_layer_hidden_golden.py \
+  tools/accuracy/dump_higgs_one_step_golden.py \
+  tools/higgs/check_higgs_gate_summary.py
+```
+
+Linux/4090 checks:
+
+```bash
+export PEGAINFER_CUDA_SM=89
+cargo check -p pegainfer-higgs-audio --features runtime-qwen3 --bins
+tools/higgs/run_higgs_one_step_cuda_gate.sh <higgs-model-dir> <golden.safetensors> <out-dir>
+```
+
+On macOS, the runtime-Qwen3 check is expected to stop before useful Rust type
+checking because the workspace currently builds CUDA kernels and Linux RDMA
+dependencies (`rdma-mummy-sys` expects Linux headers such as `endian.h` and
+`linux/types.h`).

pegainfer-core/src/weight_loader.rs

@@ -24,6 +24,35 @@ use crate::tensor::DeviceContext;
 use crate::tensor::DeviceMatrix;
 use crate::tensor::DeviceVec;
 
+/// Optional mapping from the runtime's requested tensor name to the tensor name
+/// stored in the safetensors shard.
+///
+/// Normal model loading keeps the identity mapping. Model adapters can pass an
+/// alias table to reuse a checkpoint layout without materializing a renamed
+/// weight copy.
+#[derive(Clone, Debug, Default)]
+pub struct TensorNameAliases {
+    storage_by_requested: HashMap<String, String>,
+}
+
+impl TensorNameAliases {
+    pub fn new(storage_by_requested: HashMap<String, String>) -> Self {
+        Self {
+            storage_by_requested,
+        }
+    }
+
+    pub fn is_empty(&self) -> bool {
+        self.storage_by_requested.is_empty()
+    }
+
+    fn storage_name<'a>(&'a self, requested_name: &'a str) -> &'a str {
+        self.storage_by_requested
+            .get(requested_name)
+            .map_or(requested_name, String::as_str)
+    }
+}
+
 mod staging;
 use staging::ColShardPlan;
 use staging::WeightStager;
@@ -253,18 +282,30 @@ fn find_tensor<'a>(
     weight_map: &HashMap<String, usize>,
     name: &str,
 ) -> Result<safetensors::tensor::TensorView<'a>> {
-    if let Some(&idx) = weight_map.get(name) {
-        shards[idx]
-            .tensor(name)
-            .map_err(|e| anyhow::anyhow!("Failed to load tensor '{}': {}", name, e))
+    find_tensor_with_aliases(shards, weight_map, &TensorNameAliases::default(), name)
+}
+
+fn find_tensor_with_aliases<'a>(
+    shards: &'a [SafeTensors<'a>],
+    weight_map: &HashMap<String, usize>,
+    aliases: &TensorNameAliases,
+    name: &str,
+) -> Result<safetensors::tensor::TensorView<'a>> {
+    let storage_name = aliases.storage_name(name);
+    if let Some(&idx) = weight_map.get(storage_name) {
+        shards[idx].tensor(storage_name).map_err(|e| {
+            anyhow::anyhow!("Failed to load tensor '{name}' stored as '{storage_name}': {e}")
+        })
     } else {
         // Fallback: try all shards (single-file case)
         for shard in shards {
-            if let Ok(t) = shard.tensor(name) {
+            if let Ok(t) = shard.tensor(storage_name) {
                 return Ok(t);
             }
         }
-        Err(anyhow::anyhow!("Tensor '{}' not found in any shard", name))
+        Err(anyhow::anyhow!(
+            "Tensor '{name}' stored as '{storage_name}' not found in any shard"
+        ))
     }
 }
 
@@ -361,6 +402,7 @@ pub struct StagedWeightLoader<'a> {
     stager: WeightStager,
     shards: &'a [SafeTensors<'a>],
     weight_map: &'a HashMap<String, usize>,
+    aliases: TensorNameAliases,
     slots: Vec<Slot>,
     vec_slots: Vec<VecSlot>,
     pending: Vec<PendingUpload<'a>>,
@@ -381,6 +423,7 @@ impl<'a> StagedWeightLoader<'a> {
             stager: WeightStager::new(ctx)?,
             shards,
             weight_map,
+            aliases: TensorNameAliases::default(),
             slots: Vec::new(),
             vec_slots: Vec::new(),
             pending: Vec::new(),
@@ -391,6 +434,12 @@ impl<'a> StagedWeightLoader<'a> {
         })
     }
 
+    #[must_use]
+    pub fn with_aliases(mut self, aliases: TensorNameAliases) -> Self {
+        self.aliases = aliases;
+        self
+    }
+
     fn ensure_recording(&self) -> Result<()> {
         anyhow::ensure!(
             !self.finished && !self.failed,
@@ -482,7 +531,7 @@ impl<'a> StagedWeightLoader<'a> {
     }
 
     fn tensor_2d(&self, name: &str, rows: usize, cols: usize) -> Result<&'a [u8]> {
-        let tensor = find_tensor(self.shards, self.weight_map, name)?;
+        let tensor = find_tensor_with_aliases(self.shards, self.weight_map, &self.aliases, name)?;
         let shape = tensor.shape();
         anyhow::ensure!(
             shape.len() == 2,
@@ -584,7 +633,7 @@ impl<'a> StagedWeightLoader<'a> {
     /// Small tensors; uploaded as plain pageable copies.
     pub fn vector(&mut self, name: &str, len: usize) -> Result<VecSlotId> {
         self.ensure_recording()?;
-        let tensor = find_tensor(self.shards, self.weight_map, name)?;
+        let tensor = find_tensor_with_aliases(self.shards, self.weight_map, &self.aliases, name)?;
         let shape = tensor.shape();
         anyhow::ensure!(
             shape.len() == 1 && shape[0] == len,

pegainfer-higgs-audio/Cargo.toml

@@ -0,0 +1,45 @@
+[package]
+name = "pegainfer-higgs-audio"
+version = "0.1.0"
+edition.workspace = true
+license.workspace = true
+
+[dependencies]
+anyhow = { workspace = true }
+clap = { workspace = true }
+half = { workspace = true }
+memmap2 = { workspace = true }
+pegainfer-core = { workspace = true, optional = true }
+pegainfer-qwen3 = { workspace = true, optional = true }
+safetensors = { workspace = true }
+serde_json = { workspace = true }
+sha2 = { workspace = true }
+
+[features]
+runtime-qwen3 = ["dep:pegainfer-core", "dep:pegainfer-qwen3"]
+
+[dev-dependencies]
+tempfile = { workspace = true }
+
+[[bin]]
+name = "higgs_dump_one_step_actual"
+path = "src/bin/higgs_dump_one_step_actual.rs"
+required-features = ["runtime-qwen3"]
+
+[[bin]]
+name = "higgs_prefill_prompt_session_smoke"
+path = "src/bin/higgs_prefill_prompt_session_smoke.rs"
+required-features = ["runtime-qwen3"]
+
+[[bin]]
+name = "higgs_dump_prefill_layer_hidden"
+path = "src/bin/higgs_dump_prefill_layer_hidden.rs"
+required-features = ["runtime-qwen3"]
+
+[[bin]]
+name = "higgs_dump_layer0_stages"
+path = "src/bin/higgs_dump_layer0_stages.rs"
+required-features = ["runtime-qwen3"]
+
+[lints]
+workspace = true