| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import contextlib |
| 4 | import os |
| 5 | import random |
| 6 | |
| 7 | import pytest |
| 8 | import torch |
| 9 | from utils import ( |
| 10 | _extract_step_logprobs, |
| 11 | _random_prompt, |
| 12 | skip_if_not_cuda, |
| 13 | ) |
| 14 | |
| 15 | from vllm import LLM, SamplingParams |
| 16 | |
| 17 | pytestmark = pytest.mark.skipif( |
| 18 | not hasattr(torch, "float8_e4m3fn"), |
| 19 | reason="NVFP4 tests require torch.float8_e4m3fn support.", |
| 20 | ) |
| 21 | |
| 22 | NVFP4_TEST_MODEL = os.getenv( |
| 23 | "VLLM_TEST_NVFP4_MODEL", "nm-testing/TinyLlama-1.1B-Chat-v1.0-NVFP4" |
| 24 | ) |
| 25 | |
| 26 | |
| 27 | def _make_llm(max_num_seqs: int, backend: str) -> LLM: |
| 28 | return LLM( |
| 29 | model=NVFP4_TEST_MODEL, |
| 30 | max_num_seqs=max_num_seqs, |
| 31 | gpu_memory_utilization=float( |
| 32 | os.getenv("VLLM_NVFP4_TEST_GPU_MEMORY_UTILIZATION", "0.05") |
| 33 | ), |
| 34 | max_model_len=int(os.getenv("VLLM_NVFP4_TEST_MAX_MODEL_LEN", "2048")), |
| 35 | dtype="auto", |
| 36 | tensor_parallel_size=int(os.getenv("VLLM_NVFP4_TEST_TP_SIZE", "1")), |
| 37 | enable_prefix_caching=False, |
| 38 | enforce_eager=True, |
| 39 | attention_config={"backend": backend}, |
| 40 | ) |
| 41 | |
| 42 | |
| 43 | @skip_if_not_cuda |
| 44 | @pytest.mark.parametrize("backend", ["FLASH_ATTN"]) |
| 45 | def test_dense_nvfp4_generation_is_deterministic_across_batch_sizes_e2e(backend): |
| 46 | seed = int(os.getenv("VLLM_TEST_SEED", "12345")) |
| 47 | random.seed(seed) |
| 48 | |
| 49 | num_trials = int(os.getenv("VLLM_NVFP4_NEEDLE_TRIALS", "2")) |
| 50 | max_batch_size = int(os.getenv("VLLM_NVFP4_NEEDLE_BATCH_SIZE", "8")) |
| 51 | min_random_prompt = int(os.getenv("VLLM_NVFP4_MIN_PROMPT", "32")) |
| 52 | max_random_prompt = int(os.getenv("VLLM_NVFP4_MAX_PROMPT", "96")) |
| 53 | assert max_batch_size >= 2, "Batch size should be >= 2 to test invariance." |
| 54 | |
| 55 | sampling = SamplingParams( |
| 56 | temperature=float(os.getenv("VLLM_NVFP4_NEEDLE_TEMPERATURE", "0.6")), |
| 57 | top_p=float(os.getenv("VLLM_NVFP4_NEEDLE_TOP_P", "0.95")), |
| 58 | max_tokens=int(os.getenv("VLLM_NVFP4_NEEDLE_MAX_TOKENS", "16")), |
| 59 | seed=20240919, |
| 60 | logprobs=5, |
| 61 | ) |
| 62 | needle_prompt = "Write one factual sentence about the moon." |
| 63 | |
| 64 | llm = None |
| 65 | baseline_completion = None |
| 66 | baseline_logprobs = None |
| 67 | try: |
| 68 | llm = _make_llm(max_num_seqs=max_batch_size, backend=backend) |
| 69 | baseline_output = llm.generate([needle_prompt], sampling, use_tqdm=False)[0] |
| 70 | baseline_completion = baseline_output.outputs[0] |
| 71 | baseline_logprobs, baseline_token_ids = _extract_step_logprobs(baseline_output) |
| 72 | assert baseline_logprobs is not None |
| 73 | assert baseline_token_ids is not None |
| 74 | for _ in range(num_trials): |
| 75 | batch_size = random.randint(max_batch_size // 2, max_batch_size) |
| 76 | needle_pos = random.randint(0, batch_size - 1) |
| 77 | prompts: list[str] = [] |
| 78 | for idx in range(batch_size): |
| 79 | if idx == needle_pos: |
| 80 | prompts.append(needle_prompt) |
| 81 | else: |
| 82 | prompts.append(_random_prompt(min_random_prompt, max_random_prompt)) |
| 83 | |
| 84 | outputs = llm.generate(prompts, sampling, use_tqdm=False) |
| 85 | needle_output = outputs[needle_pos] |
| 86 | needle_completion = needle_output.outputs[0] |
| 87 | needle_logprobs, needle_token_ids = _extract_step_logprobs(needle_output) |
| 88 | assert needle_logprobs is not None |
| 89 | assert needle_token_ids is not None |
| 90 | |
| 91 | assert needle_output.prompt == needle_prompt |
| 92 | assert baseline_completion is not None |
| 93 | assert baseline_logprobs is not None |
| 94 | assert needle_completion.token_ids == baseline_completion.token_ids |
| 95 | assert needle_completion.text == baseline_completion.text |
| 96 | torch.testing.assert_close(needle_logprobs, baseline_logprobs) |
| 97 | finally: |
| 98 | if llm is not None: |
| 99 | with contextlib.suppress(Exception): |
| 100 | llm.llm_engine.engine_core.shutdown() |