Skip to content

vllm_omni.model_executor.models.audex.audex_thinker

Audex thinker: the Nemotron dense decoder-only LM that autoregressively emits <speechcodec_N> tokens for TTS.

The thinker is a plain token-in/token-out vLLM model. Speech-codec token extraction happens in the stage input processors (vllm_omni/model_executor/stage_input_processors/audex.py), which read the sampled token stream directly, so no omni-specific output plumbing is needed here.

NemotronDenseAttention

Bases: Module

attn instance-attribute 

attn = Attention(
    self.num_heads,
    self.head_dim,
    self.scaling,
    num_kv_heads=self.num_kv_heads,
    cache_config=cache_config,
    quant_config=quant_config,
    prefix=f"{prefix}.attn",
)

head_dim instance-attribute 

head_dim = getattr(config, 'head_dim', None)

hidden_size instance-attribute 

hidden_size = hidden_size

kv_size instance-attribute 

kv_size = self.num_kv_heads * self.head_dim

max_position_embeddings instance-attribute 

max_position_embeddings = max_position_embeddings

num_heads instance-attribute 

num_heads = self.total_num_heads // tp_size

num_kv_heads instance-attribute 

num_kv_heads = max(1, self.total_num_kv_heads // tp_size)

o_proj instance-attribute 

o_proj = RowParallelLinear(
    input_size=self.total_num_heads * self.head_dim,
    output_size=hidden_size,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.o_proj",
)

q_size instance-attribute 

q_size = self.num_heads * self.head_dim

qkv_proj instance-attribute 

qkv_proj = QKVParallelLinear(
    hidden_size=hidden_size,
    head_size=self.head_dim,
    total_num_heads=self.total_num_heads,
    total_num_kv_heads=self.total_num_kv_heads,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv_proj",
)

rotary_emb instance-attribute 

rotary_emb = get_rope(
    self.head_dim,
    max_position=max_position_embeddings,
    rope_parameters=config.rope_parameters,
)

scaling instance-attribute 

scaling = self.head_dim ** -0.5

total_num_heads instance-attribute 

total_num_heads = num_heads

total_num_kv_heads instance-attribute 

total_num_kv_heads = num_kv_heads

forward 

forward(positions: Tensor, hidden_states: Tensor) -> Tensor

NemotronDenseDecoderLayer

Bases: Module

hidden_size instance-attribute 

hidden_size = config.hidden_size

input_layernorm instance-attribute 

input_layernorm = NemotronDenseRMSNorm(
    config.hidden_size, eps=config.norm_eps
)

mlp instance-attribute 

mlp = NemotronDenseMLP(
    hidden_size=self.hidden_size,
    intermediate_size=config.intermediate_size,
    hidden_act=config.hidden_act,
    quant_config=quant_config,
    bias=getattr(config, "mlp_bias", False),
    prefix=f"{prefix}.mlp",
)

post_attention_layernorm instance-attribute 

post_attention_layernorm = NemotronDenseRMSNorm(
    config.hidden_size, eps=config.norm_eps
)

self_attn instance-attribute 

self_attn = NemotronDenseAttention(
    config=config,
    hidden_size=self.hidden_size,
    num_heads=config.num_attention_heads,
    num_kv_heads=getattr(
        config,
        "num_key_value_heads",
        config.num_attention_heads,
    ),
    max_position_embeddings=max_position_embeddings,
    quant_config=quant_config,
    bias=attention_bias,
    cache_config=cache_config,
    prefix=f"{prefix}.self_attn",
)

forward 

forward(positions: Tensor, hidden_states: Tensor) -> Tensor

NemotronDenseForCausalLM

Bases: Module, SupportsLoRA, SupportsPP

config instance-attribute 

config = config

embedding_modules class-attribute instance-attribute 

embedding_modules = {
    "embed_tokens": "input_embeddings",
    "lm_head": "output_embeddings",
}

lm_head instance-attribute 

lm_head = ParallelLMHead(
    config.vocab_size,
    config.hidden_size,
    quant_config=quant_config,
    prefix=maybe_prefix(prefix, "lm_head"),
)

logits_processor instance-attribute 

logits_processor = LogitsProcessor(
    config.vocab_size, scale=logit_scale
)

make_empty_intermediate_tensors instance-attribute 

make_empty_intermediate_tensors = (
    self.model.make_empty_intermediate_tensors
)

model instance-attribute 

model = NemotronDenseModel(
    vllm_config=vllm_config,
    prefix=maybe_prefix(prefix, "model"),
)

packed_modules_mapping class-attribute instance-attribute 

packed_modules_mapping = {
    "qkv_proj": ["q_proj", "k_proj", "v_proj"]
}

quant_config instance-attribute 

quant_config = quant_config

compute_logits 

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

embed_input_ids 

embed_input_ids(input_ids: Tensor) -> Tensor

forward 

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

load_weights 

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

NemotronDenseMLP

Bases: Module

act_fn instance-attribute 

act_fn = get_act_fn(hidden_act)

down_proj instance-attribute 

down_proj = RowParallelLinear(
    input_size=intermediate_size,
    output_size=hidden_size,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.down_proj",
)

up_proj instance-attribute 

up_proj = ColumnParallelLinear(
    input_size=hidden_size,
    output_size=intermediate_size,
    bias=bias,
    quant_config=quant_config,
    prefix=f"{prefix}.up_proj",
)

forward 

forward(x: Tensor) -> Tensor

NemotronDenseModel

Bases: Module

config instance-attribute 

config = config

embed_tokens instance-attribute 

embed_tokens = VocabParallelEmbedding(
    self.vocab_size, config.hidden_size
)

make_empty_intermediate_tensors instance-attribute 

make_empty_intermediate_tensors = (
    make_empty_intermediate_tensors_factory(
        ["hidden_states"], config.hidden_size
    )
)

norm instance-attribute 

norm = NemotronDenseRMSNorm(
    config.hidden_size, eps=config.norm_eps
)

quant_config instance-attribute 

quant_config = quant_config

vocab_size instance-attribute 

vocab_size = config.vocab_size

embed_input_ids 

embed_input_ids(input_ids: Tensor) -> Tensor

forward 

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

load_weights 

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

NemotronDenseRMSNorm

Bases: Module

variance_epsilon instance-attribute 

variance_epsilon = eps

weight instance-attribute 

weight = nn.Parameter(torch.ones(hidden_size))

forward 

forward(hidden_states: Tensor) -> Tensor