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.

[xiaguan]

Thanks for pushing this checkpoint-loading path forward. The tensor-name alias idea and the checkpoint-backed one-step parity gate both look useful, and I think there is a good minimal foundation inside this PR.

Would you be open to reducing the scope before we review the implementation in detail? The current diff is about 10k lines across 48 files, and it combines several different concerns:

• the generic weight-loader alias seam;
• the Higgs config/manifest/name mapping and fused audio head;
• multiple Qwen3 runtime-source paths, including a materialized body view;
• public Qwen3 per-layer/per-stage diagnostic execution APIs;
• a fairly large shell/Python readiness and summary framework.

A smaller shape would be much easier to validate and maintain:

1. Land the generic tensor-alias loader seam separately, with one real safetensors load integration test.
2. Keep the Higgs PR focused on config/manifest validation, name mapping, the fused audio head, one narrow final-hidden hook, and one checkpoint-backed one-step golden gate.
3. Move or remove the materialized-body fallback, per-layer/stage tracing, retained prompt session support (until decode uses it), the handwritten kernel-plan metadata, and the isolated/readiness summary scaffolding.

In particular, run_isolated_higgs_checks.sh removes the real Qwen/core runtime dependencies and then validates fabricated artifact/summary files. That checks the helper parsers, but it does not increase confidence in the production checkpoint load or CUDA path. I would prefer the smaller set of real boundary tests over this additional framework.

One wording detail: the committed reference path uses a Transformers Qwen3 backbone plus the SGLang tokenizer/head semantics, not the full SGLang serving runtime, and the current Rust path stops after one prefill/audio-head prediction. Describing it as a one-step compatibility foundation would set the right expectation until delay-pattern decode, codec/vocoder, and serving are added.

[ywh555hhh]

Thanks for the detailed review. I understand the concern now.

You're right that this PR has grown too broad: it mixes the generic loader alias work, Higgs-specific checkpoint/config validation, Qwen3 diagnostics, and several readiness/debugging helpers. I'm happy to narrow the scope and make the follow-up work easier to review.

I'll convert this PR to draft for now, then split out the generic tensor-name alias loader into a smaller PR with a real safetensors loading test. After that, I'll reduce this PR to the minimal Higgs-Audio foundation: config/manifest validation, name mapping, fused audio head loading, a narrow final-hidden hook, and the checkpoint-backed one-step golden/parity gate.

Thanks again for the clear direction.