Compare commits

..
Author SHA1 Message Date
khluuandClaude Opus 4.6 5d236e5cf8 [Bugfix] Add MigCudaPlatform to avoid CUDA init before fork
NonNvmlCudaPlatform uses torch.cuda for get_device_capability(),
which initializes the CUDA context. When this runs at import time
(e.g. via requires_fp8 in test utils), it prevents fork-based tests
from working ("Cannot re-initialize CUDA in forked subprocess").

MigCudaPlatform inherits from NvmlCudaPlatform (NVML-based capability
queries don't init CUDA) and only overrides the memory methods that
fail on MIG partitions.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Signed-off-by: khluu <khluu000@gmail.com>
2026-04-16 11:49:08 -07:00
khluu f683266536 lint
Signed-off-by: khluu <khluu000@gmail.com>
2026-04-16 02:35:54 -07:00
khluuandClaude Opus 4.6 8dde2576c0 [Bugfix] Disable expandable_segments on MIG to fix PyTorch NVML assert
The previous commit added MIG detection to use NonNvmlCudaPlatform, which
fixes vllm's own NVML calls. However, PyTorch's CUDACachingAllocator has
a separate NVML code path via the expandable_segments feature that still
crashes with "NVML_SUCCESS == r INTERNAL ASSERT FAILED" on MIG devices.

This change automatically sets PYTORCH_CUDA_ALLOC_CONF=expandable_segments:False
when a MIG device is detected, so PyTorch's allocator falls back to non-NVML
memory management and real torch.cuda.OutOfMemoryError surfaces instead of
the misleading NVML assert.

Signed-off-by: khluu <khluu000@gmail.com>
Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-15 16:09:29 -07:00
khluuandClaude Opus 4.6 ae60bf5ea7 [Bugfix] Use NonNvmlCudaPlatform on MIG devices to fix NVML masking OOM errors
On MIG partitions, NVML device-level queries fail with NVMLError_NoPermission.
This causes PyTorch's CUDACachingAllocator (via expandable_segments) to crash
with an internal assert before the real torch.cuda.OutOfMemoryError can surface,
and causes NVML memory queries to report the full GPU's memory (e.g. 141 GB for
H200) instead of the MIG partition's memory (e.g. 16 GB).

This change:
1. Adds MIG detection in vllm/platforms/cuda.py that checks CUDA_VISIBLE_DEVICES
   for MIG UUIDs and probes NVML for permission errors, then falls back to
   NonNvmlCudaPlatform which uses torch.cuda APIs that correctly report MIG
   partition memory.
2. Patches test utilities (wait_for_gpu_memory_to_clear, _get_gpu_memory_used)
   to fall back to torch.cuda.mem_get_info() when NVML fails with permission
   errors on MIG.

Signed-off-by: khluu <khluu000@gmail.com>
Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
2026-04-14 23:44:59 -07:00
24 changed files with 200 additions and 406 deletions
-2
View File
@@ -1,5 +1,3 @@
lmcache >= 0.3.9
nixl[cu13] >= 0.7.1, < 0.10.0 # Required for disaggregated prefill
nixl-cu12 >= 0.7.1, < 0.10.0
nixl-cu13 >= 0.7.1, < 0.10.0
mooncake-transfer-engine >= 0.3.8
@@ -222,47 +222,3 @@ 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()
@@ -1,150 +0,0 @@
# 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
-1
View File
@@ -12,7 +12,6 @@ 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"]
+20 -8
View File
@@ -383,9 +383,14 @@ class RemoteVLLMServer:
total_used = 0
device_count = current_platform.device_count()
for i in range(device_count):
handle = nvmlDeviceGetHandleByIndex(i)
mem_info = nvmlDeviceGetMemoryInfo(handle)
total_used += mem_info.used
try:
handle = nvmlDeviceGetHandleByIndex(i)
mem_info = nvmlDeviceGetMemoryInfo(handle)
total_used += mem_info.used
except Exception:
# NVML fails on MIG — fall back to torch.cuda
free, total = torch.cuda.mem_get_info(i)
total_used += total - free
return total_used
except Exception as e:
print(f"[RemoteOpenAIServer] Could not query GPU memory: {e}")
@@ -1202,7 +1207,7 @@ def wait_for_gpu_memory_to_clear(
) -> None:
assert threshold_bytes is not None or threshold_ratio is not None
# Use nvml instead of pytorch to reduce measurement error from torch cuda
# context.
# context. Falls back to torch.cuda on MIG where NVML lacks permissions.
devices = get_physical_device_indices(devices)
start_time = time.time()
while True:
@@ -1215,10 +1220,17 @@ def wait_for_gpu_memory_to_clear(
gb_used = mem_info["vram_used"] / 2**10
gb_total = mem_info["vram_total"] / 2**10
else:
dev_handle = nvmlDeviceGetHandleByIndex(device)
mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
gb_used = mem_info.used / 2**30
gb_total = mem_info.total / 2**30
try:
dev_handle = nvmlDeviceGetHandleByIndex(device)
mem_info = nvmlDeviceGetMemoryInfo(dev_handle)
gb_used = mem_info.used / 2**30
gb_total = mem_info.total / 2**30
except Exception:
# NVML fails on MIG with NoPermission — fall back to
# torch.cuda which correctly reports MIG partition memory.
free, total = torch.cuda.mem_get_info(device)
gb_total = total / 2**30
gb_used = (total - free) / 2**30
output_raw[device] = (gb_used, gb_total)
output[device] = f"{gb_used:.02f}/{gb_total:.02f}"
@@ -124,8 +124,12 @@ 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]), medium="A", removed=False)
yield OffloadingEvent(keys=to_keys([4, 5, 6]), medium="B", removed=True)
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
)
runner.manager.take_events.side_effect = take_events
events = list(runner.scheduler_connector.take_events())
@@ -133,7 +137,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 == 0
assert event.block_size == 16
assert event.medium == "A"
assert event.token_ids == []
assert event.parent_block_hash is None
+20 -6
View File
@@ -59,6 +59,7 @@ def verify_load_output(
def verify_events(
events: Iterable[OffloadingEvent],
block_size: int,
expected_stores: tuple[set[int], ...] = (),
expected_evictions: tuple[set[int], ...] = (),
):
@@ -66,6 +67,7 @@ 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:
@@ -96,7 +98,9 @@ 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,
@@ -134,9 +138,10 @@ def test_cpu_manager():
"""
Tests CPUOffloadingManager with lru policy.
"""
# initialize a CPU manager with a capacity of 4 blocks
# initialize a CPU backend with a capacity of 4 blocks
block_size = 256
cpu_manager = CPUOffloadingManager(
num_blocks=4, cache_policy="lru", enable_events=True
block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
)
# prepare store [1, 2]
@@ -158,7 +163,9 @@ def test_cpu_manager():
# complete store [1, 2]
cpu_manager.complete_store(to_keys([1, 2]))
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
verify_events(
cpu_manager.take_events(), block_size=block_size, expected_stores=({1, 2},)
)
# lookup [1, 2]
assert cpu_manager.lookup(to_keys([1])) == 1
@@ -177,7 +184,9 @@ def test_cpu_manager():
)
# verify eviction event
verify_events(cpu_manager.take_events(), expected_evictions=({1},))
verify_events(
cpu_manager.take_events(), block_size=block_size, expected_evictions=({1},)
)
# prepare store with no space
assert cpu_manager.prepare_store(to_keys([1, 6])) is None
@@ -232,6 +241,7 @@ 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}),
)
@@ -244,6 +254,7 @@ 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,
@@ -278,7 +289,9 @@ class TestARCPolicy:
# complete store [1, 2]
cpu_manager.complete_store(to_keys([1, 2]))
verify_events(cpu_manager.take_events(), expected_stores=({1, 2},))
verify_events(
cpu_manager.take_events(), block_size=256, expected_stores=({1, 2},)
)
# lookup [1, 2]
assert cpu_manager.lookup(to_keys([1])) == 1
@@ -534,8 +547,9 @@ def test_filter_reused_manager():
"""
Tests FilterReusedOffloadingManager with a CPUOffloadingManager.
"""
block_size = 256
lru_manager = CPUOffloadingManager(
num_blocks=4, cache_policy="lru", enable_events=True
block_size=block_size, num_blocks=4, cache_policy="lru", enable_events=True
)
manager = FilterReusedOffloadingManager(
-8
View File
@@ -26,7 +26,6 @@ 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,
@@ -36,7 +35,6 @@ 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(
@@ -57,7 +55,6 @@ 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,
@@ -67,7 +64,6 @@ 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(
@@ -88,7 +84,6 @@ 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,
@@ -101,13 +96,11 @@ 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,
@@ -119,7 +112,6 @@ 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():
-46
View File
@@ -144,46 +144,6 @@ 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
@@ -595,12 +555,6 @@ 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
+5 -15
View File
@@ -354,22 +354,12 @@ class PiecewiseBackend:
return None
def __call__(self, *args: Any) -> Any:
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]
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: {self.compile_ranges}"
)
assert range_entry.compiled, (
"All ranges should be compiled or loaded up front in "
"PiecewiseBackend.__init__. "
@@ -424,7 +424,7 @@ class OffloadingConnectorScheduler:
parent_block_hash=None,
token_ids=[],
lora_id=None,
block_size=0,
block_size=event.block_size,
medium=event.medium,
lora_name=None,
)
+2 -2
View File
@@ -290,7 +290,7 @@ class CustomChatCompletionMessageParam(TypedDict, total=False):
tool_call_id: str | None
"""Tool call that this message is responding to."""
tool_calls: list[ChatCompletionMessageToolCallParam] | None
tool_calls: Iterable[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: list[ChatCompletionMessageToolCallParam] | None
tool_calls: Iterable[ChatCompletionMessageToolCallParam] | None
"""The tool calls generated by the model, such as function calls."""
reasoning: str | None
@@ -357,47 +357,6 @@ 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,
+1 -33
View File
@@ -414,7 +414,6 @@ 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()
@@ -438,27 +437,6 @@ 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:
@@ -468,21 +446,11 @@ class INCConfig(QuantizationConfig):
return UnquantizedLinearMethod()
if current_platform.is_xpu():
return self.apply_xpu_w4a16_quant_layer(layer, prefix)
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:
if "gptq" in self.packing_format or "gptq" in self.backend:
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,7 +162,3 @@ 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)
+24 -27
View File
@@ -1050,17 +1050,9 @@ 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,
@@ -1179,9 +1171,7 @@ class MiniCPMVBaseModel(nn.Module, SupportsMultiModal, SupportsPP):
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(self)
loaded = loader.load_weights(weights)
self._ensure_resampler_device()
return loaded
return loader.load_weights(weights)
def get_mm_mapping(self) -> MultiModelKeys:
"""
@@ -1286,7 +1276,9 @@ class MiniCPMV2_0(MiniCPMVBaseModel):
prefix=prefix,
)
return resampler.to(dtype=torch.get_default_dtype())
return resampler.to(
device=current_platform.device_type, dtype=torch.get_default_dtype()
)
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
pixel_values = data["pixel_values"]
@@ -1367,7 +1359,9 @@ class MiniCPMV2_5(MiniCPMVBaseModel, SupportsLoRA):
prefix=prefix,
)
return resampler.to(dtype=torch.get_default_dtype())
return resampler.to(
device=current_platform.device_type, dtype=torch.get_default_dtype()
)
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
pixel_values = data["pixel_values"]
@@ -1458,8 +1452,11 @@ class MiniCPMV2_6(MiniCPMVBaseModel, SupportsLoRA):
quant_config=quant_config,
prefix=prefix,
)
return resampler.to(dtype=torch.get_default_dtype())
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)
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
pixel_values = data["pixel_values"]
@@ -1494,9 +1491,7 @@ 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"])
loaded = loader.load_weights(weights)
self._ensure_resampler_device()
return loaded
return loader.load_weights(weights)
class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
@@ -1556,7 +1551,10 @@ class MiniCPMV4_0(MiniCPMVBaseModel, SupportsLoRA):
quant_config=quant_config,
prefix=prefix,
)
return resampler.to(dtype=torch.get_default_dtype())
return resampler.to(
device=current_platform.device_type, dtype=torch.get_default_dtype()
)
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
pixel_values = data["pixel_values"]
@@ -1591,9 +1589,7 @@ 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"])
loaded = loader.load_weights(weights)
self._ensure_resampler_device()
return loaded
return loader.load_weights(weights)
class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
@@ -1653,8 +1649,11 @@ class MiniCPMV4_5(MiniCPMVBaseModel, SupportsLoRA):
quant_config=quant_config,
prefix=prefix,
)
return resampler.to(dtype=torch.get_default_dtype())
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)
def get_vision_hidden_states(self, data: MiniCPMVImagePixelInputs) -> torch.Tensor:
pixel_values = data["pixel_values"]
@@ -1693,9 +1692,7 @@ 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"])
loaded = loader.load_weights(weights)
self._ensure_resampler_device()
return loaded
return loader.load_weights(weights)
_SUPPORT_VERSION = {
@@ -37,7 +37,6 @@ 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,
)
@@ -904,12 +903,6 @@ 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"):
+109 -1
View File
@@ -844,6 +844,37 @@ class NonNvmlCudaPlatform(CudaPlatformBase):
return None
class MigCudaPlatform(NvmlCudaPlatform):
"""Platform for MIG (Multi-Instance GPU) devices.
Inherits from NvmlCudaPlatform so that device queries like
get_device_capability() and get_device_name() use NVML instead of
torch.cuda, avoiding premature CUDA context initialization (which
would break fork-based test frameworks).
Overrides only the methods that fail on MIG partitions due to NVML
NoPermission errors (memory queries report full-GPU memory instead
of partition memory).
"""
@classmethod
def get_device_total_memory(cls, device_id: int = 0) -> int:
# NVML returns the *full* GPU's memory on MIG partitions, not the
# partition's memory. Use torch.cuda which reports correctly.
device_props = torch.cuda.get_device_properties(device_id)
return device_props.total_memory
@classmethod
def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:
# P2P / NVLink queries are not meaningful on MIG partitions.
return False
@classmethod
def log_warnings(cls):
# Skip NVML-based heterogeneous-GPU warnings on MIG partitions.
pass
# Autodetect either NVML-enabled or non-NVML platform
# based on whether NVML is available.
nvml_available = False
@@ -858,6 +889,83 @@ finally:
if nvml_available:
pynvml.nvmlShutdown()
CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
def _is_mig_device() -> bool:
"""Detect if the current CUDA device is a MIG (Multi-Instance GPU) partition.
On MIG partitions, most NVML device-level queries fail with
NVMLError_NoPermission. This causes PyTorch's CUDACachingAllocator to crash
with an internal assert (via expandable_segments) before the real OOM error
can surface. It also causes NVML memory queries to report the full GPU's
memory instead of the MIG partition's memory.
Detection methods:
1. CUDA_VISIBLE_DEVICES containing MIG UUIDs (MIG-GPU-...)
2. Probing NVML device queries for permission errors
"""
# Method 1: Check CUDA_VISIBLE_DEVICES for MIG UUID format
cuda_visible = os.environ.get("CUDA_VISIBLE_DEVICES", "")
if cuda_visible and any(
dev.strip().startswith("MIG-") for dev in cuda_visible.split(",")
):
logger.info(
"Detected MIG device from CUDA_VISIBLE_DEVICES=%s, "
"using MigCudaPlatform to avoid NVML memory errors.",
cuda_visible,
)
return True
# Method 2: Probe NVML for permission errors (catches MIG even when
# CUDA_VISIBLE_DEVICES uses numeric indices mapped by the container runtime)
if nvml_available:
try:
pynvml.nvmlInit()
try:
handle = pynvml.nvmlDeviceGetHandleByIndex(0)
pynvml.nvmlDeviceGetMemoryInfo(handle)
except pynvml.NVMLError as e:
if "NoPermission" in str(
type(e).__name__
) or "Insufficient Permissions" in str(e):
logger.info(
"Detected MIG device via NVML permission error: %s, "
"using MigCudaPlatform.",
e,
)
return True
finally:
pynvml.nvmlShutdown()
except Exception:
pass
return False
_mig_device = _is_mig_device()
CudaPlatform: type[NvmlCudaPlatform] | type[NonNvmlCudaPlatform] | type[MigCudaPlatform]
if _mig_device:
CudaPlatform = MigCudaPlatform
# Disable expandable_segments on MIG devices. PyTorch's
# CUDACachingAllocator internally calls NVML (for the
# expandable_segments feature) which fails on MIG with
# "NVML_SUCCESS == r INTERNAL ASSERT FAILED", crashing before
# a real torch.cuda.OutOfMemoryError can surface.
_alloc_conf = os.environ.get("PYTORCH_CUDA_ALLOC_CONF", "")
if "expandable_segments" not in _alloc_conf:
new_val = (
f"{_alloc_conf},expandable_segments:False"
if _alloc_conf
else "expandable_segments:False"
)
os.environ["PYTORCH_CUDA_ALLOC_CONF"] = new_val
logger.info(
"MIG device detected: set PYTORCH_CUDA_ALLOC_CONF=%s "
"to prevent NVML assert in CUDACachingAllocator.",
new_val,
)
else:
CudaPlatform = NvmlCudaPlatform if nvml_available else NonNvmlCudaPlatform
CudaPlatform.log_warnings()
+1 -3
View File
@@ -1031,9 +1031,7 @@ 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)
if self.supports_quant_query_input
else None,
q_descale=layer._q_scale.expand(descale_shape),
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,
-1
View File
@@ -799,7 +799,6 @@ 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,
+1
View File
@@ -79,6 +79,7 @@ class PrepareStoreOutput:
@dataclass
class OffloadingEvent:
keys: list[OffloadKey]
block_size: int
medium: str
# True if blocks are removed, False if stored
removed: bool
+4
View File
@@ -33,10 +33,12 @@ 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
@@ -143,6 +145,7 @@ class CPUOffloadingManager(OffloadingManager):
self.events.append(
OffloadingEvent(
keys=to_evict,
block_size=self.block_size,
medium=self.medium,
removed=True,
)
@@ -185,6 +188,7 @@ class CPUOffloadingManager(OffloadingManager):
self.events.append(
OffloadingEvent(
keys=stored_keys,
block_size=self.block_size,
medium=self.medium,
removed=False,
)
+5
View File
@@ -60,7 +60,12 @@ 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,
-3
View File
@@ -225,7 +225,6 @@ 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
@@ -428,7 +427,6 @@ 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,
@@ -460,7 +458,6 @@ 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,