From 30d2072670abdfb8c5329e6ad243e8b333cf83e2 Mon Sep 17 00:00:00 2001 From: mgoin Date: Fri, 21 Aug 2026 02:35:25 +0000 Subject: [PATCH] [Test] Add focused hybrid MTP prefix-cache regressions Signed-off-by: mgoin --- .buildkite/test_areas/engine.yaml | 11 +- tests/v1/e2e/test_hybrid_chunked_prefill.py | 139 +++++++++++++++++--- 2 files changed, 128 insertions(+), 22 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 7cf03c781602..6bd7f92ef16c 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -133,14 +133,15 @@ steps: depends_on: - image-build-amd -- label: ":nvidia: (H100) V1 E2E Hybrid Chunked Prefill" - key: v1-e2e-4xh100 +- label: ":nvidia: (B200) V1 E2E Hybrid Chunked Prefill" + key: v1-e2e-hybrid-chunked-prefill-b200 timeout_in_minutes: 35 - device: h100 - num_devices: 4 + device: b200-k8s + num_devices: 1 optional: true source_file_dependencies: - vllm/v1/attention/backends/utils.py + - vllm/v1/core/ - vllm/v1/worker/gpu_model_runner.py - tests/v1/e2e/test_hybrid_chunked_prefill.py commands: @@ -149,7 +150,7 @@ steps: amd: label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" dind: false - device: mi300_4 + device: mi300_1 timeout_in_minutes: 35 depends_on: - image-build-amd diff --git a/tests/v1/e2e/test_hybrid_chunked_prefill.py b/tests/v1/e2e/test_hybrid_chunked_prefill.py index dd8a5f5cb6d4..dd24f3884ea2 100644 --- a/tests/v1/e2e/test_hybrid_chunked_prefill.py +++ b/tests/v1/e2e/test_hybrid_chunked_prefill.py @@ -5,8 +5,30 @@ from vllm import SamplingParams from vllm.platforms import current_platform +from vllm.v1.kv_cache_interface import MambaSpec -from ...utils import large_gpu_mark, multi_gpu_marks +from ...utils import large_gpu_mark + +QWEN_MODEL = "Qwen/Qwen3.5-4B" +QWEN_KV_CACHE_BYTES = 12 << 30 +HYBRID_MTP_MODELS = [ + pytest.param( + QWEN_MODEL, + marks=[large_gpu_mark(min_gb=30)], + id="qwen", + ), + pytest.param( + "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-NVFP4", + marks=[large_gpu_mark(min_gb=80)] + + [ + pytest.mark.skipif( + not current_platform.is_cuda(), + reason="modelopt quantization is supported only on CUDA", + ) + ], + id="nemotron", + ), +] # A trivial request with a short prompt to ensure we run a mixed batch SMALL_MESSAGE = [ @@ -38,20 +60,7 @@ @pytest.mark.parametrize( "model_name", - [ - pytest.param("Qwen/Qwen3.5-4B", marks=[large_gpu_mark(min_gb=40)]), - pytest.param( - "nvidia/NVIDIA-Nemotron-3-Super-120B-A12B-FP8", - marks=[large_gpu_mark(min_gb=80)] - + multi_gpu_marks(num_gpus=4) - + [ - pytest.mark.skipif( - not current_platform.is_cuda(), - reason="modelopt quantization is supported only on CUDA", - ) - ], - ), - ], + HYBRID_MTP_MODELS, ) @pytest.mark.parametrize("enable_prefix_caching", [False, True]) def test_mtp_speculative_mixed_batch_short_prefill( @@ -62,7 +71,7 @@ def test_mtp_speculative_mixed_batch_short_prefill( # Set so large that both prefills will be classified as decodes in a mixed batch # note, with prefix caching we require chunk_size >= mamba_block_size - chunk_size = 256 if not enable_prefix_caching else 16384 + chunk_size = 256 if not enable_prefix_caching else 2048 num_draft_tokens = 100 with vllm_runner( @@ -72,9 +81,11 @@ def test_mtp_speculative_mixed_batch_short_prefill( "num_speculative_tokens": num_draft_tokens, }, max_num_batched_tokens=chunk_size, + max_num_seqs=4, max_model_len=512, + kv_cache_memory_bytes=QWEN_KV_CACHE_BYTES if model_name == QWEN_MODEL else None, enforce_eager=True, - tensor_parallel_size=4, + tensor_parallel_size=1, trust_remote_code=True, enable_chunked_prefill=True, enable_prefix_caching=enable_prefix_caching, @@ -108,3 +119,97 @@ def test_mtp_speculative_mixed_batch_short_prefill( assert "43" in responses[1], ( "The second response should contain the correct value of 42+1=43." ) + + +def _get_mamba_block_size(llm) -> int: + scheduler = llm.llm_engine.engine_core.engine_core.scheduler + block_sizes = { + group.kv_cache_spec.block_size + for group in scheduler.kv_cache_config.kv_cache_groups + if isinstance(group.kv_cache_spec, MambaSpec) + } + assert len(block_sizes) == 1 + block_size = block_sizes.pop() + assert scheduler.cache_config.block_size == block_size + return block_size + + +def _build_access_code_manual(tokenizer, target_tokens: int) -> tuple[str, list[str]]: + codes = ["605341", "693278", "597596", "751982"] + header = "Memorize this facility manual and its access codes.\n\n" + facts = "".join( + f"The access code for vault-{i:02d} is {code}.\n" + for i, code in enumerate(codes) + ) + filler = "Routine facility inspections are recorded in the audit ledger every day. " + manual = header + facts + while len(tokenizer.encode(manual)) < target_tokens: + manual = header + filler + manual.removeprefix(header) + return manual, codes + + +@large_gpu_mark(min_gb=30) +def test_qwen_mtp_mamba_prefix_cache_hit_is_bounded(vllm_runner, monkeypatch): + """MTP must not extend a hybrid Mamba hit past the attention hit.""" + monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0") + + with vllm_runner( + QWEN_MODEL, + tensor_parallel_size=1, + max_model_len=8192, + kv_cache_memory_bytes=QWEN_KV_CACHE_BYTES, + block_size=None, + enforce_eager=True, + enable_chunked_prefill=True, + enable_prefix_caching=True, + mamba_cache_mode="align", + speculative_config={"method": "mtp", "num_speculative_tokens": 2}, + ) as runner: + llm = runner.get_llm() + block_size = _get_mamba_block_size(llm) + manual, codes = _build_access_code_manual( + llm.get_tokenizer(), 2 * block_size - 192 + ) + manual_tokens = len(llm.get_tokenizer().encode(manual)) + assert block_size < manual_tokens < 2 * block_size + + wave1_prompts = [ + manual + + f"\nDescribe the audit procedure for vault-{i:02d} in detail.\nAnswer:" + for i in range(len(codes)) + ] + decode_tokens = 2 * block_size - manual_tokens + 96 + wave1_params = SamplingParams( + temperature=0.0, + min_tokens=decode_tokens, + max_tokens=decode_tokens, + ignore_eos=True, + ) + wave1_outputs = llm.generate(wave1_prompts, wave1_params) + + wave2_prompts = [ + prompt + + output.outputs[0].text + + f"\n\nWhat is the access code for vault-{i:02d}? Answer:" + for i, (prompt, output) in enumerate(zip(wave1_prompts, wave1_outputs)) + ] + warm_params = SamplingParams(temperature=0.0, max_tokens=24, stop=["\n"]) + cold_params = SamplingParams( + temperature=0.0, + max_tokens=24, + stop=["\n"], + skip_reading_prefix_cache=True, + ) + warm_outputs = llm.generate(wave2_prompts, warm_params) + cold_outputs = llm.generate(wave2_prompts, cold_params) + + cached_tokens = [output.num_cached_tokens for output in warm_outputs] + assert max(cached_tokens) == block_size + + for arm, outputs in (("warm", warm_outputs), ("cold", cold_outputs)): + missed = [ + i + for i, (code, output) in enumerate(zip(codes, outputs)) + if code not in output.outputs[0].text + ] + assert not missed, f"{arm} cache missed access codes for prompts {missed}"