iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[CI][Mamba] Use well-conditioned prompts for FlashInfer ReplaySSM parity tests (#60515)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com>
This commit is contained in:
Jiangyun Zhu
2026-10-08 17:39:56 +08:00
committed by GitHub
co-authored by Claude Opus 5.5
parent 7816afdb25
commit e265b77982
1 changed files with 16 additions and 4 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tests/v1/e2e/test_replayssm_decode.py
+16 -4
View File
@@ -25,6 +25,12 @@ PROMPTS = [
"The capital of France is",
"Once upon a time, in a small village,",
]
# The NVFP4 model is not reproducible across engine instances; these prompts
# avoid near-ties and decode confidently past the replay window.
FLASHINFER_PROMPTS = [
"The capital of France is",
"It was the best of times, it was the worst of times,",
]
@pytest.fixture(autouse=True)
@@ -43,6 +49,7 @@ def _check_replayssm_parity(
*,
tensor_parallel_size=1,
mamba_backend: str = "triton",
prompts: list[str] = PROMPTS,
name_1: str = "replayssm",
expected_v2: bool | None = None,
):
@@ -60,7 +67,7 @@ def _check_replayssm_parity(
with vllm_runner(model_name, **common) as llm:
if expected_v2 is not None:
assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2
baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
baseline = llm.generate_greedy_logprobs(prompts, max_tokens=32, num_logprobs=5)
with vllm_runner(
model_name,
use_replayssm=True,
@@ -69,7 +76,7 @@ def _check_replayssm_parity(
) as llm:
if expected_v2 is not None:
assert llm.llm.llm_engine.vllm_config.use_v2_model_runner is expected_v2
replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
replay = llm.generate_greedy_logprobs(prompts, max_tokens=32, num_logprobs=5)
check_logprobs_close(
outputs_0_lst=baseline,
@@ -105,6 +112,7 @@ def test_replayssm_flashinfer_decode_matches_baseline(
vllm_runner,
model_name,
mamba_backend="flashinfer",
prompts=FLASHINFER_PROMPTS,
name_1="replayssm_flashinfer",
expected_v2=use_v2_model_runner,
)
@@ -127,7 +135,9 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na
},
)
with vllm_runner(model_name, mamba_backend="flashinfer", **common) as llm:
baseline = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
baseline = llm.generate_greedy_logprobs(
FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5
)
with vllm_runner(
model_name,
use_replayssm=True,
@@ -135,7 +145,9 @@ def test_replayssm_flashinfer_spec_decode_matches_baseline(vllm_runner, model_na
mamba_backend="flashinfer",
**common,
) as llm:
replay = llm.generate_greedy_logprobs(PROMPTS, max_tokens=32, num_logprobs=5)
replay = llm.generate_greedy_logprobs(
FLASHINFER_PROMPTS, max_tokens=32, num_logprobs=5
)
check_logprobs_close(
outputs_0_lst=baseline,