Skip to content

vllm_omni.model_executor.models.audex.audex_omni

Audex audio-understanding model (checkpoint_folder_full): S2T / audio QA.

Mirrors the Qwen2-Audio pattern with the 2B dense backbone
  • language_model = the vllm-omni NemotronDenseForCausalLM port (wrapped via init_vllm_registered_model; resolved through the omni model registry).
  • audio_tower = transformers Qwen2AudioEncoder ("NV-Whisper"); a parameter-free avg pool maps 1500 → 750 frames per padded 30 s clip.
  • audio_projector = RMSNorm → fc1 → relu² → fc2 (no bias), matching the release NemotronHAudexProjector.

embed_multimodal(...) runs encoder+projector and returns one embedding tensor per audio item; vLLM's SupportsMultiModal.embed_input_ids merges them into the <so_embedding> placeholder positions. The encoder runs in microbatches over clips so peak memory is bounded regardless of audio length (up to the 900 s / 30-clip cap enforced by the processor).

AudexAudioFeatureInputs

Bases: TensorSchema

Dimensions
  • c: total number of 30s clips (flattened across audio items)
  • b: number of audio items

audio_num_clips class-attribute instance-attribute 

audio_num_clips: Annotated[list[int], None] = None

input_features instance-attribute 

input_features: Annotated[
    Tensor, TensorShape("c", 128, 3000)
]

type class-attribute instance-attribute 

type: Literal['audio_features'] = 'audio_features'

AudexProjector

Bases: Module

Megatron sound_projection equivalent (matches HF NemotronHAudexProjector).

activation instance-attribute 

activation = activation

fc1 instance-attribute 

fc1 = nn.Linear(
    audio_encoder_hidden_size, intermediate_size, bias=False
)

fc2 instance-attribute 

fc2 = nn.Linear(
    intermediate_size, text_hidden_size, bias=False
)

norm_eps instance-attribute 

norm_eps = norm_eps

norm_weight instance-attribute 

norm_weight = nn.Parameter(
    torch.ones(audio_encoder_hidden_size)
)

forward 

forward(hidden_states: Tensor) -> Tensor

NemotronDenseAudexForConditionalGeneration

Bases: Module, SupportsMultiModal, SupportsPP

audio_projector instance-attribute 

audio_projector = AudexProjector(
    audio_encoder_hidden_size=config.audio_encoder_hidden_size,
    intermediate_size=config.audio_projector_intermediate_size,
    text_hidden_size=config.hidden_size,
    activation=config.audio_projector_activation,
    norm_eps=config.audio_projector_norm_eps,
).to(llm_dtype)

audio_tower instance-attribute 

audio_tower = _build_qwen2_audio_encoder(
    config.audio_config
).to(llm_dtype)

config instance-attribute 

config = config

language_model instance-attribute 

language_model = init_vllm_registered_model(
    vllm_config=vllm_config,
    hf_config=config,
    prefix=maybe_prefix(prefix, "language_model"),
    architectures=[self._LM_ARCHITECTURE],
)

llm_dtype instance-attribute 

llm_dtype = llm_dtype

make_empty_intermediate_tensors instance-attribute 

make_empty_intermediate_tensors = (
    self.language_model.make_empty_intermediate_tensors
)

compute_logits 

compute_logits(
    hidden_states: Tensor, sampling_metadata=None
) -> Tensor | None

embed_multimodal 

embed_multimodal(**kwargs: object) -> MultiModalEmbeddings

forward 

forward(
    input_ids: Tensor | None,
    positions: Tensor,
    intermediate_tensors: IntermediateTensors | None = None,
    inputs_embeds: Tensor | None = None,
    **kwargs: object,
) -> Tensor | IntermediateTensors

get_language_model 

get_language_model() -> Module

get_placeholder_str classmethod 

get_placeholder_str(modality: str, i: int) -> str | None

load_weights 

load_weights(
    weights: Iterable[tuple[str, Tensor]],
) -> set[str]

Split the combined weight stream into LLM / audio_tower / projector.

vLLM streams all .safetensors in the model dir; route audio_encoder.*/audio_projector.* to the audio modules and leave the remaining model.*/lm_head.* names untouched so the wrapped NemotronDenseForCausalLM loader applies its own qkv fusion and remapping.

NemotronHAudexForConditionalGeneration

Bases: NemotronDenseAudexForConditionalGeneration, IsHybrid, HasInnerState, SupportsMambaPrefixCaching

Audex 30B-A3B audio-understanding model (checkpoint_folder_full).

Mirrors nvidia/Nemotron-Labs-Audex-30B-A3B's inference_scripts_vllm/audioqa_scripts/audex_30b_a3b_vllm/modeling_audex_vllm.py: the same encoder/projector/processing as the dense 2B wrapper around vLLM's native hybrid Mamba+MoE NemotronHForCausalLM, whose state hooks are delegated below.

has_inner_state class-attribute instance-attribute 

has_inner_state = True

is_hybrid class-attribute instance-attribute 

is_hybrid = True

supports_mamba_prefix_caching class-attribute instance-attribute 

supports_mamba_prefix_caching = True

get_mamba_state_copy_func classmethod 

get_mamba_state_copy_func()

get_mamba_state_dtype_from_config classmethod 

get_mamba_state_dtype_from_config(vllm_config)

get_mamba_state_shape_from_config classmethod 

get_mamba_state_shape_from_config(vllm_config)