Skip to content

vllm_omni.model_executor.models.personaplex.modeling_helium

vLLM-native PersonaPlex Helium temporal transformer.

The module maps the Moshi temporal LM backbone onto vLLM's Llama-style decoder components. It intentionally does not include Moshi's text/audio input embeddings or depformer; PersonaPlex feeds temporal inputs_embeds directly.

HeliumAttention

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,
    per_layer_sliding_window=config.sliding_window,
    prefix=f"{prefix}.attn",
    attn_type=attn_type,
)

head_dim instance-attribute 

head_dim = config.head_dim

hidden_size instance-attribute 

hidden_size = hidden_size

kv_size instance-attribute 

kv_size = self.num_kv_heads * self.head_dim

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=False,
    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=False,
    quant_config=quant_config,
    prefix=f"{prefix}.qkv_proj",
)

rotary_emb instance-attribute 

rotary_emb = _get_helium_rope(config)

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

HeliumDecoderLayer

Bases: Module

hidden_size instance-attribute 

hidden_size = config.hidden_size

input_layernorm instance-attribute 

input_layernorm = HeliumRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps
)

mlp instance-attribute 

mlp = HeliumMLP(
    hidden_size=self.hidden_size,
    intermediate_size=config.intermediate_size,
    hidden_act=config.hidden_act,
    quant_config=quant_config,
    prefix=f"{prefix}.mlp",
)

post_attention_layernorm instance-attribute 

post_attention_layernorm = HeliumRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps
)

self_attn instance-attribute 

self_attn = HeliumAttention(
    config=config,
    hidden_size=self.hidden_size,
    num_heads=config.num_attention_heads,
    num_kv_heads=config.num_key_value_heads,
    cache_config=cache_config,
    quant_config=quant_config,
    prefix=f"{prefix}.self_attn",
    attn_type=AttentionType.DECODER,
)

forward 

forward(
    positions: Tensor,
    hidden_states: Tensor,
    residual: Tensor | None,
) -> tuple[Tensor, Tensor]

HeliumForCausalLM

Bases: Module, SupportsLoRA, SupportsPP

config instance-attribute 

config = config

embedding_modules class-attribute instance-attribute 

embedding_modules = {'lm_head': 'output_embeddings'}

lm_head instance-attribute 

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

logits_processor instance-attribute 

logits_processor = LogitsProcessor(config.vocab_size)

make_empty_intermediate_tensors instance-attribute 

make_empty_intermediate_tensors = (
    self.model.make_empty_intermediate_tensors
)

model instance-attribute 

model = HeliumModel(
    vllm_config=vllm_config,
    prefix=maybe_prefix(prefix, "model"),
    layer_type=layer_type,
)

packed_modules_mapping class-attribute instance-attribute 

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

quant_config instance-attribute 

quant_config = vllm_config.quant_config

compute_logits 

compute_logits(hidden_states: Tensor) -> Tensor | None

embed_input_ids 

embed_input_ids(input_ids: Tensor) -> Tensor

forward 

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

load_weights 

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

HeliumMLP

Bases: Module

act_fn instance-attribute 

act_fn = SiluAndMul()

down_proj instance-attribute 

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

gate_up_proj instance-attribute 

gate_up_proj = MergedColumnParallelLinear(
    input_size=hidden_size,
    output_sizes=[intermediate_size] * 2,
    bias=False,
    quant_config=quant_config,
    prefix=f"{prefix}.gate_up_proj",
)

forward 

forward(x: Tensor) -> Tensor

HeliumModel

Bases: Module

config instance-attribute 

config = config

make_empty_intermediate_tensors instance-attribute 

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

norm instance-attribute 

norm = HeliumRMSNorm(
    config.hidden_size, eps=config.rms_norm_eps
)

quant_config instance-attribute 

quant_config = vllm_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 = None,
    inputs_embeds: Tensor | None = None,
) -> Tensor | IntermediateTensors

HeliumRMSNorm

Bases: RMSNorm

Moshi-compatible RMSNorm.

Moshi stores the parameter as alpha with shape [1, 1, hidden] and computes variance and the alpha multiply in fp32 before casting back to the input dtype. vLLM's parameter is named weight; the loader squeezes alpha.

forward 

forward(
    x: Tensor, residual: Tensor | None = None
) -> Tensor | tuple[Tensor, Tensor]