Skip to content

Commit 4795ff0

Browse files
committed
test: use max_batch_size=2 for test_free_generation
Parameterize _make_engine_config with max_batch_size (default=1) so that test_free_generation can create its own pipeline with max_batch_size=2, while all other tests continue using the shared pipe fixture with max_batch_size=1.
1 parent 4ed439c commit 4795ff0

1 file changed

Lines changed: 16 additions & 7 deletions

File tree

tests/test_lmdeploy/test_mtp_guided_decoding.py

Lines changed: 16 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -76,10 +76,10 @@ def _make_spec_config():
7676
return SpeculativeConfig(method='qwen3_5_mtp', num_speculative_tokens=1)
7777

7878

79-
def _make_engine_config():
79+
def _make_engine_config(max_batch_size=1):
8080
"""PytorchEngineConfig suitable for MTP + guided decoding tests."""
8181
return PytorchEngineConfig(
82-
max_batch_size=1,
82+
max_batch_size=max_batch_size,
8383
session_len=1024,
8484
cache_max_entry_count=0.1,
8585
)
@@ -242,11 +242,20 @@ def test_mixed_json_and_free(self, pipe):
242242
class TestSpecDecodeNoGuided:
243243
"""MTP speculative decoding without guided decoding — baseline sanity."""
244244

245-
def test_free_generation(self, pipe):
246-
gen_config = GenerationConfig(max_new_tokens=50)
247-
response = pipe(PROMPT, gen_config=gen_config)
248-
assert response.text, 'Free generation should produce text'
249-
assert response.generate_token_len > 0
245+
def test_free_generation(self):
246+
p = pipeline(
247+
MTP_MODEL_ID,
248+
backend_config=_make_engine_config(max_batch_size=2),
249+
speculative_config=_make_spec_config(),
250+
log_level='INFO',
251+
)
252+
try:
253+
gen_config = GenerationConfig(max_new_tokens=50)
254+
response = p(PROMPT, gen_config=gen_config)
255+
assert response.text, 'Free generation should produce text'
256+
assert response.generate_token_len > 0
257+
finally:
258+
p.close()
250259

251260

252261
# ---------------------------------------------------------------------------

0 commit comments

Comments
 (0)