| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import asyncio |
| 5 | from contextlib import AsyncExitStack |
| 6 | from dataclasses import dataclass, replace |
| 7 | from typing import Any |
| 8 | |
| 9 | import pytest |
| 10 | import torch |
| 11 | |
| 12 | from tests.utils import multi_gpu_marks |
| 13 | from vllm import LLM, SamplingParams |
| 14 | from vllm.config import CompilationConfig |
| 15 | from vllm.distributed import cleanup_dist_env_and_memory |
| 16 | from vllm.engine.arg_utils import AsyncEngineArgs |
| 17 | from vllm.platforms import current_platform |
| 18 | from vllm.sampling_params import RequestOutputKind |
| 19 | from vllm.v1.engine.async_llm import AsyncLLM |
| 20 | from vllm.v1.metrics.reader import Metric |
| 21 | |
| 22 | DEEPSEEK_MTP_MAIN_RANDOM = "luccafong/deepseek_mtp_main_random" |
| 23 | DEEPSEEK_MTP_DRAFT_RANDOM = "luccafong/deepseek_mtp_draft_random" |
| 24 | |
| 25 | PROMPT = "The capital of France is" |
| 26 | MAX_TOKENS = 8 |
| 27 | MAX_MODEL_LEN = 2048 |
| 28 | GPU_MEM_UTIL = 0.85 |
| 29 | # Spec / no-spec greedy output should match, but exact match is not guaranteed |
| 30 | # across parallelism layouts (reduction order can flip near-tie argmaxes), so |
| 31 | # gate on a similarity ratio like the spec-decode E2E tests. |
| 32 | MIN_MATCH_RATIO = 0.8 |
| 33 | |
| 34 | |
| 35 | def _token_match_ratio(a: tuple[int, ...], b: tuple[int, ...]) -> float: |
| 36 | n = min(len(a), len(b)) |
| 37 | return sum(x == y for x, y in zip(a, b)) / n if n else 0.0 |
| 38 | |
| 39 | |
| 40 | @dataclass(frozen=True) |
| 41 | class InlineConfig: |
| 42 | id: str |
| 43 | tp_size: int = 1 |
| 44 | pp_size: int = 1 |
| 45 | enable_expert_parallel: bool = False |
| 46 | enable_eplb: bool = False |
| 47 | skip_reason: str | None = None |
| 48 | |
| 49 | |
| 50 | INLINE_CONFIGS = [ |
| 51 | InlineConfig(id="tp2", tp_size=2), |
| 52 | InlineConfig(id="ep2", tp_size=2, enable_expert_parallel=True), |
| 53 | InlineConfig( |
| 54 | id="ep2_eplb", |
| 55 | tp_size=2, |
| 56 | enable_expert_parallel=True, |
| 57 | enable_eplb=True, |
| 58 | ), |
| 59 | InlineConfig(id="pp2", pp_size=2), |
| 60 | ] |
| 61 | |
| 62 | |
| 63 | def _inline_kwargs(config: InlineConfig, *, with_spec: bool) -> dict[str, Any]: |
| 64 | kwargs: dict[str, Any] = { |
| 65 | "model": DEEPSEEK_MTP_MAIN_RANDOM, |
| 66 | "tensor_parallel_size": config.tp_size, |
| 67 | "pipeline_parallel_size": config.pp_size, |
| 68 | "enable_expert_parallel": config.enable_expert_parallel, |
| 69 | "max_model_len": MAX_MODEL_LEN, |
| 70 | "gpu_memory_utilization": GPU_MEM_UTIL, |
| 71 | "enforce_eager": True, |
| 72 | "trust_remote_code": True, |
| 73 | "disable_log_stats": False, |
| 74 | # MLA + batch invariance currently disables prefix caching at runtime |
| 75 | # (see vllm/v1/attention/backends/mla/common.py); set explicitly so the |
| 76 | # configuration is unambiguous. |
| 77 | "enable_prefix_caching": False, |
| 78 | } |
| 79 | if config.enable_eplb: |
| 80 | kwargs["enable_eplb"] = True |
| 81 | # Rearrangement first fires after step_interval // 4 forward steps |
| 82 | # (the step counter starts at 3/4 of step_interval, see |
| 83 | # vllm/distributed/eplb/eplb_state.py). Keep these small so the EPLB |
| 84 | # routine actually runs within the short MAX_TOKENS generation instead |
| 85 | # of never triggering. |
| 86 | kwargs["eplb_config"] = { |
| 87 | "num_redundant_experts": config.tp_size, |
| 88 | "window_size": 2, |
| 89 | "step_interval": 4, |
| 90 | "log_balancedness": False, |
| 91 | } |
| 92 | if with_spec: |
| 93 | kwargs["speculative_config"] = { |
| 94 | "method": "mtp", |
| 95 | "model": DEEPSEEK_MTP_DRAFT_RANDOM, |
| 96 | "num_speculative_tokens": 1, |
| 97 | } |
| 98 | return kwargs |
| 99 | |
| 100 | |
| 101 | def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]: |
| 102 | outputs = llm.generate( |
| 103 | [PROMPT], |
| 104 | SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True), |
| 105 | ) |
| 106 | assert outputs and outputs[0].outputs, "expected one completion" |
| 107 | token_ids = tuple(outputs[0].outputs[0].token_ids) |
| 108 | assert len(token_ids) == MAX_TOKENS, ( |
| 109 | f"expected {MAX_TOKENS} generated tokens, got {len(token_ids)}: {token_ids}" |
| 110 | ) |
| 111 | return token_ids |
| 112 | |
| 113 | |
| 114 | def _spec_decode_num_drafts(metrics: list[Metric]) -> int: |
| 115 | name2metric = {m.name: m for m in metrics} |
| 116 | counter = name2metric.get("vllm:spec_decode_num_drafts") |
| 117 | assert counter is not None, ( |
| 118 | "spec_decode_num_drafts metric missing; check disable_log_stats=False" |
| 119 | ) |
| 120 | return int(counter.value) |
| 121 | |
| 122 | |
| 123 | def _enable_batch_invariance(monkeypatch: pytest.MonkeyPatch) -> None: |
| 124 | # This test exercises DeepSeek MLA prefill. ROCm's upstream FlashAttention |
| 125 | # varlen API lacks the num_splits argument used by batch invariance, so |
| 126 | # explicitly override any ambient setting on that path. |
| 127 | if current_platform.is_cuda(): |
| 128 | monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") |
| 129 | else: |
| 130 | monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0") |
| 131 | |
| 132 | |
| 133 | @pytest.mark.parametrize( |
| 134 | "config", |
| 135 | [ |
| 136 | pytest.param(c, id=c.id, marks=multi_gpu_marks(num_gpus=2)) |
| 137 | for c in INLINE_CONFIGS |
| 138 | ], |
| 139 | ) |
| 140 | def test_deepseek_mtp_load_inline( |
| 141 | monkeypatch: pytest.MonkeyPatch, |
| 142 | config: InlineConfig, |
| 143 | vllm_runner, |
| 144 | ): |
| 145 | """MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy.""" |
| 146 | if config.skip_reason is not None: |
| 147 | pytest.skip(config.skip_reason) |
| 148 | |
| 149 | # Reduces run-to-run nondeterminism so spec / no-spec stay close; |
| 150 | # see tests/v1/distributed/test_eagle_dp.py for the same pattern. |
| 151 | _enable_batch_invariance(monkeypatch) |
| 152 | |
| 153 | spec_kwargs = _inline_kwargs(config, with_spec=True) |
| 154 | spec_model = spec_kwargs.pop("model") |
| 155 | with vllm_runner( |
| 156 | spec_model, |
| 157 | block_size=None, |
| 158 | enable_chunked_prefill=None, |
| 159 | compilation_config=CompilationConfig(), |
| 160 | **spec_kwargs, |
| 161 | ) as spec_runner: |
| 162 | spec_tokens = _generate_token_ids_inline(spec_runner.llm) |
| 163 | n_drafts = _spec_decode_num_drafts(spec_runner.llm.get_metrics()) |
| 164 | |
| 165 | no_spec_kwargs = _inline_kwargs(config, with_spec=False) |
| 166 | no_spec_model = no_spec_kwargs.pop("model") |
| 167 | with vllm_runner( |
| 168 | no_spec_model, |
| 169 | block_size=None, |
| 170 | enable_chunked_prefill=None, |
| 171 | compilation_config=CompilationConfig(), |
| 172 | **no_spec_kwargs, |
| 173 | ) as no_spec_runner: |
| 174 | no_spec_tokens = _generate_token_ids_inline(no_spec_runner.llm) |
| 175 | |
| 176 | # Non-vacuity: a silently-broken MTP drafter that falls back to |
| 177 | # verifier-only decoding would still produce matching output below. |
| 178 | assert n_drafts > 0, ( |
| 179 | f"MTP drafter never fired under {config.id}: vllm:spec_decode_num_drafts == 0" |
| 180 | ) |
| 181 | match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens) |
| 182 | print(f"\n{config.id}: spec/no-spec match_ratio={match_ratio:.3f}") |
| 183 | assert match_ratio >= MIN_MATCH_RATIO, ( |
| 184 | f"Spec / no-spec output divergence under {config.id}: " |
| 185 | f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n" |
| 186 | f" spec_tokens = {spec_tokens}\n" |
| 187 | f" no_spec_tokens= {no_spec_tokens}" |
| 188 | ) |
| 189 | |
| 190 | |
| 191 | @pytest.mark.asyncio |
| 192 | @pytest.mark.parametrize( |
| 193 | "dp_size", |
| 194 | [ |
| 195 | pytest.param(2, marks=pytest.mark.distributed(num_gpus=2), id="dp2"), |
| 196 | pytest.param(1, id="dp1"), |
| 197 | ], |
| 198 | ) |
| 199 | async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: int): |
| 200 | """MTP loads under DP (on and off) via AsyncLLM; spec matches no-spec greedy.""" |
| 201 | if torch.accelerator.device_count() < dp_size: |
| 202 | pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs") |
| 203 | |
| 204 | _enable_batch_invariance(monkeypatch) |
| 205 | |
| 206 | base_args = AsyncEngineArgs( |
| 207 | model=DEEPSEEK_MTP_MAIN_RANDOM, |
| 208 | tensor_parallel_size=1, |
| 209 | data_parallel_size=dp_size, |
| 210 | data_parallel_backend="mp", |
| 211 | max_model_len=MAX_MODEL_LEN, |
| 212 | gpu_memory_utilization=GPU_MEM_UTIL, |
| 213 | enforce_eager=True, |
| 214 | trust_remote_code=True, |
| 215 | enable_prefix_caching=False, |
| 216 | ) |
| 217 | spec_args = replace( |
| 218 | base_args, |
| 219 | speculative_config={ |
| 220 | "method": "mtp", |
| 221 | "model": DEEPSEEK_MTP_DRAFT_RANDOM, |
| 222 | "num_speculative_tokens": 1, |
| 223 | }, |
| 224 | ) |
| 225 | |
| 226 | sampling_params = SamplingParams( |
| 227 | max_tokens=MAX_TOKENS, |
| 228 | ignore_eos=True, |
| 229 | output_kind=RequestOutputKind.FINAL_ONLY, |
| 230 | temperature=0.0, |
| 231 | ) |
| 232 | |
| 233 | async def _generate(args: AsyncEngineArgs, request_id: str) -> tuple[int, ...]: |
| 234 | try: |
| 235 | async with AsyncExitStack() as stack: |
| 236 | engine = AsyncLLM.from_engine_args(args) |
| 237 | stack.callback(engine.shutdown) |
| 238 | async for out in engine.generate( |
| 239 | request_id=request_id, |
| 240 | prompt=PROMPT, |
| 241 | sampling_params=sampling_params, |
| 242 | ): |
| 243 | token_ids = tuple(out.outputs[0].token_ids) |
| 244 | assert len(token_ids) == MAX_TOKENS, ( |
| 245 | f"{request_id}: expected {MAX_TOKENS} generated tokens, " |
| 246 | f"got {len(token_ids)}: {token_ids}" |
| 247 | ) |
| 248 | return token_ids |
| 249 | raise AssertionError("AsyncLLM produced no output") |
| 250 | finally: |
| 251 | torch.accelerator.empty_cache() |
| 252 | cleanup_dist_env_and_memory() |
| 253 | |
| 254 | spec_tokens = await asyncio.wait_for( |
| 255 | _generate(spec_args, f"deepseek-mtp-dp{dp_size}-spec"), timeout=600 |
| 256 | ) |
| 257 | no_spec_tokens = await asyncio.wait_for( |
| 258 | _generate(base_args, f"deepseek-mtp-dp{dp_size}-no-spec"), timeout=600 |
| 259 | ) |
| 260 | |
| 261 | match_ratio = _token_match_ratio(spec_tokens, no_spec_tokens) |
| 262 | print(f"\ndp{dp_size}: spec/no-spec match_ratio={match_ratio:.3f}") |
| 263 | assert match_ratio >= MIN_MATCH_RATIO, ( |
| 264 | f"Spec / no-spec output divergence under dp{dp_size}: " |
| 265 | f"match_ratio={match_ratio:.2f} < {MIN_MATCH_RATIO}.\n" |
| 266 | f" spec_tokens = {spec_tokens}\n" |
| 267 | f" no_spec_tokens= {no_spec_tokens}" |
| 268 | ) |