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

[ywh555hhh]

Summary

This PR adds the first Higgs-Audio integration foundation for PegaInfer.

The scope is intentionally limited to the checkpoint-backed validation path: parse Higgs-Audio config and checkpoint metadata, map Higgs checkpoint tensor names into the existing Qwen3 loader, run Qwen3 text-backbone prefill, run the audio-head projection validation path, and add golden-trace tools for debugging hidden-state drift.

This is not full Higgs-Audio generation yet. Delay-pattern decode, multi-codebook autoregressive audio decode, codec/vocoder integration, and waveform output are left for follow-up PRs.

Why this shape

Higgs-Audio reuses a Qwen3-style text backbone, but its checkpoint layout is not a drop-in Qwen3 checkpoint. To avoid materializing a renamed multi-GB checkpoint copy, this PR adds a tensor-name alias path in the weight loader. The Qwen3 config view stays small and explicit, while the original Higgs checkpoint remains the weight source.

I have been calling this the first integration layer during bring-up: it proves that checkpoint/config loading, Qwen3 backbone prefill, audio-head projection, and trace-based validation can work before expanding into full audio generation.

The Qwen3 diagnostic prefill hooks added here are for model bring-up and golden-trace validation. They are kept separate from the normal serving scheduler path.

What changed

• Added pegainfer-higgs-audio with Higgs config, manifest, load-plan, and validation utilities.
• Added Qwen3 weight-source loading with tensor-name aliases.
• Added Qwen3 prefill diagnostics for final hidden, per-layer hidden, and selected layer stages.
• Added comparison tools and a small offline fixture for the audio-head validation path.
• Documented scope, validation method, and current limitations in docs/models/higgs-audio/a-layer-validation.md.

Validation

Local checks run:

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/*.py tools/higgs/check_higgs_gate_summary.py

4090 validation evidence from the checkpoint-backed validation path:

Metric	Result
prompt ids	exact
audio argmax ids	exact
hidden cosine	0.999994874
logits cosine	0.999998987
top64 min overlap	58 / 64
top64 mean overlap	61.00 / 64
final hidden bf16 mean abs	0.006735
audio logits f32 mean abs	0.040235

Strict elementwise parity is still diagnostic-only, but the semantic gate passes.

I also ran an RMSNorm rounding ablation. The HF/Qwen3-style fused-add-RMSNorm variant slightly improves strict trace parity, but it is not required for the current Higgs-Audio one-step semantic gate. To reduce shared-layer scope, this PR no longer changes pegainfer-kernels; the rounding question can be discussed separately as a Qwen3 numeric-parity topic if needed.

Known limitations

• Full audio generation is not included in this PR.
• Diagnostic prefill hooks currently target the single-GPU bring-up path.
• Runtime CUDA validation needs a Linux CUDA machine; macOS local checks stop at CUDA/RDMA build dependencies.

Maintainer question

Is the tensor-name alias loader an acceptable way to reuse the original Higgs checkpoint layout, instead of materializing a renamed Qwen3-body checkpoint copy?

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 547bab9dca

ℹ️ About Codex in GitHub

Codex has been enabled to automatically review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

When you sign up for Codex through ChatGPT, Codex can also answer questions or update the PR, like "@codex address that feedback".

[chatgpt-codex-connector]

Move the single-prompt plan helper onto Qwen3Model

Because prefill_plan_for_single_prompt was added inside impl PrefillBuffers, self is a scratch-buffer object with only HiddenStates fields; it has no ctx, config, or local_* members. Any build of the default pegainfer-qwen3 crate will type-check this body and fail before the Higgs runtime can be built, so this helper needs to live on Qwen3Model or receive the needed context explicitly.

Useful? React with 👍 / 👎.

[ywh555hhh]

Small self-review note before detailed maintainer review:

I realized this PR is broader than a pure pegainfer-higgs-audio addition. It adds the new Higgs-Audio validation crate/tools/docs, but it also touches shared/Qwen3-facing seams:

• pegainfer-core: optional tensor-name aliases in the weight loader.
• pegainfer-qwen3: diagnostic prefill surfaces used for golden-trace validation.
• pegainfer-kernels: an HF/Qwen3-style fused add RMSNorm rounding variant.

The intent was to reuse the existing Qwen3 backbone and avoid duplicating loading/forward logic inside the Higgs-Audio crate. That said, I recognize these are public/shared-layer changes and may be too much for one PR.

I’m going to double-check whether each shared change is truly necessary and whether this should be split into smaller PRs before asking for full review. If maintainers prefer, I can split it into: Higgs crate/tools/docs, weight-loader alias support, Qwen3 diagnostics, and RMSNorm rounding as separate PRs.

CI is currently passing, but I do not want to hide the scope concern.

[ywh555hhh]

I want to discuss one design tradeoff in this PR before asking for deeper review.

This PR no longer strictly matches the original guardrail in #408, where I said the Higgs-Audio bring-up would avoid shared crate modifications. During implementation, I found that validating against the real Qwen3-backed runtime path creates a tradeoff:

• keep pegainfer-higgs-audio isolated, but duplicate Qwen3-like loading/prefill logic and risk drifting from the main Qwen3 implementation;
• or reuse the existing Qwen3 backbone path, but add small shared/Qwen3-facing seams such as tensor-name aliases and diagnostic prefill hooks.

I chose the second path in this PR, but I realize this is exactly the kind of shared-layer complexity that maintainers may not want bundled into an initial Higgs-Audio PR.

This also relates to the earlier guidance in #203: Higgs should be one model / one crate, and audio-specific concepts should not leak into shared layers. I think the current question is whether these changes are acceptable as general reusable seams, or whether I should split them out / keep the initial Higgs PR more isolated.

Concretely, I would like maintainer input on the first tradeoff:

Should Higgs-Audio reuse the existing Qwen3 loading/prefill path via small shared seams, or should the initial implementation duplicate/adapt Qwen3-like logic inside pegainfer-higgs-audio to keep other model lines untouched?

I can split this PR if the current shape is too broad.

[ywh555hhh]

Follow-up after the RMSNorm rounding ablation:

I removed the pegainfer-kernels shared rounding change from this PR. The ablation supports that HF/Qwen3-style fused-add-RMSNorm rounding slightly improves strict trace parity, but is not required for the current Higgs-Audio one-step semantic gate.

So the foundation PR is now narrower: it keeps the Higgs crate/tools/docs, tensor-name alias support, and Qwen3 diagnostic prefill surface, but no longer changes shared kernel code. If the rounding behavior is worth pursuing later, I will treat it as a separate Qwen3 numeric-parity discussion rather than coupling it to this Higgs-Audio foundation PR.