[Bugfix][ROCm] Muse Spark: apply output multiplier and softcap in compute_logits_local (#62756)

Signed-off-by: Oxana Korzh <okorzh@amd.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
Co-authored-by: Fangzhou Ai <31551580+Fangzhou-Ai@users.noreply.github.com>

.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
 

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)
+    )

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: