forked from Karylab-cklius/vllm
Compare commits
24
Commits
@@ -98,8 +98,15 @@ steps:
|
||||
commands:
|
||||
- "bash .buildkite/scripts/generate-and-upload-nightly-index.sh"
|
||||
|
||||
- block: "Unblock to build release Docker images"
|
||||
depends_on: ~
|
||||
key: block-build-release-images
|
||||
if: build.env("NIGHTLY") != "1"
|
||||
|
||||
- group: "Build release Docker images"
|
||||
key: "build-release-images"
|
||||
depends_on: block-build-release-images
|
||||
allow_dependency_failure: true
|
||||
steps:
|
||||
- label: "Build release image - x86_64 - CUDA 12.9"
|
||||
depends_on: ~
|
||||
@@ -617,6 +624,8 @@ steps:
|
||||
- label: ":docker: Build release image - x86_64 - ROCm"
|
||||
id: build-rocm-release-image
|
||||
depends_on:
|
||||
- step: block-build-release-images
|
||||
allow_failure: true
|
||||
- step: build-rocm-base-wheels
|
||||
allow_failure: false
|
||||
agents:
|
||||
|
||||
@@ -196,7 +196,6 @@ steps:
|
||||
- VLLM_ALLOW_INSECURE_SERIALIZATION=1 python3 examples/rl/rlhf_async_new_apis.py
|
||||
- VLLM_USE_DEEP_GEMM=1 VLLM_LOGGING_LEVEL=DEBUG python3 examples/offline_inference/data_parallel.py --model=Qwen/Qwen1.5-MoE-A2.7B -tp=1 -dp=2 --max-model-len=2048 --all2all-backend=deepep_high_throughput
|
||||
- pytest -v -s tests/v1/distributed/test_dbo.py
|
||||
- TP_SIZE=1 DP_SIZE=2 pytest -v -s tests/v1/distributed/test_eagle_dp.py
|
||||
|
||||
- label: Distributed Tests (2 GPUs)(B200)
|
||||
device: b200
|
||||
|
||||
@@ -42,3 +42,16 @@ steps:
|
||||
- tests/v1/e2e/spec_decode/
|
||||
commands:
|
||||
- pytest -v -s v1/e2e/spec_decode -k "draft_model or no_sync or batch_inference"
|
||||
|
||||
- label: DFlash Speculators Correctness
|
||||
timeout_in_minutes: 30
|
||||
device: h100
|
||||
optional: true
|
||||
num_devices: 1
|
||||
source_file_dependencies:
|
||||
- vllm/v1/spec_decode/
|
||||
- vllm/model_executor/models/qwen3_dflash.py
|
||||
- tests/v1/spec_decode/test_speculators_dflash.py
|
||||
commands:
|
||||
- export VLLM_ALLOW_INSECURE_SERIALIZATION=1
|
||||
- pytest -v -s v1/spec_decode/test_speculators_dflash.py -m slow_test
|
||||
|
||||
@@ -37,6 +37,7 @@ th {
|
||||
| HuggingFace-Blazedit | ✅ | ✅ | `vdaita/edit_5k_char`, `vdaita/edit_10k_char` |
|
||||
| HuggingFace-ASR | ✅ | ✅ | `openslr/librispeech_asr`, `facebook/voxpopuli`, `LIUM/tedlium`, `edinburghcstr/ami`, `speechcolab/gigaspeech`, `kensho/spgispeech` |
|
||||
| Spec Bench | ✅ | ✅ | `wget https://raw.githubusercontent.com/hemingkx/Spec-Bench/refs/heads/main/data/spec_bench/question.jsonl` |
|
||||
| SPEED-Bench | ✅ | ✅ | `curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -` |
|
||||
| Custom | ✅ | ✅ | Local file: `data.jsonl` |
|
||||
| Custom MM | ✅ | ✅ | Local file: `mm_data.jsonl` |
|
||||
|
||||
@@ -239,6 +240,69 @@ vllm bench serve \
|
||||
--spec-bench-category "summarization"
|
||||
```
|
||||
|
||||
#### SPEED-Bench Benchmark with Speculative Decoding
|
||||
|
||||
[SPEED-Bench](https://huggingface.co/datasets/nvidia/SPEED-Bench) is a unified and diverse dataset for speculative decoding, supporting acceptance rate and length measurements using the Qualitative split and throughput measurements using the Throughput splits in 5 configuration of input sequence length (1k, 2k, 8k, 16k, 32k).
|
||||
|
||||
!!! note
|
||||
This dataset is governed by the [NVIDIA Evaluation Dataset License Agreement](https://huggingface.co/datasets/nvidia/SPEED-Bench/blob/main/License.pdf). For each dataset a user elects to use, the user is responsible for checking if the dataset license is fit for the intended purpose. The `prepare.py` script automatically fetches data from all the source datasets.
|
||||
|
||||
First, download the dataset to a folder, using this one liner:
|
||||
|
||||
```bash
|
||||
curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py \| python3 -
|
||||
```
|
||||
|
||||
The command supports also the following arguments:
|
||||
|
||||
- `--config`: download only a subset of the dataset: `qualitative`, `throughput_1k`, `throughput_2k`, `throughput_8k`, `throughput_16k` and `throughput_32k`. By default, it will download all subsets.
|
||||
- `--output_dir`: download to a specified folder. By default, it will download to the current directory.
|
||||
|
||||
Start a server with speculative decoding:
|
||||
|
||||
```bash
|
||||
vllm serve meta-llama/Llama-3.3-70B-Instruct \
|
||||
--speculative-config $'{"method": "eagle3",
|
||||
"num_speculative_tokens": 3,
|
||||
"model": "nvidia/Llama-3.3-70B-Instruct-Eagle3"}'
|
||||
```
|
||||
|
||||
Run all categories in the Qualitative split:
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model meta-llama/Llama-3.3-70B-Instruct \
|
||||
--dataset-name speed_bench \
|
||||
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
|
||||
--num-prompts -1
|
||||
```
|
||||
|
||||
Available categories include `[writing, roleplay, reasoning, math, coding, stem, humanities, multilingual, summarization, qa, rag]`.
|
||||
|
||||
Run only a specific category like "multilingual":
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model meta-llama/Llama-3.3-70B-Instruct \
|
||||
--dataset-name speed_bench \
|
||||
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench" \
|
||||
--num-prompts -1
|
||||
--speed-bench-category "multilingual"
|
||||
```
|
||||
|
||||
Run all categories in the Throughput split (2k ISL):
|
||||
|
||||
```bash
|
||||
vllm bench serve \
|
||||
--model meta-llama/Llama-3.3-70B-Instruct \
|
||||
--dataset-name speed_bench \
|
||||
--speed-bench-dataset-subset throughput_2k
|
||||
--dataset-path "<YOUR_DOWNLOADED_PATH>/data/speed_bench/" \
|
||||
--num-prompts -1
|
||||
```
|
||||
|
||||
Available categories include `[high_entropy, mixed, low_entropy]`, where high entropy data contains unstructued data such as creative writing while low entropy data contains more structured data such as coding, more details are in the dataset card.
|
||||
|
||||
#### Other HuggingFaceDataset Examples
|
||||
|
||||
```bash
|
||||
|
||||
@@ -693,6 +693,12 @@ class precompiled_wheel_utils:
|
||||
flash_attn_regex = re.compile(
|
||||
r"vllm/vllm_flash_attn/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
|
||||
)
|
||||
# __init__.py and flash_attn_interface.py are source-controlled
|
||||
# in vllm and should not be overwritten (matches cmake exclusions)
|
||||
flash_attn_files_to_skip = {
|
||||
"vllm/vllm_flash_attn/__init__.py",
|
||||
"vllm/vllm_flash_attn/flash_attn_interface.py",
|
||||
}
|
||||
triton_kernels_regex = re.compile(
|
||||
r"vllm/third_party/triton_kernels/(?:[^/.][^/]*/)*(?!\.)[^/]*\.py"
|
||||
)
|
||||
@@ -705,7 +711,11 @@ class precompiled_wheel_utils:
|
||||
filter(lambda x: x.filename in files_to_copy, wheel.filelist)
|
||||
)
|
||||
file_members += list(
|
||||
filter(lambda x: flash_attn_regex.match(x.filename), wheel.filelist)
|
||||
filter(
|
||||
lambda x: flash_attn_regex.match(x.filename)
|
||||
and x.filename not in flash_attn_files_to_skip,
|
||||
wheel.filelist,
|
||||
)
|
||||
)
|
||||
file_members += list(
|
||||
filter(
|
||||
|
||||
@@ -3,6 +3,7 @@
|
||||
import pytest
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch._library.triton import set_wrap_triton_enabled
|
||||
|
||||
import vllm.kernels # noqa: F401 to register kernels
|
||||
from vllm import ir
|
||||
@@ -47,6 +48,7 @@ def test_lowering_rms_norm(rms_provider, default_vllm_config):
|
||||
with (
|
||||
ops.rms_norm.set_priority([rms_provider, "native"]),
|
||||
ir.enable_torch_wrap(True),
|
||||
set_wrap_triton_enabled(False), # set by default in forward context
|
||||
):
|
||||
compiled_model = torch.compile(model, backend=backend, fullgraph=True)
|
||||
compiled_unlowered_model = torch.compile(
|
||||
|
||||
@@ -222,3 +222,47 @@ def test_model_specialization_with_evaluate_guards(
|
||||
torch.randn(1, 10).cuda(),
|
||||
is_01_specialization=True,
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.skipif(not is_torch_equal_or_newer("2.10.0"), reason="requires torch 2.10")
|
||||
def test_piecewise_backend_empty_sym_shape_indices():
|
||||
"""Test that PiecewiseBackend handles empty sym_shape_indices correctly.
|
||||
|
||||
When all inputs have static shapes (no torch.SymInt), sym_shape_indices
|
||||
will be empty. The fix in PiecewiseBackend.__call__ handles this case
|
||||
by using the first compiled range_entry.
|
||||
"""
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
# Use small max_model_len and max_num_batched_tokens to encourage
|
||||
# static shape compilation with empty sym_shape_indices
|
||||
llm = LLM(
|
||||
model="Qwen/Qwen3-0.6B",
|
||||
max_model_len=512,
|
||||
max_num_batched_tokens=1,
|
||||
compilation_config={
|
||||
"mode": CompilationMode.VLLM_COMPILE,
|
||||
"dynamic_shapes_config": {
|
||||
"type": DynamicShapesType.BACKED.value,
|
||||
},
|
||||
},
|
||||
)
|
||||
|
||||
sampling_params = SamplingParams(temperature=0, top_p=0.95, max_tokens=10)
|
||||
|
||||
# Generate with static shape inputs
|
||||
output = llm.generate("Hello, my name is", sampling_params=sampling_params)
|
||||
result = output[0].outputs[0].text
|
||||
assert len(result) > 0, "Should generate non-empty output"
|
||||
|
||||
# Generate again to verify compilation works with empty sym_shape_indices
|
||||
output = llm.generate("The capital of France is", sampling_params=sampling_params)
|
||||
result = output[0].outputs[0].text
|
||||
assert len(result) > 0, "Should generate non-empty output on second run"
|
||||
|
||||
del llm
|
||||
gc.collect()
|
||||
torch.accelerator.empty_cache()
|
||||
torch.accelerator.synchronize()
|
||||
|
||||
@@ -0,0 +1,150 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Unit tests for tool_calls Iterable → list materialisation.
|
||||
|
||||
Regression tests for https://github.com/vllm-project/vllm/issues/34792.
|
||||
|
||||
Setting VLLM_LOGGING_LEVEL=debug caused tool calling to break for Mistral
|
||||
models because:
|
||||
1. The OpenAI Python SDK types tool_calls as Iterable[...] in
|
||||
ChatCompletionAssistantMessageParam.
|
||||
2. Pydantic v2, when validating from Python objects (not from raw JSON),
|
||||
wraps Iterable fields in a one-shot lazy iterator.
|
||||
3. Debug logging called model_dump_json() which consumed that iterator.
|
||||
4. The Mistral tokenizer then saw empty tool_calls and raised
|
||||
"ValueError: Unexpected tool call id ...".
|
||||
"""
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest
|
||||
|
||||
|
||||
def _make_tool_call(tc_id: str, name: str, args: str) -> dict:
|
||||
return {
|
||||
"id": tc_id,
|
||||
"type": "function",
|
||||
"function": {"name": name, "arguments": args},
|
||||
}
|
||||
|
||||
|
||||
def _make_request(messages: list) -> ChatCompletionRequest:
|
||||
return ChatCompletionRequest(
|
||||
model="test-model",
|
||||
messages=messages,
|
||||
)
|
||||
|
||||
|
||||
def test_tool_calls_list_preserved_after_model_dump():
|
||||
"""tool_calls in assistant messages must be readable after model_dump_json.
|
||||
|
||||
When the request is built from Python dicts (as in the Anthropic → OpenAI
|
||||
conversion path), Pydantic v2 previously wrapped the Iterable tool_calls
|
||||
in a one-shot iterator. model_dump_json() consumed it, leaving subsequent
|
||||
readers (e.g. the Mistral tokenizer) with an empty sequence.
|
||||
"""
|
||||
tool_call = _make_tool_call("call_abc123", "get_weather", '{"city": "Paris"}')
|
||||
messages = [
|
||||
{"role": "user", "content": "What is the weather in Paris?"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
|
||||
{
|
||||
"role": "tool",
|
||||
"tool_call_id": "call_abc123",
|
||||
"content": '{"temperature": 20}',
|
||||
},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
|
||||
# Simulate debug logging: serialize the model (this was the trigger)
|
||||
_ = req.model_dump_json()
|
||||
|
||||
# The assistant message must still have accessible tool_calls afterwards
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
tool_calls = assistant_msg.get("tool_calls")
|
||||
assert tool_calls is not None, "tool_calls must not be None after model_dump_json"
|
||||
assert isinstance(tool_calls, list), "tool_calls must be a list"
|
||||
assert len(tool_calls) > 0, "tool_calls must not be empty after model_dump_json"
|
||||
|
||||
|
||||
def test_tool_calls_from_generator_are_materialised():
|
||||
"""tool_calls passed as a generator must be converted to list on validation."""
|
||||
tool_call = _make_tool_call("call_gen1", "search", '{"query": "vllm"}')
|
||||
|
||||
def tool_calls_gen():
|
||||
yield tool_call
|
||||
|
||||
messages = [
|
||||
{"role": "user", "content": "Search for vllm"},
|
||||
{
|
||||
"role": "assistant",
|
||||
"content": None,
|
||||
"tool_calls": tool_calls_gen(), # one-shot generator
|
||||
},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
|
||||
# Iterate twice — must not raise or return empty on second pass
|
||||
tool_calls_first = list(assistant_msg.get("tool_calls", []))
|
||||
tool_calls_second = list(assistant_msg.get("tool_calls", []))
|
||||
|
||||
assert len(tool_calls_first) == 1, "First read must return the tool call"
|
||||
assert len(tool_calls_second) == 1, "Second read must also return the tool call"
|
||||
|
||||
|
||||
def test_tool_calls_list_passthrough():
|
||||
"""tool_calls already provided as a list must remain a list."""
|
||||
tool_call = _make_tool_call("call_list1", "calculate", '{"expr": "2+2"}')
|
||||
messages = [
|
||||
{"role": "user", "content": "Calculate 2+2"},
|
||||
{"role": "assistant", "content": None, "tool_calls": [tool_call]},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
assert isinstance(assistant_msg.get("tool_calls"), list)
|
||||
|
||||
|
||||
def test_messages_without_tool_calls_unaffected():
|
||||
"""Messages without tool_calls must be handled correctly."""
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a helpful assistant."},
|
||||
{"role": "user", "content": "Hello!"},
|
||||
{"role": "assistant", "content": "Hi there!"},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
# None of the messages should have tool_calls injected
|
||||
for msg in req.messages:
|
||||
assert isinstance(msg, dict)
|
||||
assert msg.get("tool_calls") is None or msg.get("tool_calls") == []
|
||||
|
||||
|
||||
@pytest.mark.parametrize("num_tool_calls", [1, 3])
|
||||
def test_multiple_tool_calls_materialised(num_tool_calls: int):
|
||||
"""Multiple tool calls in a single message are all preserved."""
|
||||
tool_calls = [
|
||||
_make_tool_call(f"call_{i}", f"func_{i}", f'{{"arg": {i}}}')
|
||||
for i in range(num_tool_calls)
|
||||
]
|
||||
messages = [
|
||||
{"role": "user", "content": "Do things"},
|
||||
{"role": "assistant", "content": None, "tool_calls": iter(tool_calls)},
|
||||
]
|
||||
|
||||
req = _make_request(messages)
|
||||
assistant_msg = req.messages[1]
|
||||
assert isinstance(assistant_msg, dict)
|
||||
|
||||
result_tool_calls = assistant_msg.get("tool_calls")
|
||||
assert isinstance(result_tool_calls, list)
|
||||
assert len(result_tool_calls) == num_tool_calls
|
||||
|
||||
# Verify after model_dump_json too
|
||||
_ = req.model_dump_json()
|
||||
assert len(assistant_msg.get("tool_calls", [])) == num_tool_calls
|
||||
@@ -1038,6 +1038,50 @@ _MULTIMODAL_EXAMPLE_MODELS = {
|
||||
},
|
||||
trust_remote_code=True,
|
||||
),
|
||||
# NemotronH_Nano_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2
|
||||
# Use the same registry test as NemotronH_Nano_VL_V2 above
|
||||
"NemotronH_Nano_Omni_Reasoning_V3": _HfExamplesInfo(
|
||||
"nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16",
|
||||
max_model_len=4096,
|
||||
use_original_num_layers=True,
|
||||
hf_overrides={
|
||||
"vision_config": PretrainedConfig(
|
||||
args={
|
||||
"min_num_patches": 1,
|
||||
"max_num_patches": 12,
|
||||
"model": "vit_huge_patch16_224",
|
||||
},
|
||||
video_temporal_patch_size=2,
|
||||
),
|
||||
"text_config": {
|
||||
"num_hidden_layers": 2,
|
||||
"hybrid_override_pattern": "M*",
|
||||
},
|
||||
},
|
||||
trust_remote_code=True,
|
||||
),
|
||||
# NemotronH_Super_Omni_Reasoning_V3 is an alias for NemotronH_Nano_VL_V2 as well
|
||||
# Use the same registry test as NemotronH_Nano_VL_V2 above
|
||||
"NemotronH_Super_Omni_Reasoning_V3": _HfExamplesInfo(
|
||||
"nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16",
|
||||
max_model_len=4096,
|
||||
use_original_num_layers=True,
|
||||
hf_overrides={
|
||||
"vision_config": PretrainedConfig(
|
||||
args={
|
||||
"min_num_patches": 1,
|
||||
"max_num_patches": 12,
|
||||
"model": "vit_huge_patch16_224",
|
||||
},
|
||||
video_temporal_patch_size=2,
|
||||
),
|
||||
"text_config": {
|
||||
"num_hidden_layers": 2,
|
||||
"hybrid_override_pattern": "M*",
|
||||
},
|
||||
},
|
||||
trust_remote_code=True,
|
||||
),
|
||||
"OpenCUAForConditionalGeneration": _HfExamplesInfo(
|
||||
"xlangai/OpenCUA-7B", trust_remote_code=True
|
||||
),
|
||||
|
||||
@@ -12,6 +12,7 @@ MODELS = [
|
||||
"TheBloke/TinyLlama-1.1B-Chat-v1.0-GPTQ", # with g_idx
|
||||
"Qwen/Qwen1.5-0.5B-Chat-GPTQ-Int4", # without g_idx
|
||||
"RedHatAI/Qwen3-1.7B-quantized.w4a16", # with zp
|
||||
"OPEA/Qwen2.5-0.5B-Instruct-int4-sym-inc",
|
||||
]
|
||||
DTYPE = ["bfloat16"]
|
||||
|
||||
|
||||
+8
-10
@@ -1261,9 +1261,6 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non
|
||||
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args: _P.args, **kwargs: _P.kwargs) -> None:
|
||||
# Make the process the leader of its own process group
|
||||
# to avoid sending SIGTERM to the parent process
|
||||
os.setpgrp()
|
||||
from _pytest.outcomes import Skipped
|
||||
|
||||
# Create a unique temporary file to store exception info from child
|
||||
@@ -1283,6 +1280,9 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non
|
||||
pid = os.fork()
|
||||
print(f"Fork a new process to run a test {pid}")
|
||||
if pid == 0:
|
||||
# Make the child process the leader of its own process group
|
||||
# to avoid sending SIGTERM to the parent process
|
||||
os.setpgrp()
|
||||
# Parent process responsible for deleting, don't delete
|
||||
# in child.
|
||||
delete_after.pop_all()
|
||||
@@ -1322,14 +1322,12 @@ def fork_new_process_for_each_test(func: Callable[_P, None]) -> Callable[_P, Non
|
||||
else:
|
||||
os._exit(0)
|
||||
else:
|
||||
pgid = os.getpgid(pid)
|
||||
# After setpgrp(), the child's pgid equals its pid
|
||||
pgid = pid
|
||||
_pid, _exitcode = os.waitpid(pid, 0)
|
||||
# ignore SIGTERM signal itself
|
||||
old_signal_handler = signal.signal(signal.SIGTERM, signal.SIG_IGN)
|
||||
# kill all child processes
|
||||
os.killpg(pgid, signal.SIGTERM)
|
||||
# restore the signal handler
|
||||
signal.signal(signal.SIGTERM, old_signal_handler)
|
||||
# kill all child processes - but they may already have exited cleanly
|
||||
with contextlib.suppress(ProcessLookupError):
|
||||
os.killpg(pgid, signal.SIGTERM)
|
||||
if _exitcode != 0:
|
||||
# Try to read the exception from the child process
|
||||
exc_info = {}
|
||||
|
||||
@@ -36,10 +36,10 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
using the high-level v1 LLM() API only (no manual batching).
|
||||
|
||||
Strategy:
|
||||
- Create two LLM engines with identical config except max_num_seqs: 1 vs N.
|
||||
- Compute a baseline output for the needle prompt with the bs=1 engine.
|
||||
- For many trials, generate a batch (size N) where the needle appears at a
|
||||
random position among random filler prompts using the bs=N engine.
|
||||
- Create a single LLM engine configured for the larger batch limit (N).
|
||||
- Compute a baseline output for the needle prompt when it is run alone.
|
||||
- For many trials, generate a mixed batch (size N) where the needle appears
|
||||
at a random position among random filler prompts using the same engine.
|
||||
- Track how many trials match vs mismatch, and report totals at the end.
|
||||
The test fails if any mismatches occur, but we still dump pass/fail
|
||||
counts.
|
||||
@@ -83,11 +83,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
|
||||
needle_prompt = "There once was a "
|
||||
|
||||
llm_bs1 = None
|
||||
llm_bsN = None
|
||||
llm = None
|
||||
try:
|
||||
# Engine with bs=1 behavior
|
||||
llm_bs1 = LLM_with_max_seqs(
|
||||
llm = LLM_with_max_seqs(
|
||||
model=model,
|
||||
max_num_seqs=max_batch_size,
|
||||
gpu_memory_utilization=gpu_mem_util,
|
||||
@@ -96,20 +94,11 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
)
|
||||
|
||||
# Baseline generation for the needle prompt alone.
|
||||
baseline_out = llm_bs1.generate([needle_prompt], sampling)
|
||||
baseline_out = llm.generate([needle_prompt], sampling)
|
||||
assert len(baseline_out) == 1
|
||||
assert len(baseline_out[0].outputs) >= 1
|
||||
baseline_text = baseline_out[0].outputs[0].text
|
||||
|
||||
# Engine with larger batch limit (e.g., 64)
|
||||
llm_bsN = LLM_with_max_seqs(
|
||||
model=model,
|
||||
max_num_seqs=max_batch_size,
|
||||
gpu_memory_utilization=gpu_mem_util,
|
||||
max_model_len=max_model_len,
|
||||
attention_config=attention_config,
|
||||
)
|
||||
|
||||
mismatches = 0
|
||||
|
||||
for trial in range(num_trials):
|
||||
@@ -124,8 +113,8 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
else:
|
||||
prompts.append(_random_prompt(min_random_prompt, max_random_prompt))
|
||||
|
||||
# Generate with the larger-batch engine
|
||||
outputs = llm_bsN.generate(prompts, sampling)
|
||||
# Generate with the same engine but in a larger batch.
|
||||
outputs = llm.generate(prompts, sampling)
|
||||
# Find the needle output by position
|
||||
needle_output = outputs[needle_pos]
|
||||
assert needle_output.prompt == needle_prompt
|
||||
@@ -151,12 +140,9 @@ def test_v1_generation_is_deterministic_across_batch_sizes_with_needle(
|
||||
|
||||
finally:
|
||||
# Ensure engines are shutdown to free GPU/VRAM across test sessions
|
||||
if llm_bs1 is not None:
|
||||
if llm is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
llm_bs1.shutdown()
|
||||
if llm_bsN is not None:
|
||||
with contextlib.suppress(Exception):
|
||||
llm_bsN.shutdown()
|
||||
llm.shutdown()
|
||||
|
||||
|
||||
@skip_unsupported
|
||||
|
||||
@@ -124,12 +124,8 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
|
||||
return [BlockHash(str(i).encode()) for i in int_hashes]
|
||||
|
||||
def take_events() -> Iterable[OffloadingEvent]:
|
||||
yield OffloadingEvent(
|
||||
keys=to_keys([1, 2, 3]), block_size=16, medium="A", removed=False
|
||||
)
|
||||
yield OffloadingEvent(
|
||||
keys=to_keys([4, 5, 6]), block_size=32, medium="B", removed=True
|
||||
)
|
||||
yield OffloadingEvent(keys=to_keys([1, 2, 3]), medium="A", removed=False)
|
||||
yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True)
|
||||
|
||||
runner.manager.take_events.side_effect = take_events
|
||||
events = list(runner.scheduler_connector.take_events())
|
||||
@@ -137,7 +133,7 @@ def test_offloading_connector(request_runner, async_scheduling: bool):
|
||||
event = events[0]
|
||||
assert isinstance(event, BlockStored)
|
||||
assert event.block_hashes == to_hashes([1, 2, 3])
|
||||
assert event.block_size == 16
|
||||
assert event.block_size == 0
|
||||
assert event.medium == "A"
|
||||
assert event.token_ids == []
|
||||
assert event.parent_block_hash is None
|
||||
|
||||
@@ -609,6 +609,51 @@ def test_register_kv_caches():
|
||||
assert bl == tensor1[0].nbytes // tensor1.shape[1]
|
||||
|
||||
|
||||
def test_register_kv_caches_supports_mixed_mla_and_eagle_shapes():
|
||||
"""Mixed MLA+Eagle caches should register by byte length, not shape."""
|
||||
|
||||
vllm_config = create_vllm_config(
|
||||
kv_connector="MooncakeConnector", kv_role="kv_consumer"
|
||||
)
|
||||
|
||||
with (
|
||||
set_current_vllm_config(vllm_config),
|
||||
patch_worker_dependencies(),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Event"
|
||||
),
|
||||
patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake.mooncake_connector.threading.Thread"
|
||||
) as mock_thread,
|
||||
):
|
||||
connector = MooncakeConnector(vllm_config, KVConnectorRole.WORKER)
|
||||
worker = connector.connector_worker
|
||||
mock_thread.return_value.is_alive.return_value = False
|
||||
|
||||
worker.use_mla = True
|
||||
worker.kv_topo.is_mla = True
|
||||
|
||||
# MLA cache tensor: shape[-2] is the block size.
|
||||
mla_cache = torch.zeros((2, 16, 96), dtype=torch.float16)
|
||||
# Eagle3/GQA-like cache tensor: shape[-2] is num_kv_heads, not block size.
|
||||
eagle_cache = torch.zeros((2, 16, 8, 64), dtype=torch.float16)
|
||||
kv_caches = {"mla_layer": mla_cache, "eagle_layer": eagle_cache}
|
||||
|
||||
with patch.object(
|
||||
worker.engine, "batch_register_memory", return_value=0
|
||||
) as mock_batch_register:
|
||||
connector.register_kv_caches(kv_caches)
|
||||
|
||||
mock_batch_register.assert_called_once()
|
||||
registered_ptrs, registered_lens = mock_batch_register.call_args[0]
|
||||
assert registered_ptrs == [mla_cache.data_ptr(), eagle_cache.data_ptr()]
|
||||
assert registered_lens == [mla_cache.nbytes, eagle_cache.nbytes]
|
||||
assert worker.block_len_per_layer == [
|
||||
mla_cache.nbytes // mla_cache.shape[0],
|
||||
eagle_cache.nbytes // eagle_cache.shape[0],
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.asyncio
|
||||
@patch(
|
||||
"vllm.distributed.kv_transfer.kv_connector.v1.mooncake."
|
||||
|
||||
@@ -89,6 +89,99 @@ def test_logical_to_kernel_block_ids_with_hma():
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.cpu_test
|
||||
@pytest.mark.parametrize(
|
||||
"has_mamba,swa_enabled,mamba_enabled,remote_ratio,"
|
||||
"remote_block_ids,expected_remote_block_ids",
|
||||
[
|
||||
# Non-mamba (FA+SWA): both groups expanded via _logical_to_kernel_block_ids.
|
||||
# Regression for https://github.com/vllm-project/vllm/pull/39724
|
||||
(
|
||||
False,
|
||||
True,
|
||||
False,
|
||||
1,
|
||||
([0, 1, 2], [3, 4]),
|
||||
[[0, 1, 2, 3, 4, 5], [6, 7, 8, 9]],
|
||||
),
|
||||
# Mamba (FA+Mamba): FA expanded via _logical_to_remote_kernel_block_ids,
|
||||
# Mamba passed through unchanged.
|
||||
# remote_ratio=261 (Nemotron 30B TP=1) != local_ratio=2 so that using
|
||||
# the wrong conversion method produces different FA results.
|
||||
(
|
||||
True,
|
||||
False,
|
||||
True,
|
||||
261,
|
||||
([0, 1, 2], [10, 11]),
|
||||
[[0, 1, 261, 262, 522, 523], [10, 11]],
|
||||
),
|
||||
],
|
||||
ids=["non_mamba_fa_swa", "mamba_fa_ssm"],
|
||||
)
|
||||
def test_read_blocks_for_req_expands_remote_ids(
|
||||
has_mamba,
|
||||
swa_enabled,
|
||||
mamba_enabled,
|
||||
remote_ratio,
|
||||
remote_block_ids,
|
||||
expected_remote_block_ids,
|
||||
):
|
||||
"""_read_blocks_for_req must expand remote logical block IDs to kernel
|
||||
block IDs when kernel block size != logical block size.
|
||||
|
||||
Non-mamba path uses _logical_to_kernel_block_ids (all groups expanded).
|
||||
Mamba path uses _logical_to_remote_kernel_block_ids (FA expanded, Mamba
|
||||
passed through).
|
||||
"""
|
||||
from unittest.mock import MagicMock
|
||||
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.metadata import (
|
||||
NixlConnectorMetadata,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.nixl.worker import (
|
||||
NixlConnectorWorker,
|
||||
)
|
||||
|
||||
worker = object.__new__(NixlConnectorWorker)
|
||||
worker._has_mamba = has_mamba
|
||||
worker._physical_blocks_per_logical_kv_block = 2
|
||||
worker.kv_cache_config = make_kv_cache_config(
|
||||
block_size=16, swa_enabled=swa_enabled, mamba_enabled=mamba_enabled
|
||||
)
|
||||
|
||||
remote_engine_id = "remote-engine"
|
||||
if has_mamba:
|
||||
worker._mamba_phys_ratio = {remote_engine_id: remote_ratio}
|
||||
|
||||
# Mock kv_topo: empty remote ranks skips the transfer machinery entirely,
|
||||
# isolating the block-ID expansion logic.
|
||||
worker.kv_topo = MagicMock()
|
||||
worker.kv_topo.get_target_remote_ranks_from_engine_id.return_value = []
|
||||
worker.kv_topo.tp_ratio_from_engine_id.return_value = 1
|
||||
|
||||
metadata = NixlConnectorMetadata()
|
||||
metadata.add_new_req_to_recv(
|
||||
request_id="test-req",
|
||||
local_block_ids=([0, 1], [2, 3]),
|
||||
kv_transfer_params={
|
||||
"remote_block_ids": remote_block_ids,
|
||||
"remote_engine_id": remote_engine_id,
|
||||
"remote_request_id": "prefill-test-req",
|
||||
"remote_host": "localhost",
|
||||
"remote_port": 1234,
|
||||
"tp_size": 1,
|
||||
},
|
||||
)
|
||||
|
||||
meta = metadata.reqs_to_recv["test-req"]
|
||||
worker._read_blocks_for_req("test-req", meta)
|
||||
|
||||
assert meta.remote.block_ids == expected_remote_block_ids, (
|
||||
f"Expected {expected_remote_block_ids}, got {meta.remote.block_ids}"
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("model_name, sw_size", [("google/gemma-3-1b-it", 512)])
|
||||
def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
|
||||
"""Test that a prefill instance returns fewer "remote blocks" for the SWA groups
|
||||
@@ -102,7 +195,7 @@ def test_fewer_blocks_with_hma(monkeypatch, model_name, sw_size):
|
||||
llm_kwargs = {
|
||||
"model": model_name,
|
||||
"enforce_eager": True,
|
||||
"gpu_memory_utilization": 0.5,
|
||||
"gpu_memory_utilization": 0.47,
|
||||
"kv_transfer_config": kv_transfer_config,
|
||||
"max_model_len": 2048,
|
||||
# NOTE: Make sure HMA is enabled
|
||||
|
||||
@@ -59,7 +59,6 @@ def verify_load_output(
|
||||
|
||||
def verify_events(
|
||||
events: Iterable[OffloadingEvent],
|
||||
block_size: int,
|
||||
expected_stores: tuple[set[int], ...] = (),
|
||||
expected_evictions: tuple[set[int], ...] = (),
|
||||
):
|
||||
@@ -67,7 +66,6 @@ def verify_events(
|
||||
evictions: list[set[OffloadKey]] = []
|
||||
for event in events:
|
||||
assert event.medium == CPULoadStoreSpec.medium()
|
||||
assert event.block_size == block_size
|
||||
if event.removed:
|
||||
evictions.append(set(event.keys))
|
||||
else:
|
||||
@@ -98,9 +96,7 @@ def test_already_stored_block_not_evicted_during_prepare_store(eviction_policy):
|
||||
candidate to make room for [3, 4, 5]
|
||||
- After complete_store([2, 3, 4, 5]), block 2 must still be present.
|
||||
"""
|
||||
block_size = 256
|
||||
manager = CPUOffloadingManager(
|
||||
block_size=block_size,
|
||||
num_blocks=4,
|
||||
cache_policy=eviction_policy,
|
||||
enable_events=True,
|
||||
@@ -138,10 +134,9 @@ def test_cpu_manager():
|
||||
"""
|
||||
Tests CPUOffloadingManager with lru policy.
|
||||
"""
|
||||
# initialize a CPU backend with a capacity of 4 blocks
|
||||
block_size = 256
|
||||
# initialize a CPU manager with a capacity of 4 blocks
|
||||
cpu_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
|
||||
num_blocks=4, cache_policy="lru", enable_events=True
|
||||
)
|
||||
|
||||
# prepare store [1, 2]
|
||||
@@ -163,9 +158,7 @@ def test_cpu_manager():
|
||||
|
||||
# complete store [1, 2]
|
||||
cpu_manager.complete_store(to_keys([1, 2]))
|
||||
verify_events(
|
||||
cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
|
||||
)
|
||||
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
|
||||
|
||||
# lookup [1, 2]
|
||||
assert cpu_manager.lookup(to_keys([1])) == 1
|
||||
@@ -184,9 +177,7 @@ def test_cpu_manager():
|
||||
)
|
||||
|
||||
# verify eviction event
|
||||
verify_events(
|
||||
cpu_manager.take_events(), block_size=block_size, expected_evictions=({1},)
|
||||
)
|
||||
verify_events(cpu_manager.take_events(), expected_evictions=({1},))
|
||||
|
||||
# prepare store with no space
|
||||
assert cpu_manager.prepare_store(to_keys([1, 6])) is None
|
||||
@@ -241,7 +232,6 @@ def test_cpu_manager():
|
||||
|
||||
verify_events(
|
||||
cpu_manager.take_events(),
|
||||
block_size=block_size,
|
||||
expected_stores=({3, 4, 5}, {6, 7, 8}),
|
||||
expected_evictions=({2, 3, 4}, {8}),
|
||||
)
|
||||
@@ -254,7 +244,6 @@ class TestARCPolicy:
|
||||
self, num_blocks: int = 4, enable_events: bool = True
|
||||
) -> tuple[CPUOffloadingManager, ARCCachePolicy]:
|
||||
manager = CPUOffloadingManager(
|
||||
block_size=256,
|
||||
num_blocks=num_blocks,
|
||||
cache_policy="arc",
|
||||
enable_events=enable_events,
|
||||
@@ -289,9 +278,7 @@ class TestARCPolicy:
|
||||
|
||||
# complete store [1, 2]
|
||||
cpu_manager.complete_store(to_keys([1, 2]))
|
||||
verify_events(
|
||||
cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},)
|
||||
)
|
||||
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
|
||||
|
||||
# lookup [1, 2]
|
||||
assert cpu_manager.lookup(to_keys([1])) == 1
|
||||
@@ -547,9 +534,8 @@ def test_filter_reused_manager():
|
||||
"""
|
||||
Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
|
||||
"""
|
||||
block_size = 256
|
||||
lru_manager = CPUOffloadingManager(
|
||||
block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
|
||||
num_blocks=4, cache_policy="lru", enable_events=True
|
||||
)
|
||||
|
||||
manager = FilterReusedOffloadingManager(
|
||||
|
||||
@@ -26,6 +26,7 @@ def test_prefill_kv_computed_with_cache():
|
||||
# Case 1: With prefix cache (1200 tokens cached)
|
||||
iteration_stats.update_from_finished_request(
|
||||
finish_reason=FinishReason.STOP,
|
||||
request_id="test-req-001",
|
||||
num_prompt_tokens=10000,
|
||||
max_tokens_param=100,
|
||||
req_stats=req_stats,
|
||||
@@ -35,6 +36,7 @@ def test_prefill_kv_computed_with_cache():
|
||||
finished_req = iteration_stats.finished_requests[0]
|
||||
assert finished_req.num_prompt_tokens == 10000
|
||||
assert finished_req.num_cached_tokens == 1200
|
||||
assert finished_req.request_id == "test-req-001"
|
||||
|
||||
# Verify calculation: prefill KV = prompt tokens - cached tokens
|
||||
prefill_kv_computed = finished_req.num_prompt_tokens - max(
|
||||
@@ -55,6 +57,7 @@ def test_prefill_kv_computed_no_cache():
|
||||
# Case 2: No prefix cache
|
||||
iteration_stats.update_from_finished_request(
|
||||
finish_reason=FinishReason.STOP,
|
||||
request_id="test-req-002",
|
||||
num_prompt_tokens=2000,
|
||||
max_tokens_param=100,
|
||||
req_stats=req_stats,
|
||||
@@ -64,6 +67,7 @@ def test_prefill_kv_computed_no_cache():
|
||||
finished_req = iteration_stats.finished_requests[0]
|
||||
assert finished_req.num_prompt_tokens == 2000
|
||||
assert finished_req.num_cached_tokens == 0
|
||||
assert finished_req.request_id == "test-req-002"
|
||||
|
||||
# Verify calculation: prefill KV = full prompt when no cache
|
||||
prefill_kv_computed = finished_req.num_prompt_tokens - max(
|
||||
@@ -84,6 +88,7 @@ def test_prefill_kv_computed_edge_cases():
|
||||
# Case 3: Negative num_cached_tokens (shouldn't happen, but handle gracefully)
|
||||
iteration_stats.update_from_finished_request(
|
||||
finish_reason=FinishReason.STOP,
|
||||
request_id="test-req-003",
|
||||
num_prompt_tokens=100,
|
||||
max_tokens_param=10,
|
||||
req_stats=req_stats,
|
||||
@@ -96,11 +101,13 @@ def test_prefill_kv_computed_edge_cases():
|
||||
finished_req.num_cached_tokens, 0
|
||||
)
|
||||
assert prefill_kv_computed == 100 # Should treat negative as 0
|
||||
assert finished_req.request_id == "test-req-003"
|
||||
|
||||
# Case 4: All tokens cached (shouldn't happen in practice)
|
||||
iteration_stats2 = IterationStats()
|
||||
iteration_stats2.update_from_finished_request(
|
||||
finish_reason=FinishReason.STOP,
|
||||
request_id="test-req-004",
|
||||
num_prompt_tokens=100,
|
||||
max_tokens_param=10,
|
||||
req_stats=req_stats,
|
||||
@@ -112,6 +119,7 @@ def test_prefill_kv_computed_edge_cases():
|
||||
finished_req2.num_cached_tokens, 0
|
||||
)
|
||||
assert prefill_kv_computed2 == 0 # All cached, nothing computed
|
||||
assert finished_req2.request_id == "test-req-004"
|
||||
|
||||
|
||||
def test_prompt_token_stats_all_computed():
|
||||
|
||||
@@ -0,0 +1,171 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline
|
||||
from tests.utils import large_gpu_mark
|
||||
from vllm import LLM
|
||||
from vllm.config import SpeculativeConfig
|
||||
from vllm.distributed import cleanup_dist_env_and_memory
|
||||
|
||||
MODEL_PATH = "nm-testing/dflash-qwen3-8b-speculators"
|
||||
|
||||
EXPECTED_GSM8K_ACCURACY = 0.885
|
||||
ACCURACY_RTOL = 0.03
|
||||
EXPECTED_ACCEPTANCE_LEN = 3.45
|
||||
ACCEPTANCE_LEN_RTOL = 0.15
|
||||
|
||||
# Expected per-position acceptance rates (accepted_at_pos / num_drafts)
|
||||
# Based on GSM8K evaluation with Qwen3-8B dflash speculators.
|
||||
EXPECTED_PER_POS_ACCEPTANCE_RATES = [0.795, 0.611, 0.429, 0.282]
|
||||
PER_POS_RTOL = 0.15
|
||||
|
||||
|
||||
def compute_spec_decode_stats(
|
||||
metrics,
|
||||
) -> dict:
|
||||
"""Extract all spec-decode metrics and compute derived stats."""
|
||||
name2metric = {m.name: m for m in metrics}
|
||||
|
||||
n_drafts = name2metric["vllm:spec_decode_num_drafts"].value
|
||||
n_draft_tokens = name2metric["vllm:spec_decode_num_draft_tokens"].value
|
||||
n_accepted = name2metric["vllm:spec_decode_num_accepted_tokens"].value
|
||||
|
||||
per_pos_vec = name2metric["vllm:spec_decode_num_accepted_tokens_per_pos"].values
|
||||
|
||||
acceptance_len = 1 + (n_accepted / n_drafts) if n_drafts > 0 else 1.0
|
||||
draft_tokens_per_step = (n_draft_tokens / n_drafts) if n_drafts > 0 else 0
|
||||
overall_acceptance_rate = (n_accepted / n_draft_tokens) if n_draft_tokens > 0 else 0
|
||||
per_pos_rates = [v / n_drafts for v in per_pos_vec] if n_drafts > 0 else []
|
||||
|
||||
return {
|
||||
"num_drafts": n_drafts,
|
||||
"num_draft_tokens": n_draft_tokens,
|
||||
"num_accepted_tokens": n_accepted,
|
||||
"acceptance_len": acceptance_len,
|
||||
"draft_tokens_per_step": draft_tokens_per_step,
|
||||
"overall_acceptance_rate": overall_acceptance_rate,
|
||||
"per_pos_accepted": list(per_pos_vec),
|
||||
"per_pos_acceptance_rates": per_pos_rates,
|
||||
}
|
||||
|
||||
|
||||
def print_spec_decode_stats(stats: dict) -> None:
|
||||
"""Print all spec-decode metrics and derived values."""
|
||||
print("\n===== Spec Decode Metrics =====")
|
||||
print(f" num_drafts: {stats['num_drafts']}")
|
||||
print(f" num_draft_tokens: {stats['num_draft_tokens']}")
|
||||
print(f" num_accepted_tokens: {stats['num_accepted_tokens']}")
|
||||
print(f" draft_tokens_per_step: {stats['draft_tokens_per_step']:.2f}")
|
||||
print(f" overall_acceptance_rate: {stats['overall_acceptance_rate']:.4f}")
|
||||
print(f" acceptance_len (1+acc/drafts): {stats['acceptance_len']:.4f}")
|
||||
print(" per-position accepted tokens:", stats["per_pos_accepted"])
|
||||
print(" per-position acceptance rates:")
|
||||
for i, rate in enumerate(stats["per_pos_acceptance_rates"]):
|
||||
print(f" pos {i}: {rate:.4f}")
|
||||
print("===============================\n")
|
||||
|
||||
|
||||
def test_dflash_speculators_model(vllm_runner, example_prompts, monkeypatch):
|
||||
"""
|
||||
Test DFlash speculators model properly initializes speculative decoding.
|
||||
|
||||
Verifies:
|
||||
1. Speculative config is automatically initialized from speculators config
|
||||
2. Method is detected as 'dflash'
|
||||
3. The draft model path is correctly set
|
||||
4. Speculative tokens count is valid (num_speculative_tokens=8)
|
||||
5. Text generation works with speculative decoding enabled
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
with vllm_runner(
|
||||
MODEL_PATH,
|
||||
dtype=torch.bfloat16,
|
||||
enforce_eager=True,
|
||||
quantization="fp8",
|
||||
) as vllm_model:
|
||||
vllm_config = vllm_model.llm.llm_engine.vllm_config
|
||||
|
||||
assert isinstance(vllm_config.speculative_config, SpeculativeConfig), (
|
||||
"Speculative config should be initialized for speculators model"
|
||||
)
|
||||
|
||||
spec_config = vllm_config.speculative_config
|
||||
assert spec_config.method == "dflash", (
|
||||
f"Expected method='dflash', got '{spec_config.method}'"
|
||||
)
|
||||
assert spec_config.num_speculative_tokens > 0, (
|
||||
f"Expected positive speculative tokens, "
|
||||
f"got {spec_config.num_speculative_tokens}"
|
||||
)
|
||||
assert spec_config.model == MODEL_PATH, (
|
||||
f"Draft model should be {MODEL_PATH}, got {spec_config.model}"
|
||||
)
|
||||
|
||||
vllm_outputs = vllm_model.generate_greedy(example_prompts, max_tokens=20)
|
||||
assert vllm_outputs, f"No outputs generated for speculators model {MODEL_PATH}"
|
||||
|
||||
|
||||
@pytest.mark.slow_test
|
||||
@large_gpu_mark(min_gb=40)
|
||||
def test_dflash_speculators_correctness(monkeypatch):
|
||||
"""
|
||||
E2E correctness test for DFlash via the speculators auto-detect path.
|
||||
|
||||
Evaluates GSM8k accuracy to ensure the speculators-format model produces
|
||||
correct outputs, and checks that acceptance length does not collapse under
|
||||
batched inference (lm-eval style).
|
||||
|
||||
Observed per-position acceptance rates on GSM8K (1319 prompts):
|
||||
pos 0: 0.795, pos 1: 0.611, pos 2: 0.429, pos 3: 0.282,
|
||||
pos 4: 0.169, pos 5: 0.093, pos 6: 0.048, pos 7: 0.023
|
||||
Observed mean AL: 3.45 (GSM8K dataset, max_num_seqs=128)
|
||||
"""
|
||||
monkeypatch.setenv("VLLM_ALLOW_INSECURE_SERIALIZATION", "1")
|
||||
|
||||
spec_llm = LLM(
|
||||
model=MODEL_PATH,
|
||||
trust_remote_code=True,
|
||||
max_model_len=4096,
|
||||
max_num_seqs=128,
|
||||
gpu_memory_utilization=0.85,
|
||||
enforce_eager=False,
|
||||
disable_log_stats=False,
|
||||
)
|
||||
|
||||
results = evaluate_gsm8k_offline(spec_llm)
|
||||
accuracy = results["accuracy"]
|
||||
accuracy_threshold = EXPECTED_GSM8K_ACCURACY * (1 - ACCURACY_RTOL)
|
||||
assert accuracy >= accuracy_threshold, (
|
||||
f"Expected GSM8K accuracy >= {accuracy_threshold:.3f}, got {accuracy:.3f}"
|
||||
)
|
||||
|
||||
current_metrics = spec_llm.get_metrics()
|
||||
stats = compute_spec_decode_stats(current_metrics)
|
||||
print_spec_decode_stats(stats)
|
||||
|
||||
acceptance_len = stats["acceptance_len"]
|
||||
al_threshold = EXPECTED_ACCEPTANCE_LEN * (1 - ACCEPTANCE_LEN_RTOL)
|
||||
assert acceptance_len >= al_threshold, (
|
||||
f"DFlash speculators acceptance length too low: "
|
||||
f"{acceptance_len:.2f} < {al_threshold:.2f}"
|
||||
)
|
||||
|
||||
# Check per-position acceptance rates for the first few positions.
|
||||
per_pos_rates = stats["per_pos_acceptance_rates"]
|
||||
for i, expected_rate in enumerate(EXPECTED_PER_POS_ACCEPTANCE_RATES):
|
||||
assert i < len(per_pos_rates), (
|
||||
f"Missing per-position acceptance rate for position {i}"
|
||||
)
|
||||
threshold = expected_rate * (1 - PER_POS_RTOL)
|
||||
assert per_pos_rates[i] >= threshold, (
|
||||
f"Per-position acceptance rate at pos {i} too low: "
|
||||
f"{per_pos_rates[i]:.4f} < {threshold:.4f} "
|
||||
f"(expected ~{expected_rate:.4f})"
|
||||
)
|
||||
|
||||
del spec_llm
|
||||
torch.accelerator.empty_cache()
|
||||
cleanup_dist_env_and_memory()
|
||||
@@ -144,6 +144,46 @@ def _xpu_mxfp8_quantize_fake(
|
||||
return x.to(dtype), x_s.to(torch.float8_e8m0fnu)
|
||||
|
||||
|
||||
def _xpu_mxfp4_quantize_impl(
|
||||
x: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
eps = 1e-10
|
||||
assert x.ndim == 2, "input must be 2-D"
|
||||
assert x.shape[-1] % MXFP4_BLOCK_SIZE == 0, (
|
||||
f"last dimension {x.shape[-1]} must be divisible by group_size "
|
||||
f"{MXFP4_BLOCK_SIZE}"
|
||||
)
|
||||
assert x.is_contiguous(), "input groups must be contiguous"
|
||||
|
||||
M, N = x.shape
|
||||
|
||||
# Packed FP4 output: two nibbles per byte
|
||||
x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
|
||||
x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
|
||||
|
||||
torch.ops._C.per_token_group_quant_mxfp4(x, x_q, x_s, MXFP4_BLOCK_SIZE, eps)
|
||||
|
||||
x_q = x_q.view(torch.float4_e2m1fn_x2)
|
||||
x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
|
||||
return x_q, x_s
|
||||
|
||||
|
||||
def _xpu_mxfp4_quantize_fake(
|
||||
x: torch.Tensor,
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
MXFP4_BLOCK_SIZE = 32
|
||||
M, N = x.shape
|
||||
|
||||
# Packed FP4 output: two nibbles per byte
|
||||
x_q = torch.empty(M, N // 2, device=x.device, dtype=torch.uint8)
|
||||
x_s = torch.empty(M, N // MXFP4_BLOCK_SIZE, device=x.device, dtype=torch.float32)
|
||||
|
||||
x_q = x_q.view(torch.float4_e2m1fn_x2)
|
||||
x_s = x_s.to(dtype=torch.float8_e8m0fnu, memory_format=torch.preserve_format)
|
||||
return x_q, x_s
|
||||
|
||||
|
||||
# Global flag to ensure ops are registered only once
|
||||
_OPS_REGISTERED = False
|
||||
|
||||
@@ -555,6 +595,12 @@ class xpu_ops:
|
||||
fake_impl=_xpu_mxfp8_quantize_fake,
|
||||
)
|
||||
|
||||
direct_register_custom_op(
|
||||
op_name="xpu_mxfp4_quantize",
|
||||
op_func=_xpu_mxfp4_quantize_impl,
|
||||
fake_impl=_xpu_mxfp4_quantize_fake,
|
||||
)
|
||||
|
||||
_OPS_REGISTERED = True
|
||||
|
||||
|
||||
|
||||
@@ -25,6 +25,7 @@ from contextlib import suppress
|
||||
from dataclasses import dataclass, replace
|
||||
from functools import cache
|
||||
from io import BytesIO
|
||||
from pathlib import Path
|
||||
from tempfile import NamedTemporaryFile
|
||||
from typing import Any, cast
|
||||
|
||||
@@ -1422,6 +1423,7 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
|
||||
"custom_mm",
|
||||
"prefix_repetition",
|
||||
"spec_bench",
|
||||
"speed_bench",
|
||||
],
|
||||
help="Name of the dataset to benchmark on.",
|
||||
)
|
||||
@@ -1606,6 +1608,34 @@ def add_dataset_parser(parser: FlexibleArgumentParser):
|
||||
"repetition dataset.",
|
||||
)
|
||||
|
||||
speed_bench_group = parser.add_argument_group("speed bench dataset options")
|
||||
speed_bench_group.add_argument(
|
||||
"--speed-bench-dataset-subset",
|
||||
type=str,
|
||||
default="qualitative",
|
||||
choices={
|
||||
"qualitative",
|
||||
"throughput_1k",
|
||||
"throughput_2k",
|
||||
"throughput_8k",
|
||||
"throughput_16k",
|
||||
"throughput_32k",
|
||||
},
|
||||
help="Subset of the SPEED-Bench dataset.",
|
||||
)
|
||||
speed_bench_group.add_argument(
|
||||
"--speed-bench-output-len",
|
||||
type=int,
|
||||
default=4096,
|
||||
help="Num of output tokens per request, used only for speed bench dataset.",
|
||||
)
|
||||
speed_bench_group.add_argument(
|
||||
"--speed-bench-category",
|
||||
type=str,
|
||||
default=None,
|
||||
help="Category for speed bench dataset. If None, use all categories.",
|
||||
)
|
||||
|
||||
|
||||
def add_random_dataset_base_args(
|
||||
parser_or_group: FlexibleArgumentParser | argparse._ArgumentGroup,
|
||||
@@ -2074,6 +2104,19 @@ def get_samples(args, tokenizer: TokenizerLike) -> list[SampleRequest]:
|
||||
request_id_prefix=args.request_id_prefix,
|
||||
no_oversample=args.no_oversample,
|
||||
),
|
||||
"speed_bench": lambda: SpeedBench(
|
||||
dataset_path=args.dataset_path,
|
||||
dataset_subset=args.speed_bench_dataset_subset,
|
||||
category=args.speed_bench_category,
|
||||
disable_shuffle=args.disable_shuffle,
|
||||
).sample(
|
||||
num_requests=args.num_prompts,
|
||||
tokenizer=tokenizer,
|
||||
output_len=args.speed_bench_output_len,
|
||||
enable_multimodal_chat=args.enable_multimodal_chat,
|
||||
request_id_prefix=args.request_id_prefix,
|
||||
no_oversample=args.no_oversample,
|
||||
),
|
||||
}
|
||||
|
||||
try:
|
||||
@@ -3551,3 +3594,48 @@ class MMStarDataset(HuggingFaceDataset):
|
||||
sampled_requests, num_requests, request_id_prefix, no_oversample
|
||||
)
|
||||
return sampled_requests
|
||||
|
||||
|
||||
# -----------------------------------------------------------------------------
|
||||
# Speed Bench Dataset Implementation
|
||||
# -----------------------------------------------------------------------------
|
||||
|
||||
|
||||
class SpeedBench(CustomDataset):
|
||||
"""
|
||||
Implements the SPEED-Bench dataset: https://huggingface.co/datasets/nvidia/SPEED-Bench
|
||||
Download the dataset using:
|
||||
curl -LsSf https://raw.githubusercontent.com/NVIDIA-NeMo/Skills/refs/heads/main/nemo_skills/dataset/speed-bench/prepare.py | python3 -
|
||||
""" # noqa: E501
|
||||
|
||||
def __init__(self, **kwargs) -> None:
|
||||
self.dataset_subset = kwargs.pop("dataset_subset", "qualitative")
|
||||
self.category = kwargs.pop("category", None)
|
||||
super().__init__(**kwargs)
|
||||
self.load_data()
|
||||
|
||||
def load_data(self) -> None:
|
||||
if self.dataset_path is None:
|
||||
raise ValueError("dataset_path must be provided for loading data.")
|
||||
|
||||
self.data = []
|
||||
|
||||
# Load the JSONL file
|
||||
jsonl_data = pd.read_json(
|
||||
path_or_buf=Path(self.dataset_path) / f"{self.dataset_subset}.jsonl",
|
||||
lines=True,
|
||||
)
|
||||
|
||||
# check if the JSONL file has a 'turns' column
|
||||
if "messages" not in jsonl_data.columns:
|
||||
raise ValueError("JSONL file must contain a 'messages' column.")
|
||||
|
||||
for _, row in jsonl_data.iterrows():
|
||||
# sample only from a specific category if specified
|
||||
if (not self.category) or (self.category == row["category"]):
|
||||
prompt = row["messages"][0]["content"]
|
||||
self.data.append({"prompt": prompt})
|
||||
|
||||
random.seed(self.random_seed)
|
||||
if not getattr(self, "disable_shuffle", False):
|
||||
random.shuffle(self.data)
|
||||
|
||||
@@ -10,6 +10,7 @@ from torch._inductor.pattern_matcher import (
|
||||
PatternMatcherPass,
|
||||
register_graph_pattern,
|
||||
)
|
||||
from torch._library.triton import set_wrap_triton_enabled
|
||||
from torch._ops import OpOverload, OpOverloadPacket
|
||||
|
||||
from vllm.config import VllmConfig
|
||||
@@ -92,14 +93,19 @@ class VllmIRLoweringPass(VllmInductorPass):
|
||||
# Defaults not present on node.args but required for replacement tracing
|
||||
bound_args = ir_op._py_signature.bind(*node.args)
|
||||
bound_args.apply_defaults()
|
||||
match.replace_by_example(ir_op_impl.impl_fn, bound_args.args)
|
||||
match.replace_by_example(
|
||||
ir_op_impl.impl_fn, bound_args.args, run_functional_passes=False
|
||||
)
|
||||
|
||||
@VllmInductorPass.time_and_log
|
||||
def __call__(self, graph: fx.Graph) -> None:
|
||||
# clear at the beginning instead of end, so that tests can inspect
|
||||
self.selected_impls.clear()
|
||||
|
||||
count = self.patterns.apply(graph)
|
||||
# Triton wrap is disabled in the forward context, enable it during lowering.
|
||||
# This way make_fx replacement tracing handles the Triton kernel correctly.
|
||||
with set_wrap_triton_enabled(True):
|
||||
count = self.patterns.apply(graph)
|
||||
logger.debug("VllmIRLoweringPass lowered %d vLLM IR nodes", count)
|
||||
|
||||
# TODO write self.selected_impls to depyf/tlparse dir
|
||||
|
||||
@@ -113,8 +113,8 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
|
||||
VllmInductorPass.dump_prefix += 1
|
||||
|
||||
# clean up after lowering again
|
||||
self.post_cleanup(graph)
|
||||
VllmInductorPass.dump_prefix += 1
|
||||
# self.post_cleanup(graph)
|
||||
# VllmInductorPass.dump_prefix += 1
|
||||
|
||||
# always run fix_functionalization last
|
||||
self.fix_functionalization(graph)
|
||||
@@ -190,7 +190,7 @@ class PostGradPassManager(CustomGraphPass): # type: ignore[misc]
|
||||
|
||||
passes.append(self.post_cleanup.uuid())
|
||||
passes.append(self.ir_lowering.uuid())
|
||||
passes.append(self.post_cleanup.uuid())
|
||||
# passes.append(self.post_cleanup.uuid())
|
||||
passes.append(self.fix_functionalization.uuid())
|
||||
|
||||
# Include the compile range in the uuid to ensure that inductor
|
||||
|
||||
@@ -354,12 +354,22 @@ class PiecewiseBackend:
|
||||
return None
|
||||
|
||||
def __call__(self, *args: Any) -> Any:
|
||||
runtime_shape = args[self.sym_shape_indices[0]]
|
||||
range_entry = self._find_range_for_shape(runtime_shape)
|
||||
if self.sym_shape_indices:
|
||||
runtime_shape = args[self.sym_shape_indices[0]]
|
||||
range_entry = self._find_range_for_shape(runtime_shape)
|
||||
assert range_entry is not None, (
|
||||
f"Shape: {runtime_shape} out of considered ranges: "
|
||||
f"{self.compile_ranges}"
|
||||
)
|
||||
else:
|
||||
# All inputs have static shapes; use the only compiled range_entry
|
||||
compiled_entries = [re for re in self.range_entries.values() if re.compiled]
|
||||
assert len(compiled_entries) == 1, (
|
||||
f"Expected exactly one compiled range_entry for static shape "
|
||||
f"compilation, but found {len(compiled_entries)}"
|
||||
)
|
||||
range_entry = compiled_entries[0]
|
||||
|
||||
assert range_entry is not None, (
|
||||
f"Shape: {runtime_shape} out of considered ranges: {self.compile_ranges}"
|
||||
)
|
||||
assert range_entry.compiled, (
|
||||
"All ranges should be compiled or loaded up front in "
|
||||
"PiecewiseBackend.__init__. "
|
||||
|
||||
@@ -300,6 +300,10 @@ class SpeculativeConfig:
|
||||
{"n_predict": n_predict, "architectures": ["ErnieMTPModel"]}
|
||||
)
|
||||
|
||||
if hf_config.architectures[0] == "NemotronH_Super_Omni_Reasoning_V3":
|
||||
# Promote VLM's text_config so MTP detection below fires correctly
|
||||
hf_config = hf_config.text_config
|
||||
|
||||
if (
|
||||
hf_config.model_type in {"nemotron_h", "nemotron_h_puzzle"}
|
||||
and hasattr(hf_config, "num_nextn_predict_layers")
|
||||
|
||||
@@ -215,8 +215,11 @@ class LMCacheMPRequestTracker:
|
||||
# Main state
|
||||
state: LMCacheMPRequestState = LMCacheMPRequestState.PREFETCHING
|
||||
|
||||
cache_salt: str = ""
|
||||
|
||||
def __init__(self, request: "Request"):
|
||||
self.request_id = request.request_id
|
||||
self.cache_salt: str = request.cache_salt or ""
|
||||
self.all_token_ids = request.all_token_ids
|
||||
self.block_hashes = ConstantList(request.block_hashes)
|
||||
self.allocated_block_ids = []
|
||||
@@ -289,6 +292,7 @@ class LMCacheMPRequestMetadata:
|
||||
request_id: str
|
||||
direction: Literal["STORE", "RETRIEVE"]
|
||||
op: LoadStoreOp
|
||||
cache_salt: str = ""
|
||||
|
||||
@staticmethod
|
||||
def GetStoreMetadata(
|
||||
@@ -355,6 +359,7 @@ class LMCacheMPRequestMetadata:
|
||||
request_id=tracker.request_id,
|
||||
direction="STORE",
|
||||
op=op,
|
||||
cache_salt=tracker.cache_salt,
|
||||
)
|
||||
|
||||
# Update the request tracker
|
||||
@@ -421,6 +426,7 @@ class LMCacheMPRequestMetadata:
|
||||
request_id=tracker.request_id,
|
||||
direction="RETRIEVE",
|
||||
op=op,
|
||||
cache_salt=tracker.cache_salt,
|
||||
)
|
||||
return ret
|
||||
|
||||
@@ -569,12 +575,14 @@ class LMCacheMPConnector(KVConnectorBase_V1):
|
||||
|
||||
request_ids = []
|
||||
ops = []
|
||||
cache_salts = []
|
||||
|
||||
for meta in metadata.requests:
|
||||
if meta.direction != "RETRIEVE":
|
||||
continue
|
||||
request_ids.append(meta.request_id)
|
||||
ops.append(meta.op)
|
||||
cache_salts.append(meta.cache_salt)
|
||||
|
||||
if len(request_ids) == 0:
|
||||
return
|
||||
@@ -583,7 +591,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):
|
||||
event = torch.cuda.Event(interprocess=True)
|
||||
event.record()
|
||||
|
||||
self.worker_adapter.batched_submit_retrieve_requests(request_ids, ops, event)
|
||||
self.worker_adapter.batched_submit_retrieve_requests(
|
||||
request_ids, ops, event, cache_salts=cache_salts
|
||||
)
|
||||
|
||||
def wait_for_layer_load(self, layer_name: str) -> None:
|
||||
"""
|
||||
@@ -640,11 +650,13 @@ class LMCacheMPConnector(KVConnectorBase_V1):
|
||||
|
||||
request_ids = []
|
||||
ops = []
|
||||
cache_salts = []
|
||||
for meta in metadata.requests:
|
||||
if meta.direction != "STORE":
|
||||
continue
|
||||
request_ids.append(meta.request_id)
|
||||
ops.append(meta.op)
|
||||
cache_salts.append(meta.cache_salt)
|
||||
|
||||
if len(request_ids) == 0:
|
||||
return
|
||||
@@ -653,7 +665,9 @@ class LMCacheMPConnector(KVConnectorBase_V1):
|
||||
event = torch.cuda.Event(interprocess=True)
|
||||
event.record()
|
||||
|
||||
self.worker_adapter.batched_submit_store_requests(request_ids, ops, event)
|
||||
self.worker_adapter.batched_submit_store_requests(
|
||||
request_ids, ops, event, cache_salts=cache_salts
|
||||
)
|
||||
|
||||
def get_finished(
|
||||
self, finished_req_ids: set[str]
|
||||
@@ -755,6 +769,7 @@ class LMCacheMPConnector(KVConnectorBase_V1):
|
||||
self.scheduler_adapter.maybe_submit_lookup_request(
|
||||
request.request_id,
|
||||
token_ids=list(request.all_token_ids),
|
||||
cache_salt=tracker.cache_salt,
|
||||
)
|
||||
|
||||
ret = self.scheduler_adapter.check_lookup_result(request.request_id)
|
||||
|
||||
@@ -23,6 +23,7 @@ from vllm.distributed.kv_transfer.kv_connector.utils import (
|
||||
EngineId,
|
||||
TpKVTopology,
|
||||
get_current_attn_backend,
|
||||
get_current_attn_backends,
|
||||
)
|
||||
from vllm.distributed.kv_transfer.kv_connector.v1.base import (
|
||||
KVConnectorBase_V1,
|
||||
@@ -47,6 +48,7 @@ from vllm.v1.attention.backend import AttentionMetadata
|
||||
from vllm.v1.attention.backends.utils import get_kv_cache_layout
|
||||
from vllm.v1.core.sched.output import SchedulerOutput
|
||||
from vllm.v1.request import RequestStatus
|
||||
from vllm.v1.worker.utils import select_common_block_size
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
@@ -751,6 +753,7 @@ class MooncakeConnectorWorker:
|
||||
self.model_config = vllm_config.model_config
|
||||
self.cache_config = vllm_config.cache_config
|
||||
self.use_mla = self.model_config.use_mla
|
||||
self._sync_block_size_with_kernel()
|
||||
|
||||
# Get the attention backend from the first layer
|
||||
# NOTE (NickLucche) models with multiple backends are not supported yet
|
||||
@@ -777,6 +780,23 @@ class MooncakeConnectorWorker:
|
||||
self._xfer_meta_decoder = msgspec.msgpack.Decoder(MooncakeXferMetadata)
|
||||
self._xfer_resp_decoder = msgspec.msgpack.Decoder(MooncakeXferResponse)
|
||||
|
||||
def _sync_block_size_with_kernel(self) -> None:
|
||||
# When speculative decoding (e.g. Eagle) is enabled, the main model
|
||||
# and draft model may use different attention backends with different
|
||||
# physical block sizes. Pick the common (smallest) block size so that
|
||||
# KV-cache registration and transfer work correctly for both models.
|
||||
backends = get_current_attn_backends(self.vllm_config)
|
||||
kernel_block_size = select_common_block_size(self.block_size, backends)
|
||||
if self.block_size != kernel_block_size:
|
||||
logger.info_once(
|
||||
"User-specified logical block size (%s) does not match"
|
||||
" physical kernel block size (%s). Using the latter.",
|
||||
self.block_size,
|
||||
kernel_block_size,
|
||||
)
|
||||
assert self.block_size > kernel_block_size
|
||||
self.block_size = kernel_block_size
|
||||
|
||||
def __del__(self):
|
||||
self.shutdown()
|
||||
|
||||
@@ -1268,9 +1288,6 @@ class MooncakeConnectorWorker:
|
||||
self.block_len_per_layer.append(
|
||||
curr_tensor_size_bytes // self.num_blocks
|
||||
)
|
||||
|
||||
kernel_block_size = cache.shape[-2 if self.use_mla else -3]
|
||||
assert self.block_size == kernel_block_size
|
||||
kv_data_ptrs.append(base_addr)
|
||||
kv_data_lens.append(curr_tensor_size_bytes)
|
||||
|
||||
|
||||
@@ -1914,6 +1914,10 @@ class NixlConnectorWorker:
|
||||
meta.remote.block_ids,
|
||||
self._mamba_phys_ratio[meta.remote.engine_id],
|
||||
)
|
||||
else:
|
||||
meta.remote.block_ids = self._logical_to_kernel_block_ids(
|
||||
meta.remote.block_ids
|
||||
)
|
||||
# D may have to perform multiple reads from different remote ranks.
|
||||
for i, remote_rank in enumerate(remote_ranks):
|
||||
if self.use_mla and tp_ratio < 0 and i > 0:
|
||||
|
||||
@@ -424,7 +424,7 @@ class OffloadingConnectorScheduler:
|
||||
parent_block_hash=None,
|
||||
token_ids=[],
|
||||
lora_id=None,
|
||||
block_size=event.block_size,
|
||||
block_size=0,
|
||||
medium=event.medium,
|
||||
lora_name=None,
|
||||
)
|
||||
|
||||
@@ -290,7 +290,7 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):
|
||||
tool_call_id: str | None
|
||||
"""Tool call that this message is responding to."""
|
||||
|
||||
tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
|
||||
tool_calls: list[ChatCompletionMessageToolCallParam] | None
|
||||
"""The tool calls generated by the model, such as function calls."""
|
||||
|
||||
reasoning: str | None
|
||||
@@ -321,7 +321,7 @@ class ConversationMessage(TypedDict, total=False):
|
||||
name: str | None
|
||||
"""The name of the function to call"""
|
||||
|
||||
tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
|
||||
tool_calls: list[ChatCompletionMessageToolCallParam] | None
|
||||
"""The tool calls generated by the model, such as function calls."""
|
||||
|
||||
reasoning: str | None
|
||||
|
||||
@@ -357,6 +357,47 @@ class ChatCompletionRequest(OpenAIBaseModel):
|
||||
|
||||
# --8<-- [end:chat-completion-extra-params]
|
||||
|
||||
@model_validator(mode="before")
|
||||
@classmethod
|
||||
def _materialize_tool_calls_before(cls, data: Any) -> Any:
|
||||
"""Eagerly convert tool_calls generators/iterators to lists.
|
||||
|
||||
Must run before Pydantic field validation so that one-shot
|
||||
generators are not consumed during union type matching of
|
||||
ChatCompletionAssistantMessageParam (which types tool_calls
|
||||
as Iterable[...]).
|
||||
"""
|
||||
if not isinstance(data, dict):
|
||||
return data
|
||||
messages = data.get("messages")
|
||||
if not isinstance(messages, list):
|
||||
return data
|
||||
for msg in messages:
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
tool_calls = msg.get("tool_calls")
|
||||
if tool_calls is not None and not isinstance(tool_calls, list):
|
||||
msg["tool_calls"] = list(tool_calls)
|
||||
return data
|
||||
|
||||
@model_validator(mode="after")
|
||||
def _materialize_tool_calls_after(self) -> "ChatCompletionRequest":
|
||||
"""Convert Pydantic ValidatorIterator wrappers back to lists.
|
||||
|
||||
Even after the "before" validator converts iterables to lists,
|
||||
Pydantic re-wraps them in a ValidatorIterator when validating
|
||||
against ChatCompletionAssistantMessageParam's Iterable[...] type.
|
||||
This "after" pass materialises those wrappers so downstream code
|
||||
(tokenizers, model_dump_json) always sees plain lists.
|
||||
"""
|
||||
for msg in self.messages:
|
||||
if not isinstance(msg, dict):
|
||||
continue
|
||||
tool_calls = msg.get("tool_calls")
|
||||
if tool_calls is not None and not isinstance(tool_calls, list):
|
||||
msg["tool_calls"] = list(tool_calls)
|
||||
return self
|
||||
|
||||
def build_chat_params(
|
||||
self,
|
||||
default_template: str | None,
|
||||
|
||||
@@ -69,7 +69,10 @@ class AuthenticationMiddleware:
|
||||
return token_match
|
||||
|
||||
def __call__(self, scope: Scope, receive: Receive, send: Send) -> Awaitable[None]:
|
||||
if scope["type"] not in ("http", "websocket") or scope["method"] == "OPTIONS":
|
||||
if (
|
||||
scope["type"] not in ("http", "websocket")
|
||||
or scope.get("method") == "OPTIONS"
|
||||
):
|
||||
# scope["type"] can be "lifespan" or "startup" for example,
|
||||
# in which case we don't need to do anything
|
||||
return self.app(scope, receive, send)
|
||||
|
||||
@@ -8,6 +8,7 @@ from dataclasses import dataclass, field
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
from torch._library.triton import set_wrap_triton_enabled
|
||||
|
||||
import vllm.envs as envs
|
||||
import vllm.ir
|
||||
@@ -326,6 +327,7 @@ def set_forward_context(
|
||||
vllm.ir.enable_torch_wrap(
|
||||
vllm_config.compilation_config.ir_enable_torch_wrap
|
||||
),
|
||||
set_wrap_triton_enabled(False),
|
||||
):
|
||||
yield
|
||||
finally:
|
||||
|
||||
@@ -414,6 +414,7 @@ class INCConfig(QuantizationConfig):
|
||||
|
||||
def apply_xpu_w4a16_quant_layer(self, layer, prefix: str):
|
||||
weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
|
||||
|
||||
if not self.check_quantized(weight_bits):
|
||||
if isinstance(layer, (LinearBase, ParallelLMHead)):
|
||||
return UnquantizedLinearMethod()
|
||||
@@ -437,6 +438,27 @@ class INCConfig(QuantizationConfig):
|
||||
)
|
||||
return None
|
||||
|
||||
def apply_cpu_w4a16_quant_layer(self, layer, prefix: str):
|
||||
weight_bits, group_size, sym = self.get_layer_config(layer, prefix)
|
||||
if not self.check_quantized(weight_bits):
|
||||
if isinstance(layer, (LinearBase, ParallelLMHead)):
|
||||
return UnquantizedLinearMethod()
|
||||
else:
|
||||
return None
|
||||
|
||||
if weight_bits != 4:
|
||||
raise NotImplementedError(
|
||||
f"INC on CPU only supports 4-bit quantization, "
|
||||
f"got weight_bits={weight_bits}."
|
||||
)
|
||||
if not sym:
|
||||
raise NotImplementedError(
|
||||
"INC W4A16 on CPU only supports symmetric quantization for now."
|
||||
)
|
||||
if isinstance(layer, (LinearBase, ParallelLMHead)):
|
||||
return self.apply_gptq_quant_layer(layer, prefix)
|
||||
return None
|
||||
|
||||
def get_quant_method(self, layer: torch.nn.Module, prefix: str):
|
||||
if prefix and self.extra_config:
|
||||
for layer_name in self.extra_config:
|
||||
@@ -446,11 +468,21 @@ class INCConfig(QuantizationConfig):
|
||||
return UnquantizedLinearMethod()
|
||||
if current_platform.is_xpu():
|
||||
return self.apply_xpu_w4a16_quant_layer(layer, prefix)
|
||||
if "gptq" in self.packing_format or "gptq" in self.backend:
|
||||
is_gptq = "gptq" in self.packing_format or "gptq" in self.backend
|
||||
if current_platform.is_cpu() and is_gptq:
|
||||
return self.apply_cpu_w4a16_quant_layer(layer, prefix)
|
||||
if is_gptq:
|
||||
return self.apply_gptq_quant_layer(layer, prefix)
|
||||
if "awq" in self.packing_format or "awq" in self.backend:
|
||||
return self.apply_awq_quant_layer(layer, prefix)
|
||||
|
||||
raise NotImplementedError(
|
||||
f"Unsupported quantization configuration for layer '{prefix}'. "
|
||||
f"Platform: CPU={current_platform.is_cpu()}. "
|
||||
f"Platform: XPU={current_platform.is_xpu()}. "
|
||||
f"Format: {self.packing_format}, Backend: {self.backend}."
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def override_quantization_method(
|
||||
cls, hf_quant_cfg, user_quant, hf_config=None
|
||||
|
||||
@@ -162,3 +162,7 @@ try:
|
||||
quant_dequant_mxfp4 = torch.ops.vllm.quant_dequant_mxfp4
|
||||
except AttributeError as error:
|
||||
raise error
|
||||
|
||||
|
||||
def xpu_mxfp4_quantize(x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
return torch.ops.vllm.xpu_mxfp4_quantize(x)
|
||||
|
||||
@@ -30,6 +30,7 @@ from .deepseek_v2 import (
|
||||
DeepseekV2DecoderLayer,
|
||||
DeepseekV2MixtureOfExperts,
|
||||
DeepseekV2MoE,
|
||||
_try_load_fp8_indexer_wk,
|
||||
get_spec_layer_idx_from_weight_name,
|
||||
)
|
||||
from .utils import maybe_prefix
|
||||
@@ -190,10 +191,6 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
|
||||
)
|
||||
# Set MoE hyperparameters
|
||||
self.set_moe_parameters()
|
||||
self.is_fp4_ckpt = (
|
||||
self.quant_config is not None
|
||||
and self.quant_config.get_name() == "modelopt_fp4"
|
||||
)
|
||||
|
||||
def set_moe_parameters(self):
|
||||
self.expert_weights = []
|
||||
@@ -248,13 +245,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
|
||||
("fused_qkv_a_proj", "kv_a_proj_with_mqa", 1),
|
||||
]
|
||||
|
||||
if self.is_fp4_ckpt:
|
||||
# Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
|
||||
indexer_fused_mapping = [
|
||||
("wk_weights_proj", "wk", 0),
|
||||
("wk_weights_proj", "weights_proj", 1),
|
||||
]
|
||||
stacked_params_mapping.extend(indexer_fused_mapping)
|
||||
# Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
|
||||
indexer_fused_mapping = [
|
||||
("wk_weights_proj", "wk", 0),
|
||||
("wk_weights_proj", "weights_proj", 1),
|
||||
]
|
||||
stacked_params_mapping.extend(indexer_fused_mapping)
|
||||
|
||||
expert_params_mapping = SharedFusedMoE.make_expert_params_mapping(
|
||||
self,
|
||||
@@ -271,6 +267,7 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
|
||||
|
||||
params_dict = dict(self.named_parameters())
|
||||
loaded_params: set[str] = set()
|
||||
_pending_wk_fp8: dict = {} # FP8 indexer wk dequant buffer
|
||||
for name, loaded_weight in weights:
|
||||
if "rotary_emb.inv_freq" in name:
|
||||
continue
|
||||
@@ -281,6 +278,12 @@ class DeepSeekMTP(nn.Module, DeepseekV2MixtureOfExperts):
|
||||
rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
|
||||
)
|
||||
name = self._rewrite_spec_layer_name(spec_layer, name)
|
||||
|
||||
if _try_load_fp8_indexer_wk(
|
||||
name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
|
||||
):
|
||||
continue
|
||||
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
# Skip non-stacked layers and experts (experts handled below).
|
||||
if weight_name not in name:
|
||||
|
||||
@@ -66,6 +66,10 @@ from vllm.model_executor.layers.quantization import QuantizationConfig
|
||||
from vllm.model_executor.layers.quantization.utils.fp8_utils import (
|
||||
per_token_group_quant_fp8,
|
||||
)
|
||||
from vllm.model_executor.layers.quantization.utils.quant_utils import (
|
||||
GroupShape,
|
||||
scaled_dequantize,
|
||||
)
|
||||
from vllm.model_executor.layers.rotary_embedding import get_rope
|
||||
from vllm.model_executor.layers.sparse_attn_indexer import (
|
||||
SparseAttnIndexer,
|
||||
@@ -628,10 +632,6 @@ class Indexer(nn.Module):
|
||||
self.vllm_config = vllm_config
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
self.is_fp4_ckpt = (
|
||||
self.quant_config is not None
|
||||
and self.quant_config.get_name() == "modelopt_fp4"
|
||||
)
|
||||
# self.indexer_cfg = config.attn_module_list_cfg[0]["attn_index"]
|
||||
self.topk_tokens = config.index_topk
|
||||
self.n_head = config.index_n_heads # 64
|
||||
@@ -646,36 +646,16 @@ class Indexer(nn.Module):
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.wq_b",
|
||||
)
|
||||
if self.is_fp4_ckpt:
|
||||
# Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
|
||||
# weights_proj does not get quantized,
|
||||
# so we run both with quant_config=None
|
||||
# wk may be upcasted from the default quant;
|
||||
# experiments show fusion is always faster unless WK proj is in FP4,
|
||||
# which is not the case for all known quants.
|
||||
self.wk_weights_proj = MergedColumnParallelLinear(
|
||||
hidden_size,
|
||||
[self.head_dim, self.n_head],
|
||||
bias=False,
|
||||
quant_config=None,
|
||||
disable_tp=True,
|
||||
prefix=f"{prefix}.wk_weights_proj",
|
||||
)
|
||||
else:
|
||||
self.wk = ReplicatedLinear(
|
||||
hidden_size,
|
||||
self.head_dim,
|
||||
bias=False,
|
||||
quant_config=quant_config,
|
||||
prefix=f"{prefix}.wk",
|
||||
)
|
||||
self.weights_proj = ReplicatedLinear(
|
||||
hidden_size,
|
||||
self.n_head,
|
||||
bias=False,
|
||||
quant_config=None,
|
||||
prefix=f"{prefix}.weights_proj",
|
||||
)
|
||||
# Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
|
||||
# FP8 wk weights are upcasted to BF16 during loading to maintain fusion.
|
||||
self.wk_weights_proj = MergedColumnParallelLinear(
|
||||
hidden_size,
|
||||
[self.head_dim, self.n_head],
|
||||
bias=False,
|
||||
quant_config=None,
|
||||
disable_tp=True,
|
||||
prefix=f"{prefix}.wk_weights_proj",
|
||||
)
|
||||
self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
|
||||
self.softmax_scale = self.head_dim**-0.5
|
||||
|
||||
@@ -716,14 +696,10 @@ class Indexer(nn.Module):
|
||||
q_pe, q_nope = torch.split(
|
||||
q, [self.rope_dim, self.head_dim - self.rope_dim], dim=-1
|
||||
)
|
||||
if self.is_fp4_ckpt:
|
||||
# Fused wk + weights_proj: one GEMM, then split
|
||||
kw, _ = self.wk_weights_proj(hidden_states)
|
||||
k = kw[:, : self.head_dim]
|
||||
weights = kw[:, self.head_dim :]
|
||||
else:
|
||||
k, _ = self.wk(hidden_states)
|
||||
weights, _ = self.weights_proj(hidden_states)
|
||||
# Fused wk + weights_proj: one GEMM, then split
|
||||
kw, _ = self.wk_weights_proj(hidden_states)
|
||||
k = kw[:, : self.head_dim]
|
||||
weights = kw[:, self.head_dim :]
|
||||
|
||||
k = self.k_norm(k)
|
||||
k_pe, k_nope = torch.split(
|
||||
@@ -761,6 +737,46 @@ class Indexer(nn.Module):
|
||||
return self.indexer_op(hidden_states, q_fp8, k, weights)
|
||||
|
||||
|
||||
def _try_load_fp8_indexer_wk(name, tensor, buf, params_dict, loaded_params):
|
||||
"""
|
||||
We fuse the WK and weights_proj projections, but in some checkpoints WK is stored
|
||||
in FP8 with a separate weight_scale_inv, while weights_proj is stored in BF16.
|
||||
Upcasting to BF16 during loading enables the fusion. This function loads the FP8 WK
|
||||
weights and scale, and when both are available, dequantizes to BF16 and stores into
|
||||
the fused wk_weights_proj.weight parameter.
|
||||
"""
|
||||
if "indexer.wk." not in name or "wk_weights" in name:
|
||||
return False # Weight is not an isolated WK weight for the indexer, ignore.
|
||||
is_weight = name.endswith(".weight") and tensor.dtype == torch.float8_e4m3fn
|
||||
is_scale = "weight_scale_inv" in name
|
||||
if not is_weight and not is_scale:
|
||||
return False # WK is not in FP8 format, ignore.
|
||||
# Buffer this tensor (weight or scale) until both have arrived.
|
||||
layer_prefix = name.rsplit(".wk.", 1)[0] # e.g. "model.layers.0.self_attn.indexer"
|
||||
entry = buf.setdefault(layer_prefix, {})
|
||||
entry["weight" if is_weight else "scale"] = tensor
|
||||
if "weight" not in entry or "scale" not in entry:
|
||||
return True # still waiting for the other param
|
||||
|
||||
# We have both weight and scale: dequantize FP8 to BF16.
|
||||
weight_fp8, scale_inv = entry["weight"], entry["scale"]
|
||||
del buf[layer_prefix]
|
||||
block_size = weight_fp8.shape[1] // scale_inv.shape[1]
|
||||
weight_bf16 = scaled_dequantize(
|
||||
weight_fp8,
|
||||
scale_inv,
|
||||
group_shape=GroupShape(block_size, block_size),
|
||||
out_dtype=torch.bfloat16,
|
||||
)
|
||||
|
||||
# Load the dequantized weight into shard 0 of the fused buffer.
|
||||
fused_name = f"{layer_prefix}.wk_weights_proj.weight"
|
||||
param = params_dict[fused_name]
|
||||
param.weight_loader(param, weight_bf16, 0)
|
||||
loaded_params.add(fused_name)
|
||||
return True
|
||||
|
||||
|
||||
def _min_latency_fused_qkv_a_proj_impl(
|
||||
input_: torch.Tensor,
|
||||
weight: torch.Tensor,
|
||||
@@ -1344,10 +1360,6 @@ class DeepseekV2ForCausalLM(
|
||||
quant_config = vllm_config.quant_config
|
||||
self.config = config
|
||||
self.quant_config = quant_config
|
||||
self.is_fp4_ckpt = (
|
||||
self.quant_config is not None
|
||||
and self.quant_config.get_name() == "modelopt_fp4"
|
||||
)
|
||||
|
||||
qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
|
||||
qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
|
||||
@@ -1473,13 +1485,13 @@ class DeepseekV2ForCausalLM(
|
||||
("qkv_proj", "k_proj", "k"),
|
||||
("qkv_proj", "v_proj", "v"),
|
||||
]
|
||||
if self.is_fp4_ckpt:
|
||||
# Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
|
||||
indexer_fused_mapping = [
|
||||
("wk_weights_proj", "wk", 0),
|
||||
("wk_weights_proj", "weights_proj", 1),
|
||||
]
|
||||
stacked_params_mapping.extend(indexer_fused_mapping)
|
||||
# Fused indexer wk + weights_proj (shard 0 = wk, shard 1 = weights_proj)
|
||||
_pending_wk_fp8: dict = {} # When WK is in FP8, we dequant to BF16 for fusion
|
||||
indexer_fused_mapping = [
|
||||
("wk_weights_proj", "wk", 0),
|
||||
("wk_weights_proj", "weights_proj", 1),
|
||||
]
|
||||
stacked_params_mapping.extend(indexer_fused_mapping)
|
||||
|
||||
if self.use_mha:
|
||||
stacked_params_mapping.extend(mha_params_mapping)
|
||||
@@ -1516,6 +1528,11 @@ class DeepseekV2ForCausalLM(
|
||||
rocm_aiter_moe_shared_expert_enabled and ("mlp.shared_experts" in name)
|
||||
)
|
||||
|
||||
if _try_load_fp8_indexer_wk(
|
||||
name, loaded_weight, _pending_wk_fp8, params_dict, loaded_params
|
||||
):
|
||||
continue
|
||||
|
||||
for param_name, weight_name, shard_id in stacked_params_mapping:
|
||||
# Skip non-stacked layers and experts (experts handled below).
|
||||
if weight_name not in name:
|
||||
|
||||
@@ -1050,9 +1050,17 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
|
||||
quant_config=quant_config,
|
||||
prefix=maybe_prefix(prefix, "resampler"),
|
||||
)
|
||||
self._resampler_moved = False
|
||||
|
||||
self.make_empty_intermediate_tensors = self.llm.make_empty_intermediate_tensors
|
||||
|
||||
def _ensure_resampler_device(self) -> None:
|
||||
if self._resampler_moved:
|
||||
return
|
||||
# Only move device, DO NOT touch dtype (fp8 quant needs its own dtype)
|
||||
self.resampler.to(current_platform.device_type)
|
||||
self._resampler_moved = True
|
||||
|
||||
def _parse_and_validate_vision_input(
|
||||
self,
|
||||
modality: str,
|
||||
@@ -1171,7 +1179,9 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self)
|
||||
return loader.load_weights(weights)
|
||||
loaded = loader.load_weights(weights)
|
||||
self._ensure_resampler_device()
|
||||
return loaded
|
||||
|
||||
def get_mm_mapping(self) -> MultiModelKeys:
|
||||
"""
|
||||
@@ -1276,9 +1286,7 @@ class MiniCPMV2_0(MiniCPMVBaseModel):
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
return resampler.to(
|
||||
device=current_platform.device_type, dtype=torch.get_default_dtype()
|
||||
)
|
||||
return resampler.to(dtype=torch.get_default_dtype())
|
||||
|
||||
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
|
||||
pixel_values = data["pixel_values"]
|
||||
@@ -1359,9 +1367,7 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
return resampler.to(
|
||||
device=current_platform.device_type, dtype=torch.get_default_dtype()
|
||||
)
|
||||
return resampler.to(dtype=torch.get_default_dtype())
|
||||
|
||||
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
|
||||
pixel_values = data["pixel_values"]
|
||||
@@ -1452,11 +1458,8 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
target_device = current_platform.device_type
|
||||
target_dtype = torch.get_default_dtype()
|
||||
if any(p.is_meta for p in resampler.parameters()):
|
||||
return resampler.to_empty(device=target_device).to(dtype=target_dtype)
|
||||
return resampler.to(device=target_device, dtype=target_dtype)
|
||||
|
||||
return resampler.to(dtype=torch.get_default_dtype())
|
||||
|
||||
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
|
||||
pixel_values = data["pixel_values"]
|
||||
@@ -1491,7 +1494,9 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"])
|
||||
return loader.load_weights(weights)
|
||||
loaded = loader.load_weights(weights)
|
||||
self._ensure_resampler_device()
|
||||
return loaded
|
||||
|
||||
|
||||
class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
|
||||
@@ -1551,10 +1556,7 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
|
||||
return resampler.to(
|
||||
device=current_platform.device_type, dtype=torch.get_default_dtype()
|
||||
)
|
||||
return resampler.to(dtype=torch.get_default_dtype())
|
||||
|
||||
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
|
||||
pixel_values = data["pixel_values"]
|
||||
@@ -1589,7 +1591,9 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"])
|
||||
return loader.load_weights(weights)
|
||||
loaded = loader.load_weights(weights)
|
||||
self._ensure_resampler_device()
|
||||
return loaded
|
||||
|
||||
|
||||
class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
|
||||
@@ -1649,11 +1653,8 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
|
||||
quant_config=quant_config,
|
||||
prefix=prefix,
|
||||
)
|
||||
target_device = current_platform.device_type
|
||||
target_dtype = torch.get_default_dtype()
|
||||
if any(p.is_meta for p in resampler.parameters()):
|
||||
return resampler.to_empty(device=target_device).to(dtype=target_dtype)
|
||||
return resampler.to(device=target_device, dtype=target_dtype)
|
||||
|
||||
return resampler.to(dtype=torch.get_default_dtype())
|
||||
|
||||
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
|
||||
pixel_values = data["pixel_values"]
|
||||
@@ -1692,7 +1693,9 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
|
||||
|
||||
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
||||
loader = AutoWeightsLoader(self, skip_prefixes=["apm.", "audio", "tts"])
|
||||
return loader.load_weights(weights)
|
||||
loaded = loader.load_weights(weights)
|
||||
self._ensure_resampler_device()
|
||||
return loaded
|
||||
|
||||
|
||||
_SUPPORT_VERSION = {
|
||||
|
||||
@@ -37,6 +37,7 @@ from vllm.model_executor.models.nemotron_h import NemotronHForCausalLM
|
||||
from vllm.model_executor.models.parakeet import ParakeetExtractor, ProjectedParakeet
|
||||
from vllm.model_executor.models.radio import RadioModel, calc_seq_lens
|
||||
from vllm.model_executor.models.utils import (
|
||||
WeightsMapper,
|
||||
init_vllm_registered_model,
|
||||
maybe_prefix,
|
||||
)
|
||||
@@ -903,6 +904,12 @@ class NemotronH_Nano_VL_V2(
|
||||
requires_sequential_video_encoding = True
|
||||
"""Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet"""
|
||||
|
||||
hf_to_vllm_mapper = WeightsMapper(
|
||||
orig_to_new_prefix={
|
||||
"language_model.backbone": "language_model.model",
|
||||
},
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def get_placeholder_str(cls, modality: str, i: int) -> str | None:
|
||||
if modality.startswith("image"):
|
||||
|
||||
@@ -523,7 +523,14 @@ class DFlashQwen3ForCausalLM(Qwen3ForCausalLM):
|
||||
self.logits_processor = LogitsProcessor(
|
||||
self.config.draft_vocab_size, scale=logit_scale
|
||||
)
|
||||
self.draft_id_to_target_id = None
|
||||
target_vocab_size = vllm_config.model_config.get_vocab_size()
|
||||
if self.config.draft_vocab_size != target_vocab_size:
|
||||
self.draft_id_to_target_id = nn.Parameter(
|
||||
torch.zeros(self.config.draft_vocab_size, dtype=torch.long),
|
||||
requires_grad=False,
|
||||
)
|
||||
else:
|
||||
self.draft_id_to_target_id = None
|
||||
|
||||
def embed_input_ids(
|
||||
self,
|
||||
|
||||
@@ -473,6 +473,8 @@ _MULTIMODAL_MODELS = {
|
||||
"MolmoForCausalLM": ("molmo", "MolmoForCausalLM"),
|
||||
"Molmo2ForConditionalGeneration": ("molmo2", "Molmo2ForConditionalGeneration"),
|
||||
"NemotronH_Nano_VL_V2": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
|
||||
"NemotronH_Nano_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
|
||||
"NemotronH_Super_Omni_Reasoning_V3": ("nano_nemotron_vl", "NemotronH_Nano_VL_V2"),
|
||||
"NVLM_D": ("nvlm_d", "NVLM_D_Model"),
|
||||
"OpenCUAForConditionalGeneration": ("opencua", "OpenCUAForConditionalGeneration"),
|
||||
"OpenPanguVLForConditionalGeneration": (
|
||||
|
||||
@@ -41,3 +41,34 @@ def update_eagle3(config_dict: dict, pre_trained_config: dict) -> None:
|
||||
pre_trained_config["eagle_aux_hidden_state_layer_ids"] = config_dict[
|
||||
"eagle_aux_hidden_state_layer_ids"
|
||||
]
|
||||
|
||||
|
||||
@register_speculator("dflash")
|
||||
def update_dflash(config_dict: dict, pre_trained_config: dict) -> None:
|
||||
"""
|
||||
Apply DFlash specific configuration transformations to the `dict` used to
|
||||
construct the Transformers PreTrainedConfig.
|
||||
|
||||
DFlash specific fields:
|
||||
- draft_vocab_size: Size of the draft model's vocabulary
|
||||
- target_hidden_size: Hidden size of the target model
|
||||
- mask_token_id (required): Token ID used for parallel drafting mask
|
||||
placeholders
|
||||
- aux_hidden_state_layer_ids (required): Layer indices from the target
|
||||
model whose intermediate hidden states are used as context for the
|
||||
DFlash drafter. Mapped to both eagle_aux_hidden_state_layer_ids
|
||||
(for gpu_model_runner) and dflash_config.target_layer_ids (for the
|
||||
DFlash model).
|
||||
"""
|
||||
pre_trained_config["architectures"] = ["DFlashDraftModel"]
|
||||
pre_trained_config["draft_vocab_size"] = config_dict.get("draft_vocab_size")
|
||||
if config_dict.get("target_hidden_size") is not None:
|
||||
pre_trained_config["target_hidden_size"] = config_dict["target_hidden_size"]
|
||||
|
||||
aux_layer_ids = config_dict["aux_hidden_state_layer_ids"]
|
||||
pre_trained_config["eagle_aux_hidden_state_layer_ids"] = aux_layer_ids
|
||||
|
||||
pre_trained_config["dflash_config"] = {
|
||||
"mask_token_id": config_dict["mask_token_id"],
|
||||
"target_layer_ids": aux_layer_ids,
|
||||
}
|
||||
|
||||
@@ -1031,7 +1031,9 @@ class FlashAttentionImpl(AttentionImpl):
|
||||
window_size=sliding_window_size,
|
||||
softcap=self.logits_soft_cap,
|
||||
fa_version=self.vllm_flash_attn_version,
|
||||
q_descale=layer._q_scale.expand(descale_shape),
|
||||
q_descale=layer._q_scale.expand(descale_shape)
|
||||
if self.supports_quant_query_input
|
||||
else None,
|
||||
k_descale=layer._k_scale.expand(descale_shape),
|
||||
v_descale=layer._v_scale.expand(descale_shape),
|
||||
num_splits=1 if self.batch_invariant_enabled else 0,
|
||||
|
||||
@@ -799,6 +799,7 @@ class OutputProcessor:
|
||||
assert req_state.stats is not None
|
||||
iteration_stats.update_from_finished_request(
|
||||
finish_reason=finish_reason,
|
||||
request_id=req_state.external_req_id,
|
||||
num_prompt_tokens=req_state.prompt_len,
|
||||
max_tokens_param=req_state.max_tokens_param,
|
||||
req_stats=req_state.stats,
|
||||
|
||||
@@ -79,7 +79,6 @@ class PrepareStoreOutput:
|
||||
@dataclass
|
||||
class OffloadingEvent:
|
||||
keys: list[OffloadKey]
|
||||
block_size: int
|
||||
medium: str
|
||||
# True if blocks are removed, False if stored
|
||||
removed: bool
|
||||
|
||||
@@ -33,12 +33,10 @@ class CPUOffloadingManager(OffloadingManager):
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
block_size: int,
|
||||
num_blocks: int,
|
||||
cache_policy: Literal["lru", "arc"] = "lru",
|
||||
enable_events: bool = False,
|
||||
):
|
||||
self.block_size: int = block_size
|
||||
self.medium: str = CPULoadStoreSpec.medium()
|
||||
self._num_blocks: int = num_blocks
|
||||
self._num_allocated_blocks: int = 0
|
||||
@@ -145,7 +143,6 @@ class CPUOffloadingManager(OffloadingManager):
|
||||
self.events.append(
|
||||
OffloadingEvent(
|
||||
keys=to_evict,
|
||||
block_size=self.block_size,
|
||||
medium=self.medium,
|
||||
removed=True,
|
||||
)
|
||||
@@ -188,7 +185,6 @@ class CPUOffloadingManager(OffloadingManager):
|
||||
self.events.append(
|
||||
OffloadingEvent(
|
||||
keys=stored_keys,
|
||||
block_size=self.block_size,
|
||||
medium=self.medium,
|
||||
removed=False,
|
||||
)
|
||||
|
||||
@@ -60,12 +60,7 @@ class CPUOffloadingSpec(OffloadingSpec):
|
||||
kv_events_config is not None and kv_events_config.enable_kv_cache_events
|
||||
)
|
||||
|
||||
assert len(self.gpu_block_size) == 1
|
||||
gpu_block_size = self.gpu_block_size[0]
|
||||
offloaded_block_size = gpu_block_size * self.block_size_factor
|
||||
|
||||
self._manager = CPUOffloadingManager(
|
||||
block_size=offloaded_block_size,
|
||||
num_blocks=self.num_blocks,
|
||||
cache_policy=self.eviction_policy, # type: ignore[arg-type]
|
||||
enable_events=enable_events,
|
||||
|
||||
@@ -225,6 +225,7 @@ class FinishedRequestStats:
|
||||
"""Stats associated with a finished request."""
|
||||
|
||||
finish_reason: "FinishReason"
|
||||
request_id: str | None = None
|
||||
e2e_latency: float = 0.0
|
||||
num_prompt_tokens: int = 0
|
||||
num_generation_tokens: int = 0
|
||||
@@ -427,6 +428,7 @@ class IterationStats:
|
||||
def update_from_finished_request(
|
||||
self,
|
||||
finish_reason: "FinishReason",
|
||||
request_id: str,
|
||||
num_prompt_tokens: int,
|
||||
max_tokens_param: int | None,
|
||||
req_stats: RequestStateStats,
|
||||
@@ -458,6 +460,7 @@ class IterationStats:
|
||||
|
||||
finished_req = FinishedRequestStats(
|
||||
finish_reason=finish_reason,
|
||||
request_id=request_id,
|
||||
e2e_latency=e2e_latency,
|
||||
num_prompt_tokens=num_prompt_tokens,
|
||||
num_generation_tokens=req_stats.num_generation_tokens,
|
||||
|
||||
Reference in New Issue
Block a user