A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 9, 21:48 UTC
[bugfix] Aria model (#32727)
Signed-off-by: Divakar Verma <divakar.verma@amd.com>
Divakar Verma committed
Jan 21, 2026 at 07:11 UTC
e14467be432bf55493c21ffd9f7d1d4c32e14e19
1 file changed
+16
-1
vllm/model_executor/models/aria.py
+16
-1
@@ -15,7 +15,9 @@ from vllm.distributed import get_tensor_model_parallel_rank
15
from vllm.model_executor.layers.activation import get_act_fn
16
from vllm.model_executor.layers.fused_moe import SharedFusedMoE
17
from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear
18
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
19
from vllm.model_executor.layers.quantization import QuantizationConfig
20
+from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead
21
from vllm.model_executor.model_loader.weight_utils import (
22
default_weight_loader,
23
maybe_remap_kv_scale_name,
@@ -554,6 +556,18 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
556
prefix=maybe_prefix(prefix, "language_model.model"),
557
)
558
559
+ self.lm_head = ParallelLMHead(
560
+ config.text_config.vocab_size,
561
+ config.text_config.hidden_size,
562
+ quant_config=quant_config,
563
+ prefix=maybe_prefix(prefix, "lm_head"),
564
+ )
565
+
566
+ logit_scale = getattr(config, "logit_scale", 1.0)
567
+ self.logits_processor = LogitsProcessor(
568
+ config.text_config.vocab_size, scale=logit_scale
569
+ )
570
+
571
def _parse_and_validate_image_input(
572
self, **kwargs: object
573
) -> AriaImagePixelInputs | None:
@@ -637,7 +651,8 @@ class AriaForConditionalGeneration(nn.Module, SupportsMultiModal):
651
self,
652
hidden_states: torch.Tensor,
653
) -> torch.Tensor | None:
640
- return self.language_model.compute_logits(hidden_states)
654
+ logits = self.logits_processor(self.lm_head, hidden_states)
655
+ return logits
656
657
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]):
658
loader = AutoWeightsLoader(self)