diff --git a/.buildkite/test_areas/models_basic.yaml b/.buildkite/test_areas/models_basic.yaml index 121cd3fa87f..39d8c7f9baf 100644 --- a/.buildkite/test_areas/models_basic.yaml +++ b/.buildkite/test_areas/models_basic.yaml @@ -78,8 +78,9 @@ steps: - tests/models/test_deepseek_v41_replay_start.py - tests/models/test_deepseek_v4_vl_stream_weights.py - tests/models/test_muse_spark_rocm_loading.py + - tests/models/test_muse_spark_logits_local.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_pcp.py models/test_deepseek_v41_pipeline.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_rocm_loading.py models/test_muse_spark_logits_local.py mirror: amd: label: ":amd: (MI355) Basic Models (Other)" @@ -105,8 +106,10 @@ steps: - tests/models/test_muse_spark_rocm_loading.py - tests/models/test_deepseek_v4_vl_stream_weights.py - tests/models/hy_v4/test_triton_ihc.py + - vllm/models/muse_spark/ + - tests/models/test_muse_spark_logits_local.py commands: - - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py + - pytest -v -s models/test_terratorch.py models/transformers/test_backend.py models/test_registry.py models/test_deepseek_v4_vl_rocm.py models/transformers/test_layer_registry_e2e.py models/test_deepseek_v41_replay_start.py models/test_deepseek_v41_decoder_replay_layers.py models/test_deepseek_v41_replay_batch.py models/test_deepseek_v41_amd_replay_forward.py models/test_deepseek_v4_dspark_rocm.py models/test_deepseek_v4_rocm_wo_a.py models/test_deepseek_v4_rocm_compressor_gemm_fusion.py models/test_dsv41_mono_decode.py models/test_muse_spark_rocm_loading.py models/test_deepseek_v4_vl_stream_weights.py models/test_muse_spark_logits_local.py - VLLM_ROCM_USE_AITER=1 pytest -v -s models/test_hyv4_rocm.py -m 'not distributed' - pytest -v -s models/hy_v4/test_triton_ihc.py diff --git a/tests/models/test_muse_spark_logits_local.py b/tests/models/test_muse_spark_logits_local.py new file mode 100644 index 00000000000..bb2ff5772e7 --- /dev/null +++ b/tests/models/test_muse_spark_logits_local.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import torch +from torch import nn + +from vllm.models.muse_spark.amd.model import MuseSpark1xForCausalLM + + +def test_compute_logits_local_applies_output_transform() -> None: + model = MuseSpark1xForCausalLM.__new__(MuseSpark1xForCausalLM) + nn.Module.__init__(model) + model.lm_head = nn.Linear(4, 8, bias=False, dtype=torch.bfloat16) + model.logits_processor = lambda lm_head, h, skip_gather=False: lm_head(h) + model.metap_output_multiplier = 0.5 + model.output_soft_cap_temp = 20.0 + model.mul_tanh_soft_cap = None + hidden = torch.randn(3, 4, dtype=torch.bfloat16) * 50 + + torch.testing.assert_close( + model.compute_logits_local(hidden), model.compute_logits(hidden) + ) diff --git a/vllm/models/muse_spark/amd/model.py b/vllm/models/muse_spark/amd/model.py index 4dce224cb67..d092e6068f7 100644 --- a/vllm/models/muse_spark/amd/model.py +++ b/vllm/models/muse_spark/amd/model.py @@ -1101,6 +1101,14 @@ class MuseSpark1xForCausalLM(LlamaForCausalLM, SupportsLoRA, MixtureOfExperts): logits = self.logits_processor(self.lm_head, hidden_states) if logits is None: return None + return self._apply_output_transform(logits) + + def compute_logits_local(self, hidden_states: torch.Tensor) -> torch.Tensor: + hidden_states = hidden_states.to(self.lm_head.weight.dtype) + logits = self.logits_processor(self.lm_head, hidden_states, skip_gather=True) + return self._apply_output_transform(logits) + + def _apply_output_transform(self, logits: torch.Tensor) -> torch.Tensor: if logits.dtype != torch.float32: logits = logits.to(torch.float32) if self.mul_tanh_soft_cap is not None: