Skip to content

Commit 2be876f

Browse files
authored
fix(glm_asr): honor sampling params in vLLM generate() (#2997)
GLMASRVLLMEngine.generate() accepted **kwargs but hardcoded temperature=0 and never forwarded any sampling parameter, so callers could not control decoding. Its sibling Fun-ASR-Nano engine already exposes temperature/top_p/top_k/repetition_penalty. Mirror that surface on GLM-ASR: forward temperature, top_p and top_k into SamplingParams (preserving the previous greedy default), and guard repetition_penalty so a non-neutral value is coerced back to 1.0 (with a one-time warning). This engine runs vLLM in prompt-embeds mode, where a non-neutral repetition_penalty triggers the issue #2948 CUDA crash. Add CPU-only regression tests covering parameter forwarding, the greedy defaults, top_k normalization, and the repetition_penalty guard.
1 parent d5ea9eb commit 2be876f

2 files changed

Lines changed: 163 additions & 2 deletions

File tree

‎funasr/models/glm_asr/inference_vllm.py‎

Lines changed: 44 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,6 +26,36 @@
2626
logger = logging.getLogger(__name__)
2727
dtype_map = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32}
2828

29+
# Warn only once per process so batch loops do not spam the log.
30+
_warned_rep_penalty = False
31+
32+
33+
def _safe_repetition_penalty(repetition_penalty):
34+
"""Force ``repetition_penalty`` to the neutral value for prompt-embeds mode.
35+
36+
GLM-ASR feeds vLLM precomputed embeddings (``enable_prompt_embeds=True``), so
37+
a request carries no prompt token IDs. vLLM applies ``repetition_penalty`` by
38+
scattering over those IDs, so any value other than 1.0 indexes an empty
39+
token-id tensor and aborts the engine with a CUDA
40+
``scatter gather index out of bounds`` assertion (issue #2948). We therefore
41+
warn once and fall back to the neutral value of 1.0.
42+
"""
43+
global _warned_rep_penalty
44+
45+
if repetition_penalty is None or repetition_penalty == 1.0:
46+
return 1.0
47+
48+
if not _warned_rep_penalty:
49+
logger.warning(
50+
"repetition_penalty=%s is not supported in vLLM prompt-embeds mode "
51+
"(no prompt token IDs to penalize) and would trigger a CUDA scatter "
52+
"index-out-of-bounds crash; using repetition_penalty=1.0 instead. "
53+
"See https://github.com/modelscope/FunASR/issues/2948.",
54+
repetition_penalty,
55+
)
56+
_warned_rep_penalty = True
57+
return 1.0
58+
2959

3060
def prepare_glmasr_vllm_dir(model_dir: str) -> str:
3161
"""Extract language_model weights into vLLM-compatible Llama format."""
@@ -182,13 +212,22 @@ def _build_prompt_embeds(self, audio_embeds, prompt="转录以下音频内容"):
182212
audio_emb = audio_embeds[0] if audio_embeds.dim() == 3 else audio_embeds
183213
return torch.cat([prefix_emb, audio_emb, suffix_emb], dim=0)
184214

185-
def generate(self, inputs, prompt="转录以下音频内容", max_new_tokens=500, **kwargs):
215+
def generate(self, inputs, prompt="转录以下音频内容", max_new_tokens=500,
216+
temperature=0.0, top_p=1.0, top_k=-1, repetition_penalty=1.0,
217+
**kwargs):
186218
"""Run batch ASR inference.
187219
188220
Args:
189221
inputs: Audio file path(s), numpy arrays, or tensors.
190222
prompt: Instruction prompt for ASR.
191223
max_new_tokens: Maximum tokens to generate per sample.
224+
temperature: Sampling temperature (0 = greedy decoding).
225+
top_p: Nucleus sampling parameter.
226+
top_k: Top-k sampling (-1 = disabled).
227+
repetition_penalty: Repetition penalty factor. Non-neutral values are
228+
forced back to 1.0 here because this engine feeds vLLM precomputed
229+
embeddings (``enable_prompt_embeds=True``); see
230+
``resolve_repetition_penalty`` and issue #2948.
192231
193232
Returns:
194233
List of {"key": str, "text": str}
@@ -204,7 +243,10 @@ def generate(self, inputs, prompt="转录以下音频内容", max_new_tokens=500
204243

205244
sampling_params = SamplingParams(
206245
max_tokens=max_new_tokens,
207-
temperature=0,
246+
temperature=temperature,
247+
top_p=top_p,
248+
top_k=top_k if top_k > 0 else -1,
249+
repetition_penalty=_safe_repetition_penalty(repetition_penalty),
208250
skip_special_tokens=True,
209251
)
210252

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
"""Unit tests for GLM-ASR vLLM sampling-parameter handling.
2+
3+
These tests exercise ``GLMASRVLLMEngine.generate`` without a GPU or a real
4+
vLLM installation: the vLLM entry points are stubbed in ``sys.modules`` and the
5+
audio/encoder/engine collaborators are mocked, so only the sampling-parameter
6+
wiring is under test.
7+
"""
8+
9+
import re
10+
import sys
11+
import types
12+
import unittest
13+
from unittest import mock
14+
15+
16+
def _install_vllm_stub():
17+
"""Install a minimal ``vllm`` stub whose SamplingParams records kwargs."""
18+
19+
captured = {}
20+
21+
class _RecordingSamplingParams:
22+
def __init__(self, **kwargs):
23+
captured.clear()
24+
captured.update(kwargs)
25+
26+
class _EmbedsPrompt:
27+
def __init__(self, **kwargs):
28+
self.kwargs = kwargs
29+
30+
vllm_mod = types.ModuleType("vllm")
31+
vllm_mod.SamplingParams = _RecordingSamplingParams
32+
vllm_mod.LLM = object
33+
inputs_mod = types.ModuleType("vllm.inputs")
34+
inputs_mod.EmbedsPrompt = _EmbedsPrompt
35+
data_mod = types.ModuleType("vllm.inputs.data")
36+
data_mod.EmbedsPrompt = _EmbedsPrompt
37+
38+
sys.modules["vllm"] = vllm_mod
39+
sys.modules["vllm.inputs"] = inputs_mod
40+
sys.modules["vllm.inputs.data"] = data_mod
41+
return captured
42+
43+
44+
class GLMASRSamplingParamsTest(unittest.TestCase):
45+
def setUp(self):
46+
self.captured = _install_vllm_stub()
47+
from funasr.models.glm_asr.inference_vllm import GLMASRVLLMEngine
48+
49+
# Build an engine without running __init__ (no model load / GPU needed).
50+
engine = GLMASRVLLMEngine.__new__(GLMASRVLLMEngine)
51+
engine.device = "cpu"
52+
engine._encode_audio = mock.Mock(return_value="audio_embeds")
53+
engine._build_prompt_embeds = mock.Mock(
54+
return_value=mock.Mock(float=lambda: "embeds")
55+
)
56+
57+
token_out = types.SimpleNamespace(token_ids=[1, 2, 3])
58+
vllm_output = types.SimpleNamespace(outputs=[token_out])
59+
engine.vllm_engine = mock.Mock()
60+
engine.vllm_engine.generate = mock.Mock(return_value=[vllm_output])
61+
engine.tokenizer = mock.Mock()
62+
engine.tokenizer.decode = mock.Mock(return_value="hello world")
63+
self.engine = engine
64+
65+
def test_defaults_preserve_greedy_behavior(self):
66+
results = self.engine.generate("a.wav")
67+
self.assertEqual(results, [{"key": "a", "text": "hello world"}])
68+
self.assertEqual(self.captured["temperature"], 0.0)
69+
self.assertEqual(self.captured["top_p"], 1.0)
70+
self.assertEqual(self.captured["top_k"], -1)
71+
self.assertEqual(self.captured["repetition_penalty"], 1.0)
72+
self.assertEqual(self.captured["max_tokens"], 500)
73+
74+
def test_caller_sampling_params_are_forwarded(self):
75+
self.engine.generate(
76+
"a.wav", max_new_tokens=128, temperature=0.7, top_p=0.9, top_k=20
77+
)
78+
self.assertEqual(self.captured["max_tokens"], 128)
79+
self.assertEqual(self.captured["temperature"], 0.7)
80+
self.assertEqual(self.captured["top_p"], 0.9)
81+
self.assertEqual(self.captured["top_k"], 20)
82+
83+
def test_non_positive_top_k_is_normalized_to_disabled(self):
84+
self.engine.generate("a.wav", top_k=0)
85+
self.assertEqual(self.captured["top_k"], -1)
86+
87+
def test_repetition_penalty_is_forced_neutral_in_prompt_embeds_mode(self):
88+
# A non-neutral repetition_penalty would crash vLLM prompt-embeds mode
89+
# (issue #2948), so it must be coerced back to 1.0 rather than forwarded.
90+
self.engine.generate("a.wav", repetition_penalty=1.3)
91+
self.assertEqual(self.captured["repetition_penalty"], 1.0)
92+
93+
def test_neutral_repetition_penalty_passes_through(self):
94+
self.engine.generate("a.wav", repetition_penalty=1.0)
95+
self.assertEqual(self.captured["repetition_penalty"], 1.0)
96+
97+
98+
class SafeRepetitionPenaltyTest(unittest.TestCase):
99+
def setUp(self):
100+
_install_vllm_stub()
101+
import funasr.models.glm_asr.inference_vllm as mod
102+
103+
self.mod = mod
104+
# Reset the process-wide warn-once flag between tests.
105+
mod._warned_rep_penalty = False
106+
107+
def test_neutral_and_none_map_to_one(self):
108+
self.assertEqual(self.mod._safe_repetition_penalty(1.0), 1.0)
109+
self.assertEqual(self.mod._safe_repetition_penalty(None), 1.0)
110+
111+
def test_non_neutral_is_coerced_and_warns_once(self):
112+
with self.assertLogs(self.mod.logger, level="WARNING") as cm:
113+
self.assertEqual(self.mod._safe_repetition_penalty(1.5), 1.0)
114+
self.assertTrue(any("2948" in line for line in cm.output))
115+
self.assertTrue(self.mod._warned_rep_penalty)
116+
117+
118+
if __name__ == "__main__":
119+
unittest.main()

0 commit comments

Comments
 (0)