inference_scripts_vllm/audioqa_scripts/README.md

Audex-2B audio understanding in vLLM 0.20 (Nemotron Nano Omni-style)

Production-style integration: audio -> text through vLLM's own multimodal pipeline. A registered model wraps the vLLM-native NemotronDense (2B dense) LLM and adds the NV-Whisper audio encoder + Audex projector; vLLM owns placeholder merging and serving (offline LLM.generate and the OpenAI audio_url server).

This is the 2B-dense counterpart to Audex-30B-A3B/.../audioqa_scripts. The only differences from the 30B package: the backbone is the plain attention NemotronDenseForCausalLM (no Mamba/hybrid cache) and the LLM hidden size is 2048 (the Audex projector output matches config.hidden_size). Audio preprocessing, 30s chunking, <so_embedding> expansion, caps, and the vLLM processor logic are identical.

Architecture name (from checkpoint_folder_full/config.json): NemotronDenseAudexForConditionalGeneration (model_type=nemotron_dense_audex).

Layout

inference_scripts_vllm/audioqa_scripts/
  README.md
  pyproject.toml                       # installs the model as a vLLM plugin
  run_audioqa_vllm.py                  # offline LLM.generate runner
  serve_audioqa_vllm.sh                # OpenAI-compatible server
  client_audioqa.py                    # OpenAI audio_url client (one request)
  audex_2b_vllm/                       # import package (distinct from 30B audex_30b_a3b_vllm)
    modeling_audex_vllm.py   # NemotronDenseAudex model (SupportsMultiModal + SupportsPP)
    processing_audex_vllm.py # processor: 30s chunking, <so_embedding> expansion, caps
    audio_features.py        # waveform -> NV-Whisper features
    audio_encoder.py         # Audex projector + Qwen2AudioEncoder factory
    plugin.py / register.py  # registration in every TP worker

The import package is audex_2b_vllm (the 30B package is audex_30b_a3b_vllm), so both audioqa plugins can be co-installed without one shadowing the other.

Install (once)

The model must be registered in every tensor-parallel worker, which vLLM does via the vllm.general_plugins entry points. Install the dense backbone plugin and this folder as editable plugins:

pip install -e ../../nemotron_dense_vllm_plugin --no-deps --no-build-isolation
pip install -e . --no-deps --no-build-isolation

Run offline

python run_audioqa_vllm.py \
  --model-path "$(cd ../.. && pwd)/checkpoint_folder_full" \
  --input-json ./inputs.json \
  --output-jsonl ./audioqa_outputs/results.jsonl \
  --tensor-parallel-size 1

inputs.json: [{"id", "sound": "/abs/path.wav", "conversations":[{"from":"human","value":"<sound>\nDescribe this audio."}]}].

Serve + query (OpenAI audio_url)

# Safe defaults: HOST=127.0.0.1, local audio restricted to the Audex-2B release root.
bash serve_audioqa_vllm.sh "$(cd ../.. && pwd)/checkpoint_folder_full" 8000
python client_audioqa.py --audio /path/to/audio.wav --prompt "Describe this audio."

To expose externally or widen file access (advanced):

HOST=0.0.0.0 ALLOWED_MEDIA_PATH=/data bash serve_audioqa_vllm.sh ... 8000
# ALLOWED_MEDIA_PATH= (empty) disables local-file audio entirely.

Benchmark recipes

Both run_audioqa_vllm.py and client_audioqa.py take --recipe, which sets reproduction-safe sampling. No audio benchmark uses thinking mode; default is audio-understanding (non-thinking).

recipe	thinking	temperature	top_p	top_k
audio-understanding (default)	off	0.7	0.9	0
speech-recognition-translation	off	0.0	1.0	0
custom	on	0.7	0.9	0

• top_k=0 means "disabled" in vLLM (consider all tokens).
• Audio understanding is non-thinking: audio-understanding (temperature=0.7, top_p=0.9, top_k=0) covers audio understanding/reasoning.
• Greedy translation: speech-recognition-translation is true greedy for ASR/AST — temperature=0.0 triggers vLLM's greedy path (it normalizes top_p=1.0, top_k=0).
• Precedence: recipe defaults < explicit CLI override (--reasoning/--no-reasoning, --temperature, --top-p, --top-k). custom is the manual escape hatch (thinking on by default; override as needed).

Key facts / gotchas

• Backbone: 2B dense (NemotronDenseForCausalLM) — plain RMSNorm / relu^2 / GQA, no Mamba; the wrapper is SupportsMultiModal + SupportsPP (not hybrid).
• Hidden size: 2048. The Audex projector's fc2 maps the 1280-d encoder features to config.hidden_size (2048), not 2688.
• Long audio: non-overlapping 30s windows, padded tail, N = num_clips*750 placeholders. Caps (fail loud): MAX_AUDIO_SECONDS=900, MAX_AUDIO_CLIPS=30, MAX_AUDIO_TOKENS=22500 (in processing_audex_vllm.py). The offline runner and serve script default to --max-model-len 32768 so the full 22500-embedding cap fits in the context window.
• Placeholder contract: <so_embedding> -> <so_start> + N*<so_embedding>
  • <so_end>; a placeholder/token count mismatch fails loud.
• No audio-token leakage (offline and served): generation is masked to text ids — allowed_token_ids = range(131072) minus the sound placeholder ids (<so_embedding>/<so_start>/<so_end> = 29/30/31). The offline runner passes this to SamplingParams; the client passes it via extra_body. All audio codec/gen tokens are id >= 131072. Both paths also scan the output text for <audiocodec_/<speechcodec_/<audiogen_/<speechgen_/<so_*> leakage.
• Self-contained preflight: the offline runner and serve script fail early with a clear message if model.safetensors.index.json references shards that are missing/unresolvable.
• Reasoning / prompt format: the audio-understanding recipe is non-thinking; the offline runner and client default to it. These scripts use the Audex audio-understanding evaluation prompt format — the non-thinking generation prompt uses the <think></think> assistant prefix, consistently offline (run_audioqa_vllm.py) and served (checkpoint_folder_full/chat_template.jinja, enable_thinking=False). Use --recipe custom for a thinking-capable manual setup.

Integration notes (why the model code looks the way it does)

1. Out-of-tree registration must reach TP workers -> done via the plugin entry points (registering only in the main process raises "unsupported arch" in workers).
2. vLLM streams all .safetensors in the model dir; load_weights splits the stream: model.*/lm_head.* -> dense language model (its AutoWeightsLoader fuses q/k/v into qkv_proj), audio_encoder.*/audio_projector.* -> audio.
3. The dense backbone is registered by the separate nemotron-dense-vllm plugin; register_audex also registers it defensively so the wrapped language_model arch always resolves.