From 42489e43c2718674828ece00eefc0f11088e801d Mon Sep 17 00:00:00 2001 From: Bhoomit Date: Wed, 25 Feb 2026 07:30:55 -0800 Subject: [PATCH 01/43] [Misc][LoRA] Increase max vocab size limit to 258048 in logits processor (#34773) Signed-off-by: Bhoomit Vasani --- tests/lora/conftest.py | 12 ++++++------ tests/lora/test_layers.py | 27 ++++++++++++++++++++++++++- vllm/lora/layers/logits_processor.py | 4 ++-- 3 files changed, 34 insertions(+), 9 deletions(-) diff --git a/tests/lora/conftest.py b/tests/lora/conftest.py index deb1ab92d70..d0d8382acf9 100644 --- a/tests/lora/conftest.py +++ b/tests/lora/conftest.py @@ -103,14 +103,14 @@ def dummy_model(default_vllm_config) -> nn.Module: ("output", ColumnParallelLinear(50, 10)), ("outact", nn.Sigmoid()), # Special handling for lm_head & sampler - ("lm_head", ParallelLMHead(512, 10)), - ("logits_processor", LogitsProcessor(512)), + ("lm_head", ParallelLMHead(32064, 10)), + ("logits_processor", LogitsProcessor(32064)), ] ) ) model.config = MagicMock() model.embedding_modules = {"lm_head": "lm_head"} - model.unpadded_vocab_size = 32000 + model.unpadded_vocab_size = 32064 return model @@ -136,8 +136,8 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module: ("gate_up_proj", MergedColumnParallelLinear(50, [5, 5])), ("outact", nn.Sigmoid()), # Special handling for lm_head & sampler - ("lm_head", ParallelLMHead(512, 10)), - ("logits_processor", LogitsProcessor(512)), + ("lm_head", ParallelLMHead(32064, 10)), + ("logits_processor", LogitsProcessor(32064)), ] ) ) @@ -149,7 +149,7 @@ def dummy_model_gate_up(default_vllm_config) -> nn.Module: ], } model.embedding_modules = {"lm_head": "lm_head"} - model.unpadded_vocab_size = 32000 + model.unpadded_vocab_size = 32064 return model diff --git a/tests/lora/test_layers.py b/tests/lora/test_layers.py index 2a96529d889..c9c55114360 100644 --- a/tests/lora/test_layers.py +++ b/tests/lora/test_layers.py @@ -353,7 +353,7 @@ def test_embeddings( @torch.inference_mode() @pytest.mark.parametrize("num_loras", [1, 2, 4]) @pytest.mark.parametrize("device", DEVICES) -@pytest.mark.parametrize("vocab_size", [512, 32000, 64000, 256512]) +@pytest.mark.parametrize("vocab_size", [64000, 256512, 258048]) @pytest.mark.parametrize("stage", STAGES) def test_lm_head_logits_processor( default_vllm_config, dist_init, num_loras, device, vocab_size, stage @@ -468,6 +468,31 @@ def test_lm_head_logits_processor( torch.testing.assert_close(lora_result, expected_result, rtol=rtol, atol=atol) +@torch.inference_mode() +@pytest.mark.parametrize("vocab_size", [512, 32000, 258049, 300000]) +@pytest.mark.parametrize("device", DEVICES) +def test_lm_head_logits_processor_invalid_vocab_size( + default_vllm_config, dist_init, vocab_size, device +) -> None: + """Test that LogitsProcessorWithLoRA raises ValueError for invalid vocab sizes.""" + if current_platform.is_cuda_alike(): + torch.cuda.set_device(device) + + torch.set_default_device(device) + max_loras = 8 + lora_config = LoRAConfig( + max_loras=max_loras, max_lora_rank=8, lora_dtype=torch.float16 + ) + + logits_processor = LogitsProcessor(vocab_size) + lora_logits_processor = LogitsProcessorWithLoRA( + logits_processor, 1024, torch.float16, device, None + ) + + with pytest.raises(ValueError, match="vocab size must be > 32000 and <= 258048"): + lora_logits_processor.create_lora_weights(max_loras, lora_config) + + @torch.inference_mode() @pytest.mark.parametrize("num_loras", [1, 2, 4]) @pytest.mark.parametrize("device", DEVICES) diff --git a/vllm/lora/layers/logits_processor.py b/vllm/lora/layers/logits_processor.py index d7b02ec9678..217c46fbec4 100644 --- a/vllm/lora/layers/logits_processor.py +++ b/vllm/lora/layers/logits_processor.py @@ -88,9 +88,9 @@ class LogitsProcessorWithLoRA(BaseLayerWithLoRA): model_config: PretrainedConfig | None = None, ) -> None: # TODO: Verify if this condition can be further relaxed - if 32000 < self.base_layer.vocab_size > 257024: + if self.base_layer.vocab_size <= 32000 or self.base_layer.vocab_size > 258048: raise ValueError( - "When using LoRA, vocab size must be 32000 >= vocab_size <= 257024" + "When using LoRA, vocab size must be > 32000 and <= 258048" ) self.lora_a_stacked = torch.zeros( ( From d72b0be33cdd561e557df1ce5350a14451b9af13 Mon Sep 17 00:00:00 2001 From: "Chendi.Xue" Date: Wed, 25 Feb 2026 09:31:07 -0600 Subject: [PATCH 02/43] [XPU]Fix for Qwen-OMNI crash (#35249) Signed-off-by: Chendi Xue --- vllm/_xpu_ops.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/_xpu_ops.py b/vllm/_xpu_ops.py index e40b18f8151..1f64aacd421 100644 --- a/vllm/_xpu_ops.py +++ b/vllm/_xpu_ops.py @@ -105,9 +105,10 @@ class xpu_ops: assert len(window_size) == 2 real_window_size = (window_size[0], window_size[1]) # noqa: F841 - # In encode attention, v maybe not contiguous and current + # In encode attention, k and v maybe not contiguous and current # kernel can't handle it if block_table is None: + k = k.contiguous() v = v.contiguous() return flash_attn_varlen_func( out=out, From 0788ff0a153c6eb6436743d717b31c863c987761 Mon Sep 17 00:00:00 2001 From: haosdent Date: Wed, 25 Feb 2026 23:31:45 +0800 Subject: [PATCH 03/43] [Bugfix] Gracefully disable AllReduceFusionPass on GPUs without multicast support (#35085) Signed-off-by: haosdent --- .../passes/fusion/allreduce_rms_fusion.py | 28 +++++++++++++------ 1 file changed, 20 insertions(+), 8 deletions(-) diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index b613d4424ee..b6a1314af9e 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -729,14 +729,26 @@ class AllReduceFusionPass(VllmPatternMatcherPass): scope="global", ) - self.workspace = flashinfer_comm.create_allreduce_fusion_workspace( - backend="trtllm", - world_size=self.tp_size, - rank=rank, - max_token_num=self.max_token_num, - hidden_dim=self.hidden_dim, - dtype=self.model_dtype, - ) + try: + self.workspace = flashinfer_comm.create_allreduce_fusion_workspace( + backend="trtllm", + world_size=self.tp_size, + rank=rank, + max_token_num=self.max_token_num, + hidden_dim=self.hidden_dim, + dtype=self.model_dtype, + ) + except RuntimeError as e: + if "multicast" not in str(e).lower(): + raise + logger.warning_once( + "AllReduce fusion pass is disabled: flashinfer workspace " + "creation failed: %s. This is expected on GPUs without " + "NVSwitch (e.g., NVLink bridge-only or PCIe topologies). " + "Falling back to non-fused allreduce.", + str(e), + ) + return global _FI_WORKSPACE _FI_WORKSPACE = self.workspace From 5d18bf8b32837275d7656e2ae8b5c684274234d2 Mon Sep 17 00:00:00 2001 From: pushkar Date: Wed, 25 Feb 2026 21:38:16 +0530 Subject: [PATCH 04/43] [Bugfix] Fix Harmony preamble visibility in Responses API (#32114) Signed-off-by: Pushkar Patel Signed-off-by: pupa --- .../openai/parser/test_harmony_utils.py | 123 ++++++++++++++++-- .../openai/responses/test_harmony.py | 17 ++- .../openai/responses/test_mcp_tools.py | 4 +- .../test_serving_chat_stream_harmony.py | 14 +- .../openai/test_serving_responses.py | 115 ++++++++++++++++ tests/entrypoints/test_context.py | 38 ++++++ .../openai/chat_completion/stream_harmony.py | 2 +- .../openai/parser/harmony_utils.py | 75 +++++++---- vllm/entrypoints/openai/responses/context.py | 8 +- .../openai/responses/streaming_events.py | 8 +- 10 files changed, 341 insertions(+), 63 deletions(-) diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index 1d34fc51ad5..b73a0b0745c 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -2,7 +2,11 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from openai.types.responses import ResponseFunctionToolCall, ResponseReasoningItem +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseReasoningItem, +) from openai.types.responses.response_output_item import McpCall from openai_harmony import Author, Message, Role, TextContent @@ -10,6 +14,7 @@ from tests.entrypoints.openai.utils import verify_harmony_messages from vllm.entrypoints.openai.parser.harmony_utils import ( auto_drop_analysis_messages, get_encoding, + get_system_message, has_custom_tools, parse_chat_input_to_harmony_message, parse_chat_output, @@ -840,15 +845,58 @@ class TestParseChatOutput: assert reasoning == "I've thought hard about this." assert final_content == "The answer is 4." + def test_parse_chat_output_commentary_with_recipient_excluded(self) -> None: + """Commentary with a recipient (tool call) should not appear in + final_content — those are handled separately by the tool parser. + + The first message is a preamble (visible), the second is a tool + call (excluded). Only the preamble should appear in final_content. + """ + harmony_str = ( + "<|channel|>commentary" + "<|message|>Let me check the weather.<|end|>" + "<|start|>assistant to=functions.get_weather" + "<|channel|>commentary" + '<|message|>{"location": "SF"}<|end|>' + ) + token_ids = get_encoding().encode(harmony_str, allowed_special="all") + reasoning, final_content, _ = parse_chat_output(token_ids) + assert reasoning is None + assert final_content == "Let me check the weather." + + def test_parse_chat_output_interrupted_preamble(self) -> None: + """Partial/interrupted preamble (commentary without recipient) should + appear in final_content, not reasoning.""" + harmony_str = "<|channel|>commentary<|message|>I'll search for that" + token_ids = get_encoding().encode(harmony_str, allowed_special="all") + reasoning, final_content, _ = parse_chat_output(token_ids) + assert reasoning is None + assert final_content == "I'll search for that" + + def test_parse_chat_output_preamble_then_final(self) -> None: + """Preamble followed by a final message should both appear in + final_content, joined by newline.""" + harmony_str = ( + "<|channel|>commentary" + "<|message|>Let me look that up.<|end|>" + "<|start|>assistant<|channel|>final" + "<|message|>The answer is 42.<|end|>" + ) + token_ids = get_encoding().encode(harmony_str, allowed_special="all") + reasoning, final_content, _ = parse_chat_output(token_ids) + assert reasoning is None + assert final_content == "Let me look that up.\nThe answer is 42." + class TestParseOutputMessage: """Tests for parse_output_message function.""" - def test_commentary_with_no_recipient_creates_reasoning(self): - """Test that commentary with recipient=None (preambles) creates reasoning items. + def test_commentary_with_no_recipient_creates_message(self): + """Test that commentary with recipient=None (preambles) creates message items. - Per Harmony format, commentary channel can contain preambles to calling - multiple functions - explanatory text with no recipient. + Per Harmony format, preambles are intended to be shown to end-users, + unlike analysis channel content which is hidden reasoning. + See: https://cookbook.openai.com/articles/openai-harmony """ message = Message.from_role_and_content( Role.ASSISTANT, "I will now search for the weather information." @@ -859,13 +907,16 @@ class TestParseOutputMessage: output_items = parse_output_message(message) assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseReasoningItem) - assert output_items[0].type == "reasoning" + assert isinstance(output_items[0], ResponseOutputMessage) + assert output_items[0].type == "message" + assert output_items[0].role == "assistant" + assert output_items[0].status == "completed" + assert len(output_items[0].content) == 1 + assert output_items[0].content[0].type == "output_text" assert ( output_items[0].content[0].text == "I will now search for the weather information." ) - assert output_items[0].content[0].type == "reasoning_text" def test_commentary_with_function_recipient_creates_function_call(self): """Test commentary with recipient='functions.X' creates function calls.""" @@ -944,7 +995,7 @@ class TestParseOutputMessage: output_items = parse_output_message(message) assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseReasoningItem) + assert isinstance(output_items[0], ResponseOutputMessage) assert output_items[0].content[0].text == "" def test_commentary_with_multiple_contents_and_no_recipient(self): @@ -958,10 +1009,13 @@ class TestParseOutputMessage: output_items = parse_output_message(message) - assert len(output_items) == 2 - assert all(isinstance(item, ResponseReasoningItem) for item in output_items) + # _parse_final_message returns single ResponseOutputMessage with + # multiple contents + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseOutputMessage) + assert len(output_items[0].content) == 2 assert output_items[0].content[0].text == "Step 1: Analyze the request" - assert output_items[1].content[0].text == "Step 2: Prepare to call functions" + assert output_items[0].content[1].text == "Step 2: Prepare to call functions" def test_commentary_with_multiple_function_calls(self): """Test multiple function calls in commentary channel.""" @@ -1133,7 +1187,7 @@ def test_parse_remaining_state_commentary_channel() -> None: assert mcp_items[0].status == "in_progress" # Test 3: Built-in tool (python) - # should NOT return MCP call, falls through to reasoning + # should NOT return MCP call, returns reasoning (internal tool interaction) parser_builtin = Mock() parser_builtin.current_content = "print('hello')" parser_builtin.current_role = Role.ASSISTANT @@ -1142,11 +1196,26 @@ def test_parse_remaining_state_commentary_channel() -> None: builtin_items = parse_remaining_state(parser_builtin) - # Should fall through to reasoning logic + # Built-in tools explicitly return reasoning assert len(builtin_items) == 1 assert not isinstance(builtin_items[0], McpCall) assert builtin_items[0].type == "reasoning" + # Test 4: No recipient (preamble) → should return message, not reasoning + parser_preamble = Mock() + parser_preamble.current_content = "I'll search for that information now." + parser_preamble.current_role = Role.ASSISTANT + parser_preamble.current_channel = "commentary" + parser_preamble.current_recipient = None + + preamble_items = parse_remaining_state(parser_preamble) + + assert len(preamble_items) == 1 + assert isinstance(preamble_items[0], ResponseOutputMessage) + assert preamble_items[0].type == "message" + assert preamble_items[0].content[0].text == "I'll search for that information now." + assert preamble_items[0].status == "incomplete" # streaming + def test_parse_remaining_state_analysis_channel() -> None: """Test parse_remaining_state with analysis channel and various recipients.""" @@ -1199,3 +1268,29 @@ def test_parse_remaining_state_analysis_channel() -> None: assert len(builtin_items) == 1 assert not isinstance(builtin_items[0], McpCall) assert builtin_items[0].type == "reasoning" + + +class TestGetSystemMessage: + """Tests for get_system_message channel configuration.""" + + def test_commentary_channel_present_without_custom_tools(self) -> None: + """Commentary channel must be valid even without custom tools.""" + sys_msg = get_system_message(with_custom_tools=False) + valid_channels = sys_msg.content[0].channel_config.valid_channels + assert "commentary" in valid_channels + + def test_commentary_channel_present_with_custom_tools(self) -> None: + """Commentary channel present when custom tools are enabled.""" + sys_msg = get_system_message(with_custom_tools=True) + valid_channels = sys_msg.content[0].channel_config.valid_channels + assert "commentary" in valid_channels + + def test_all_standard_channels_present(self) -> None: + """All three standard Harmony channels should always be valid.""" + for with_tools in (True, False): + sys_msg = get_system_message(with_custom_tools=with_tools) + valid_channels = sys_msg.content[0].channel_config.valid_channels + for channel in ("analysis", "commentary", "final"): + assert channel in valid_channels, ( + f"{channel} missing when with_custom_tools={with_tools}" + ) diff --git a/tests/entrypoints/openai/responses/test_harmony.py b/tests/entrypoints/openai/responses/test_harmony.py index af7de202669..78419c92a9d 100644 --- a/tests/entrypoints/openai/responses/test_harmony.py +++ b/tests/entrypoints/openai/responses/test_harmony.py @@ -712,15 +712,14 @@ async def test_function_calling_required(client: OpenAI, model_name: str): async def test_system_message_with_tools(client: OpenAI, model_name: str): from vllm.entrypoints.openai.parser.harmony_utils import get_system_message - # Test with custom tools enabled - commentary channel should be available - sys_msg = get_system_message(with_custom_tools=True) - valid_channels = sys_msg.content[0].channel_config.valid_channels - assert "commentary" in valid_channels - - # Test with custom tools disabled - commentary channel should be removed - sys_msg = get_system_message(with_custom_tools=False) - valid_channels = sys_msg.content[0].channel_config.valid_channels - assert "commentary" not in valid_channels + # Commentary channel should always be present (needed for preambles) + # regardless of whether custom tools are enabled + for with_tools in (True, False): + sys_msg = get_system_message(with_custom_tools=with_tools) + valid_channels = sys_msg.content[0].channel_config.valid_channels + assert "commentary" in valid_channels, ( + f"commentary channel missing when with_custom_tools={with_tools}" + ) @pytest.mark.asyncio diff --git a/tests/entrypoints/openai/responses/test_mcp_tools.py b/tests/entrypoints/openai/responses/test_mcp_tools.py index add199b6199..310af4308c6 100644 --- a/tests/entrypoints/openai/responses/test_mcp_tools.py +++ b/tests/entrypoints/openai/responses/test_mcp_tools.py @@ -172,13 +172,13 @@ class TestMCPEnabled: recipient = message.get("recipient") if recipient and recipient.startswith("python"): tool_call_found = True - assert message.get("channel") == "analysis" + assert message.get("channel") == "commentary" author = message.get("author", {}) if author.get("role") == "tool" and (author.get("name") or "").startswith( "python" ): tool_response_found = True - assert message.get("channel") == "analysis" + assert message.get("channel") == "commentary" assert tool_call_found, ( f"No Python tool call found. " diff --git a/tests/entrypoints/openai/test_serving_chat_stream_harmony.py b/tests/entrypoints/openai/test_serving_chat_stream_harmony.py index 21d3d02ce71..9f8c36f0473 100644 --- a/tests/entrypoints/openai/test_serving_chat_stream_harmony.py +++ b/tests/entrypoints/openai/test_serving_chat_stream_harmony.py @@ -180,20 +180,13 @@ class TestExtractHarmonyStreamingDelta: assert delta_message.tool_calls[0].index == 1 - @pytest.mark.parametrize( - "channel,recipient", - [ - ("commentary", None), - ("commentary", "browser.search"), - ], - ) - def test_returns_tool_call_preambles(self, channel, recipient): - """Test that invalid tool recipient on commentary is treated as content.""" + def test_returns_preambles_as_content(self): + """Test that commentary with no recipient (preamble) is user content.""" parser = MockStreamableParser() delta_text = "some text" token_states = [ - TokenState(channel=channel, recipient=recipient, text=delta_text) + TokenState(channel="commentary", recipient=None, text=delta_text) ] delta_message, tools_streamed = extract_harmony_streaming_delta( @@ -211,6 +204,7 @@ class TestExtractHarmonyStreamingDelta: [ (None, None), ("unknown_channel", None), + ("commentary", "browser.search"), ], ) def test_returns_none_for_invalid_inputs(self, channel, recipient): diff --git a/tests/entrypoints/openai/test_serving_responses.py b/tests/entrypoints/openai/test_serving_responses.py index 5cf07ac0f6a..291bfd442fa 100644 --- a/tests/entrypoints/openai/test_serving_responses.py +++ b/tests/entrypoints/openai/test_serving_responses.py @@ -26,6 +26,9 @@ from vllm.entrypoints.openai.responses.serving import ( _extract_allowed_tools_from_mcp_requests, extract_tool_types, ) +from vllm.entrypoints.openai.responses.streaming_events import ( + StreamingState, +) from vllm.inputs.data import TokensPrompt from vllm.outputs import CompletionOutput, RequestOutput from vllm.sampling_params import SamplingParams @@ -439,3 +442,115 @@ class TestExtractAllowedToolsFromMcpRequests: "server1": ["tool1"], "server2": ["tool2"], } + + +class TestHarmonyPreambleStreaming: + """Tests for preamble (commentary with no recipient) streaming events.""" + + @staticmethod + def _make_ctx(*, channel, recipient, delta="hello"): + """Build a lightweight mock StreamingHarmonyContext.""" + ctx = MagicMock() + ctx.last_content_delta = delta + ctx.parser.current_channel = channel + ctx.parser.current_recipient = recipient + return ctx + + @staticmethod + def _make_previous_item(*, channel, recipient, text="preamble text"): + """Build a lightweight mock previous_item (openai_harmony Message).""" + content_part = MagicMock() + content_part.text = text + item = MagicMock() + item.channel = channel + item.recipient = recipient + item.content = [content_part] + return item + + def test_preamble_delta_emits_text_events(self) -> None: + """commentary + recipient=None should emit output_text.delta events.""" + from vllm.entrypoints.openai.responses.streaming_events import ( + emit_content_delta_events, + ) + + ctx = self._make_ctx(channel="commentary", recipient=None) + state = StreamingState() + + events = emit_content_delta_events(ctx, state) + + type_names = [e.type for e in events] + assert "response.output_text.delta" in type_names + assert "response.output_item.added" in type_names + + def test_preamble_delta_second_token_no_added(self) -> None: + """Second preamble token should emit delta only, not added again.""" + from vllm.entrypoints.openai.responses.streaming_events import ( + emit_content_delta_events, + ) + + ctx = self._make_ctx(channel="commentary", recipient=None, delta="w") + state = StreamingState() + state.sent_output_item_added = True + state.current_item_id = "msg_test" + state.current_content_index = 0 + + events = emit_content_delta_events(ctx, state) + + type_names = [e.type for e in events] + assert "response.output_text.delta" in type_names + assert "response.output_item.added" not in type_names + + def test_commentary_with_function_recipient_not_preamble(self) -> None: + """commentary + recipient='functions.X' must NOT use preamble path.""" + from vllm.entrypoints.openai.responses.streaming_events import ( + emit_content_delta_events, + ) + + ctx = self._make_ctx( + channel="commentary", + recipient="functions.get_weather", + ) + state = StreamingState() + + events = emit_content_delta_events(ctx, state) + + type_names = [e.type for e in events] + assert "response.output_text.delta" not in type_names + + def test_preamble_done_emits_text_done_events(self) -> None: + """Completed preamble should emit text done + content_part done + + output_item done, same shape as final channel.""" + from vllm.entrypoints.openai.responses.streaming_events import ( + emit_previous_item_done_events, + ) + + previous = self._make_previous_item(channel="commentary", recipient=None) + state = StreamingState() + state.current_item_id = "msg_test" + state.current_output_index = 0 + state.current_content_index = 0 + + events = emit_previous_item_done_events(previous, state) + + type_names = [e.type for e in events] + assert "response.output_text.done" in type_names + assert "response.content_part.done" in type_names + assert "response.output_item.done" in type_names + + def test_commentary_with_recipient_no_preamble_done(self) -> None: + """commentary + recipient='functions.X' should route to function call + done, not preamble done.""" + from vllm.entrypoints.openai.responses.streaming_events import ( + emit_previous_item_done_events, + ) + + previous = self._make_previous_item( + channel="commentary", recipient="functions.get_weather" + ) + state = StreamingState() + state.current_item_id = "fc_test" + + events = emit_previous_item_done_events(previous, state) + + type_names = [e.type for e in events] + assert "response.output_text.done" not in type_names diff --git a/tests/entrypoints/test_context.py b/tests/entrypoints/test_context.py index 1ab2b5edb6e..b1c8df4fac3 100644 --- a/tests/entrypoints/test_context.py +++ b/tests/entrypoints/test_context.py @@ -236,6 +236,44 @@ def test_reasoning_tokens_counting(mock_parser): assert context.num_output_tokens == 4 +def test_preamble_tokens_not_counted_as_reasoning(mock_parser): + """Preambles (commentary with no recipient) are visible user text, + not hidden reasoning. They must NOT inflate num_reasoning_tokens.""" + context = HarmonyContext(messages=[], available_tools=[]) + + mock_parser.current_channel = "commentary" + mock_parser.current_recipient = None # preamble + + mock_output = create_mock_request_output( + prompt_token_ids=[1, 2, 3], + output_token_ids=[4, 5, 6], + num_cached_tokens=0, + ) + context.append_output(mock_output) + + assert context.num_reasoning_tokens == 0 + assert context.num_output_tokens == 3 + + +def test_commentary_with_recipient_counted_as_reasoning(mock_parser): + """Commentary directed at a tool (recipient != None) is hidden from + the user, so it should still count as reasoning tokens.""" + context = HarmonyContext(messages=[], available_tools=[]) + + mock_parser.current_channel = "commentary" + mock_parser.current_recipient = "python" + + mock_output = create_mock_request_output( + prompt_token_ids=[1, 2, 3], + output_token_ids=[4, 5, 6], + num_cached_tokens=0, + ) + context.append_output(mock_output) + + assert context.num_reasoning_tokens == 3 + assert context.num_output_tokens == 3 + + def test_zero_tokens_edge_case(): """Test behavior with all zero token counts.""" context = HarmonyContext(messages=[], available_tools=[]) diff --git a/vllm/entrypoints/openai/chat_completion/stream_harmony.py b/vllm/entrypoints/openai/chat_completion/stream_harmony.py index 4dbdddd20e6..87f2f9b9227 100644 --- a/vllm/entrypoints/openai/chat_completion/stream_harmony.py +++ b/vllm/entrypoints/openai/chat_completion/stream_harmony.py @@ -147,7 +147,7 @@ def extract_harmony_streaming_delta( function=DeltaFunctionCall(arguments=group.text), ) ) - elif group.channel == "commentary": + elif group.channel == "commentary" and group.recipient is None: # Tool call preambles meant to be shown to the user combined_content += group.text content_encountered = True diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index 486873db809..9dfd5f518f7 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -26,7 +26,6 @@ from openai.types.responses.response_reasoning_item import ( from openai.types.responses.tool import Tool from openai_harmony import ( Author, - ChannelConfig, Conversation, DeveloperContent, HarmonyEncodingName, @@ -126,13 +125,6 @@ def get_system_message( sys_msg_content = sys_msg_content.with_tools(python_description) if container_description is not None: sys_msg_content = sys_msg_content.with_tools(container_description) - if not with_custom_tools: - channel_config = sys_msg_content.channel_config - invalid_channel = "commentary" - new_config = ChannelConfig.require_channels( - [c for c in channel_config.valid_channels if c != invalid_channel] - ) - sys_msg_content = sys_msg_content.with_channel_config(new_config) sys_msg = Message.from_role_and_content(Role.SYSTEM, sys_msg_content) return sys_msg @@ -686,6 +678,22 @@ def _parse_mcp_call(message: Message, recipient: str) -> list[ResponseOutputItem return output_items +def _parse_message_no_recipient( + message: Message, +) -> list[ResponseOutputItem]: + """Parse a Harmony message with no recipient based on its channel.""" + if message.channel == "analysis": + return _parse_reasoning(message) + + if message.channel in ("commentary", "final"): + # Per Harmony format, preambles (commentary with no recipient) and + # final channel content are both intended to be shown to end-users. + # See: https://cookbook.openai.com/articles/openai-harmony + return [_parse_final_message(message)] + + raise ValueError(f"Unknown channel: {message.channel}") + + def parse_output_message(message: Message) -> list[ResponseOutputItem]: """ Parse a Harmony message into a list of output response items. @@ -717,19 +725,8 @@ def parse_output_message(message: Message) -> list[ResponseOutputItem]: output_items.extend(_parse_mcp_call(message, recipient)) # No recipient - handle based on channel for non-tool messages - elif message.channel == "analysis": - output_items.extend(_parse_reasoning(message)) - - elif message.channel == "commentary": - # Per Harmony format, commentary channel can contain preambles to calling - # multiple functions - explanatory text with no recipient - output_items.extend(_parse_reasoning(message)) - - elif message.channel == "final": - output_items.append(_parse_final_message(message)) - else: - raise ValueError(f"Unknown channel: {message.channel}") + output_items.extend(_parse_message_no_recipient(message)) return output_items @@ -786,7 +783,26 @@ def parse_remaining_state(parser: StreamableParser) -> list[ResponseOutputItem]: ) ] - if parser.current_channel in ("commentary", "analysis"): + if parser.current_channel == "commentary": + # Per Harmony format, preambles (commentary with no recipient) are + # intended to be shown to end-users, unlike analysis channel content. + output_text = ResponseOutputText( + text=parser.current_content, + annotations=[], + type="output_text", + logprobs=None, + ) + return [ + ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=[output_text], + role="assistant", + status="incomplete", + type="message", + ) + ] + + if parser.current_channel == "analysis": return [ ResponseReasoningItem( id=f"rs_{random_uuid()}", @@ -855,17 +871,30 @@ def parse_chat_output( is_tool_call = False # TODO: update this when tool call is supported # Get completed messages from the parser + # - analysis channel: hidden reasoning + # - commentary channel without recipient (preambles): visible to user + # - final channel: visible to user + # - commentary with recipient (tool calls): handled separately by tool parser reasoning_texts = [ msg.content[0].text for msg in output_msgs if msg.channel == "analysis" ] final_texts = [ - msg.content[0].text for msg in output_msgs if msg.channel != "analysis" + msg.content[0].text + for msg in output_msgs + if msg.channel == "final" or (msg.channel == "commentary" and not msg.recipient) ] # Extract partial messages from the parser if parser.current_channel == "analysis" and parser.current_content: reasoning_texts.append(parser.current_content) - elif parser.current_channel != "analysis" and parser.current_content: + elif parser.current_channel == "final" and parser.current_content: + final_texts.append(parser.current_content) + elif ( + parser.current_channel == "commentary" + and not parser.current_recipient + and parser.current_content + ): + # Preambles (commentary without recipient) are visible to user final_texts.append(parser.current_content) # Flatten multiple messages into a single string diff --git a/vllm/entrypoints/openai/responses/context.py b/vllm/entrypoints/openai/responses/context.py index b57adeeb846..bab59e0aa1e 100644 --- a/vllm/entrypoints/openai/responses/context.py +++ b/vllm/entrypoints/openai/responses/context.py @@ -540,8 +540,12 @@ class HarmonyContext(ConversationContext): self.first_tok_of_message = True # For streaming support def _update_num_reasoning_tokens(self): - # Count all analysis and commentary channels as reasoning tokens - if self.parser.current_channel in {"analysis", "commentary"}: + channel = self.parser.current_channel + if channel == "analysis": + self.num_reasoning_tokens += 1 + elif channel == "commentary" and self.parser.current_recipient is not None: + # Tool interactions (python/browser/container) are hidden. + # Preambles (recipient=None) are visible user text. self.num_reasoning_tokens += 1 def append_output(self, output: RequestOutput) -> None: diff --git a/vllm/entrypoints/openai/responses/streaming_events.py b/vllm/entrypoints/openai/responses/streaming_events.py index 49d2b99dace..cc242e7baa8 100644 --- a/vllm/entrypoints/openai/responses/streaming_events.py +++ b/vllm/entrypoints/openai/responses/streaming_events.py @@ -563,7 +563,9 @@ def emit_content_delta_events( channel = ctx.parser.current_channel recipient = ctx.parser.current_recipient - if channel == "final" and recipient is None: + if channel in ("final", "commentary") and recipient is None: + # Preambles (commentary with no recipient) and final messages + # are both user-visible text. return emit_text_delta_events(delta, state) elif channel == "analysis" and recipient is None: return emit_reasoning_delta_events(delta, state) @@ -607,7 +609,9 @@ def emit_previous_item_done_events( return emit_mcp_completion_events(previous_item.recipient, text, state) elif previous_item.channel == "analysis": return emit_reasoning_done_events(text, state) - elif previous_item.channel == "final": + elif previous_item.channel in ("commentary", "final"): + # Preambles (commentary with no recipient) and final messages + # are both user-visible text. return emit_text_output_done_events(text, state) return [] From 8fd69754798c50b7b07938451bd97b1e66765927 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 25 Feb 2026 10:48:37 -0600 Subject: [PATCH 05/43] [ROCm][CI] Disable skinny GEMMs in multimodal tests to fix non-deterministic results (#35049) Signed-off-by: Andreas Karatzas --- tests/models/multimodal/conftest.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/tests/models/multimodal/conftest.py b/tests/models/multimodal/conftest.py index 3f53b3fe629..d00c3df786d 100644 --- a/tests/models/multimodal/conftest.py +++ b/tests/models/multimodal/conftest.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Pytest configuration for vLLM multimodal tests.""" +import os import warnings import torch @@ -9,6 +10,23 @@ import torch from vllm.platforms import current_platform +def pytest_configure(config): + """Early ROCm configuration that must happen before test collection.""" + if not current_platform.is_rocm(): + return + + # Disable skinny GEMM on ROCm to avoid non-deterministic results + # from atomic reductions in wvSplitKrc kernel. + # See: https://github.com/vllm-project/vllm/pull/33493#issuecomment-3906083975 + os.environ["VLLM_ROCM_USE_SKINNY_GEMM"] = "0" + warnings.warn( + "ROCm: Set VLLM_ROCM_USE_SKINNY_GEMM=0 to avoid non-deterministic " + "results from skinny GEMM atomic reductions", + UserWarning, + stacklevel=1, + ) + + def pytest_collection_modifyitems(config, items): """Configure ROCm-specific settings based on collected tests.""" if not current_platform.is_rocm(): From 15d76f74e2fdb12a95ea00f0ca283acf6219a2b7 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Wed, 25 Feb 2026 12:20:15 -0500 Subject: [PATCH 06/43] Revert "[Misc] Enable weights loading tracking for quantized models" (#35309) --- .../model_loader/default_loader.py | 19 ++++--------------- 1 file changed, 4 insertions(+), 15 deletions(-) diff --git a/vllm/model_executor/model_loader/default_loader.py b/vllm/model_executor/model_loader/default_loader.py index ed201630d25..7064998af86 100644 --- a/vllm/model_executor/model_loader/default_loader.py +++ b/vllm/model_executor/model_loader/default_loader.py @@ -14,7 +14,6 @@ from transformers.utils import SAFE_WEIGHTS_INDEX_NAME from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger -from vllm.model_executor.layers.quantization.kv_cache import BaseKVCacheMethod from vllm.model_executor.layers.quantization.torchao import torchao_version_at_least from vllm.model_executor.model_loader.base_loader import BaseModelLoader from vllm.model_executor.model_loader.weight_utils import ( @@ -287,6 +286,7 @@ class DefaultModelLoader(BaseModelLoader): ): self.load_config.safetensors_load_strategy = "torchao" + weights_to_load = {name for name, _ in model.named_parameters()} loaded_weights = model.load_weights(self.get_all_weights(model_config, model)) self.counter_after_loading_weights = time.perf_counter() @@ -295,20 +295,9 @@ class DefaultModelLoader(BaseModelLoader): self.counter_after_loading_weights - self.counter_before_loading_weights, scope="local", ) - self.track_weights_loading(model, loaded_weights) - - def track_weights_loading( - self, model: nn.Module, loaded_weights: set[str] | None - ) -> None: - weights_to_load = {name for name, _ in model.named_parameters()} - if loaded_weights is not None: - for name, module in model.named_modules(): - quant_method = getattr(module, "quant_method", None) - # ignore kv_cache scale, which can be missing in checkpoints - if isinstance(quant_method, BaseKVCacheMethod): - for param_name, _ in module.named_parameters(): - full_name = f"{name}.{param_name}" if name else param_name - loaded_weights.add(full_name) + # We only enable strict check for non-quantized models + # that have loaded weights tracking currently. + if model_config.quantization is None and loaded_weights is not None: weights_not_loaded = weights_to_load - loaded_weights if weights_not_loaded: raise ValueError( From b188bab4417f7f94df0c4e84399d163e6c0db316 Mon Sep 17 00:00:00 2001 From: rasmith Date: Wed, 25 Feb 2026 13:18:00 -0600 Subject: [PATCH 07/43] [CI][AMD][BugFix] Add torch.cuda.set_device to test_punica_ops so punica kernels execute on same device as tensor (#34985) Signed-off-by: Randall Smith --- tests/lora/test_punica_ops.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/lora/test_punica_ops.py b/tests/lora/test_punica_ops.py index 96326036767..82db7fece3f 100644 --- a/tests/lora/test_punica_ops.py +++ b/tests/lora/test_punica_ops.py @@ -395,6 +395,7 @@ def test_kernels( Tests LoRA kernels. """ torch.set_default_device(device) + torch.cuda.set_device(device) set_random_seed(seed) if op_type == "shrink": @@ -447,6 +448,7 @@ def test_kernels_hidden_size( Tests SGMV and LoRA kernels. """ torch.set_default_device(device) + torch.cuda.set_device(device) set_random_seed(seed) if op_type == "shrink": From c97234c08b42326cf1e5ef024d9ac8441e0848b1 Mon Sep 17 00:00:00 2001 From: Elizabeth Thomas Date: Wed, 25 Feb 2026 15:33:42 -0600 Subject: [PATCH 08/43] fix(mxfp4): Disable monolithic path for TRITON backend with EP (#34270) Signed-off-by: Elizabeth Thomas Co-authored-by: Claude Opus 4.6 Co-authored-by: Michael Goin --- .../quantization/test_mxfp4_triton_ep.py | 194 ++++++++++++++++++ .../fused_moe/gpt_oss_triton_kernels_moe.py | 36 +++- 2 files changed, 225 insertions(+), 5 deletions(-) create mode 100644 tests/kernels/quantization/test_mxfp4_triton_ep.py diff --git a/tests/kernels/quantization/test_mxfp4_triton_ep.py b/tests/kernels/quantization/test_mxfp4_triton_ep.py new file mode 100644 index 00000000000..d4eb9105890 --- /dev/null +++ b/tests/kernels/quantization/test_mxfp4_triton_ep.py @@ -0,0 +1,194 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Tests that triton_kernel_moe_forward correctly applies expert_map +remapping when expert parallelism (EP) is enabled. + +Previously, legacy_routing was always used and it produced routing data +with global expert IDs that didn't correspond to local weight indices, +causing illegal memory access with EP. The fix splits routing: when +expert_map is provided, topk selection is performed first, expert_map is +applied to remap global→local IDs, and make_routing_data builds routing +structures from the local IDs. +""" + +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm.model_executor.layers.quantization.mxfp4 import ( + Mxfp4Backend, + Mxfp4MoEMethod, +) + + +def _make_mock_moe_config(ep_size: int = 1) -> MagicMock: + """Create a mock FusedMoEConfig with the given EP size.""" + parallel_config = MagicMock() + parallel_config.ep_size = ep_size + + moe_config = MagicMock() + moe_config.ep_size = ep_size + moe_config.is_lora_enabled = False + moe_config.moe_parallel_config = parallel_config + return moe_config + + +class TestMxfp4TritonIsMonolithic: + """Verify that is_monolithic is always True for the TRITON backend, + regardless of EP size, since triton_kernel_moe_forward now handles + expert_map remapping internally.""" + + @pytest.mark.parametrize( + "backend,ep_size,expected_monolithic", + [ + # TRITON is always monolithic (handles EP via expert_map remapping) + (Mxfp4Backend.TRITON, 1, True), + (Mxfp4Backend.TRITON, 2, True), + (Mxfp4Backend.TRITON, 4, True), + # SM100 backends are always monolithic + (Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM, 1, True), + (Mxfp4Backend.SM100_FI_MXFP4_MXFP8_TRTLLM, 2, True), + (Mxfp4Backend.SM100_FI_MXFP4_BF16, 1, True), + (Mxfp4Backend.SM100_FI_MXFP4_BF16, 2, True), + # MARLIN is never monolithic + (Mxfp4Backend.MARLIN, 1, False), + (Mxfp4Backend.MARLIN, 2, False), + ], + ids=[ + "triton-no-ep", + "triton-ep2", + "triton-ep4", + "sm100-trtllm-no-ep", + "sm100-trtllm-ep2", + "sm100-bf16-no-ep", + "sm100-bf16-ep2", + "marlin-no-ep", + "marlin-ep2", + ], + ) + @patch( + "vllm.model_executor.layers.quantization.mxfp4.get_mxfp4_backend", + ) + @patch( + "vllm.model_executor.layers.quantization.mxfp4.get_current_vllm_config", + ) + def test_is_monolithic( + self, + mock_get_config, + mock_get_backend, + backend, + ep_size, + expected_monolithic, + ): + """is_monolithic should be True for TRITON regardless of EP size.""" + mock_get_backend.return_value = backend + + mock_compilation_config = MagicMock() + mock_compilation_config.max_cudagraph_capture_size = 1024 + mock_vllm_config = MagicMock() + mock_vllm_config.compilation_config = mock_compilation_config + mock_get_config.return_value = mock_vllm_config + + moe_config = _make_mock_moe_config(ep_size=ep_size) + method = Mxfp4MoEMethod(moe_config) + + assert method.is_monolithic == expected_monolithic, ( + f"Expected is_monolithic={expected_monolithic} for " + f"backend={backend.name}, ep_size={ep_size}, " + f"but got {method.is_monolithic}." + ) + + +class TestTritonMoeForwardExpertMap: + """Test that triton_kernel_moe_forward applies expert_map remapping + when expert_map is provided (EP active).""" + + @pytest.mark.parametrize("expert_map_present", [False, True]) + def test_routing_path_selection(self, expert_map_present): + """Verify that the EP-aware routing path is taken when expert_map + is present, and the legacy_routing path is taken otherwise.""" + + device = "cuda" if torch.cuda.is_available() else "cpu" + # This is a structural test: we mock the routing functions to + # verify the correct path is exercised. + mock_expert_map = ( + torch.tensor([0, -1, 1, -1], device=device) if expert_map_present else None + ) + + with ( + patch( + "vllm.model_executor.layers.fused_moe." + "gpt_oss_triton_kernels_moe.legacy_routing" + ) as mock_legacy, + patch("triton_kernels.topk.topk") as mock_topk, + patch( + "vllm.model_executor.layers.fused_moe." + "gpt_oss_triton_kernels_moe.make_routing_data" + ) as mock_make_routing, + patch( + "vllm.model_executor.layers.fused_moe." + "gpt_oss_triton_kernels_moe.triton_kernel_fused_experts" + ) as mock_fused_experts, + ): + from vllm.model_executor.layers.fused_moe.gpt_oss_triton_kernels_moe import ( # noqa: E501 + triton_kernel_moe_forward, + ) + + # Set up return values + mock_routing_data = MagicMock() + mock_gather = MagicMock() + mock_scatter = MagicMock() + + if expert_map_present: + sparse_result = MagicMock() + sparse_result.indx = torch.tensor([[0, 2]], dtype=torch.int32) + sparse_result.vals = torch.tensor([[0.6, 0.4]]) + mock_topk.return_value = sparse_result + mock_make_routing.return_value = ( + mock_routing_data, + mock_gather, + mock_scatter, + ) + else: + mock_legacy.return_value = ( + mock_routing_data, + mock_gather, + mock_scatter, + ) + + mock_fused_experts.return_value = torch.zeros((1, 8), device=device) + + hidden = torch.randn((1, 8), device=device) + w1 = torch.randn((2, 8, 16), device=device) + w2 = torch.randn((2, 8, 8), device=device) + logits = torch.randn((1, 4), device=device) + + triton_kernel_moe_forward( + hidden_states=hidden, + w1=w1, + w2=w2, + gating_output=logits, + topk=2, + renormalize=True, + expert_map=mock_expert_map, + ) + + if expert_map_present: + # EP path: should use topk + make_routing_data, NOT + # legacy_routing + mock_topk.assert_called_once() + mock_make_routing.assert_called_once() + mock_legacy.assert_not_called() + # expert_map should be None in the fused_experts call + # (already applied) + call_kwargs = mock_fused_experts.call_args + assert call_kwargs[1].get("expert_map") is None or ( + len(call_kwargs[0]) > 0 + ) + else: + # Non-EP path: should use legacy_routing + mock_legacy.assert_called_once() + mock_topk.assert_not_called() + mock_make_routing.assert_not_called() diff --git a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py index 70d11f44f43..5617156bf2f 100644 --- a/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py +++ b/vllm/model_executor/layers/fused_moe/gpt_oss_triton_kernels_moe.py @@ -179,9 +179,35 @@ def triton_kernel_moe_forward( global_num_experts: int = -1, expert_map: torch.Tensor | None = None, ) -> torch.Tensor: - routing_data, gather_idx, scatter_idx = legacy_routing( - gating_output, topk, sm_first=not renormalize - ) + if expert_map is not None: + # With expert parallelism, legacy_routing produces routing data + # using global expert IDs which don't correspond to local weight + # indices. Split the routing into topk selection + expert_map + # remapping + local routing data construction (matching the + # approach used by OAITritonExperts.apply). + from triton_kernels.topk import topk as topk_fn + + sm_first = not renormalize + logits = gating_output + if sm_first: + logits = torch.softmax(logits, dim=-1) + sparse_logits = topk_fn(logits, topk, apply_softmax=not sm_first) + # sparse_logits.indx contains global expert IDs – remap to local. + topk_ids = expert_map[sparse_logits.indx.to(torch.long)] + topk_weights = sparse_logits.vals + local_num_experts = w1.size(0) + routing_data, gather_idx, scatter_idx = make_routing_data( + topk_ids, topk_weights, local_num_experts + ) + # expert_map already applied; pass None downstream. + effective_expert_map = None + effective_global_num_experts = local_num_experts + else: + routing_data, gather_idx, scatter_idx = legacy_routing( + gating_output, topk, sm_first=not renormalize + ) + effective_expert_map = expert_map + effective_global_num_experts = global_num_experts output = torch.empty_like(hidden_states) @@ -197,8 +223,8 @@ def triton_kernel_moe_forward( activation=activation, quant_config=quant_config, apply_router_weight_on_input=apply_router_weight_on_input, - global_num_experts=global_num_experts, - expert_map=expert_map, + global_num_experts=effective_global_num_experts, + expert_map=effective_expert_map, ) From 9571e999451a468423e97cd7e3f36e9d27a098cb Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 25 Feb 2026 16:16:18 -0600 Subject: [PATCH 09/43] [ROCm][CI] Extending attention backend coverage for Eagle spec decode tests (#35265) Signed-off-by: Andreas Karatzas --- .buildkite/test_areas/engine.yaml | 2 +- tests/utils.py | 51 ++++ tests/v1/e2e/test_async_scheduling.py | 4 + tests/v1/e2e/test_spec_decode.py | 407 ++++++++++++++++---------- 4 files changed, 314 insertions(+), 150 deletions(-) diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index 4f2380592d9..19cd91370e6 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -30,7 +30,7 @@ steps: - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py mirror: amd: - device: mi325_8 + device: mi325_1 depends_on: - image-build-amd commands: diff --git a/tests/utils.py b/tests/utils.py index 75d33e50952..4041c261788 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1327,6 +1327,57 @@ def multi_gpu_test(*, num_gpus: int): return wrapper +def gpu_tier_mark(*, min_gpus: int = 1, max_gpus: int | None = None): + """ + Mark a test to only run when the GPU count falls within [min_gpus, max_gpus]. + + Examples: + @gpu_tier_mark(min_gpus=2) # only on multi-GPU + @gpu_tier_mark(max_gpus=1) # only on single-GPU + @gpu_tier_mark(min_gpus=2, max_gpus=4) # 2-4 GPUs only + """ + gpu_count = cuda_device_count_stateless() + marks = [] + + if min_gpus > 1: + marks.append(pytest.mark.distributed(num_gpus=min_gpus)) + + reasons = [] + if gpu_count < min_gpus: + reasons.append(f"Need at least {min_gpus} GPUs (have {gpu_count})") + if max_gpus is not None and gpu_count > max_gpus: + reasons.append(f"Need at most {max_gpus} GPUs (have {gpu_count})") + + if reasons: + marks.append(pytest.mark.skipif(True, reason="; ".join(reasons))) + + return marks + + +def single_gpu_only(f=None): + """Skip this test when running in a multi-GPU environment.""" + marks = gpu_tier_mark(max_gpus=1) + + def wrapper(func): + for mark in reversed(marks): + func = mark(func) + return func + + return wrapper(f) if f is not None else wrapper + + +def multi_gpu_only(*, num_gpus: int = 2): + """Skip this test when running on fewer than num_gpus GPUs.""" + marks = gpu_tier_mark(min_gpus=num_gpus) + + def wrapper(f): + for mark in reversed(marks): + f = mark(f) + return f + + return wrapper + + async def completions_with_server_args( prompts: list[str], model_name: str, diff --git a/tests/v1/e2e/test_async_scheduling.py b/tests/v1/e2e/test_async_scheduling.py index b85f8880cf8..393c8dbeecf 100644 --- a/tests/v1/e2e/test_async_scheduling.py +++ b/tests/v1/e2e/test_async_scheduling.py @@ -6,6 +6,7 @@ from typing import Any import pytest import torch._dynamo.config as dynamo_config +from tests.utils import large_gpu_mark, single_gpu_only from vllm import SamplingParams from vllm.logprobs import Logprob from vllm.platforms import current_platform @@ -36,6 +37,7 @@ default_params = dict( ) +@single_gpu_only def test_without_spec_decoding( sample_json_schema, monkeypatch: pytest.MonkeyPatch, @@ -95,6 +97,8 @@ def test_without_spec_decoding( run_tests(monkeypatch, MODEL, test_configs, test_sampling_params) +@single_gpu_only +@large_gpu_mark(min_gb=16) def test_with_spec_decoding(sample_json_schema, monkeypatch: pytest.MonkeyPatch): """Test consistency and acceptance rates with some different combos of preemption, executor, async scheduling, prefill chunking, diff --git a/tests/v1/e2e/test_spec_decode.py b/tests/v1/e2e/test_spec_decode.py index 9289d1ce131..7f2db19a075 100644 --- a/tests/v1/e2e/test_spec_decode.py +++ b/tests/v1/e2e/test_spec_decode.py @@ -9,7 +9,13 @@ import pytest import torch from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts, evaluate_gsm8k_offline -from tests.utils import get_attn_backend_list_based_on_platform, large_gpu_mark +from tests.utils import ( + get_attn_backend_list_based_on_platform, + large_gpu_mark, + multi_gpu_marks, + multi_gpu_only, + single_gpu_only, +) from vllm import LLM, SamplingParams from vllm.assets.base import VLLM_S3_BUCKET_URL from vllm.assets.image import VLM_IMAGES_DIR @@ -160,6 +166,8 @@ def reset_torch_dynamo(): }, ], ) +@single_gpu_only +@large_gpu_mark(min_gb=20) def test_ngram_and_suffix_correctness( speculative_config: dict, model_name: str, @@ -175,6 +183,8 @@ def test_ngram_and_suffix_correctness( cleanup_dist_env_and_memory() +@single_gpu_only +@large_gpu_mark(min_gb=20) def test_suffix_decoding_acceptance( monkeypatch: pytest.MonkeyPatch, sampling_config: SamplingParams, @@ -242,6 +252,8 @@ def test_suffix_decoding_acceptance( ], ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"], ) +@single_gpu_only +@large_gpu_mark(min_gb=24) def test_speculators_model_integration( monkeypatch: pytest.MonkeyPatch, sampling_config: SamplingParams, @@ -319,137 +331,7 @@ def test_speculators_model_integration( ) -@pytest.mark.parametrize( - [ - "model_setup", - "mm_enabled", - "enable_chunked_prefill", - "model_impl", - "expected_accuracy_threshold", - ], - [ - ( - ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), - False, - False, - "auto", - 0.8, # ref: 90% - ), - ( - ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), - False, - False, - "transformers", - 0.8, # ref: 90% - ), - pytest.param( - ( - "eagle3", - "Qwen/Qwen3-VL-8B-Instruct", - "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3", - 1, - ), - False, - False, - "auto", - 0.8, # ref: 90% - marks=pytest.mark.skip( - reason="architecture of its eagle3 is LlamaForCausalLMEagle3" - ), - ), - pytest.param( - ( - "eagle3", - "Qwen/Qwen2.5-VL-7B-Instruct", - "Rayzl/qwen2.5-vl-7b-eagle3-sgl", - 1, - ), - False, - False, - "auto", - 0.7, # TODO, update this with a reference value when re-enabling this case - marks=pytest.mark.skip( - reason="Skipping due to its head_dim not being a a multiple of 32" - ), - ), - pytest.param( - ( - "eagle", - "meta-llama/Llama-3.1-8B-Instruct", - "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", - 1, - ), - False, - True, - "auto", - 0.7, # ref: 75%-80% - marks=large_gpu_mark(min_gb=40), - ), # works on 4x H100 - ( - ( - "eagle3", - "meta-llama/Llama-3.1-8B-Instruct", - "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", - 1, - ), - False, - False, - "auto", - 0.7, # ref: 75%-80% - ), - pytest.param( - ( - "eagle", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", - 4, - ), - False, - False, - "auto", - 0.8, # ref: 90% - # marks=large_gpu_mark(min_gb=80), - ), # works on 4x H100 - pytest.param( - ( - "eagle", - "meta-llama/Llama-4-Scout-17B-16E-Instruct", - "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", - 4, - ), - True, - True, - "auto", - 0.8, # ref: 90% - marks=large_gpu_mark(min_gb=80), - ), # works on 4x H100 - ( - ( - "eagle", - "eagle618/deepseek-v3-random", - "eagle618/eagle-deepseek-v3-random", - 1, - ), - False, - False, - "auto", - 0.0, # dummy model, skip gsm8k check - ), - ], - ids=[ - "qwen3_eagle3", - "qwen3_eagle3-transformers", - "qwen3_vl_eagle3", - "qwen2_5_vl_eagle3", - "llama3_eagle", - "llama3_eagle3", - "llama4_eagle", - "llama4_eagle_mm", - "deepseek_eagle", - ], -) -@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) -def test_eagle_correctness( +def _run_eagle_correctness( monkeypatch: pytest.MonkeyPatch, sampling_config: SamplingParams, model_setup: tuple[str, str, str, int], @@ -460,14 +342,10 @@ def test_eagle_correctness( attn_backend: str, ): """ - Compare the outputs of a original LLM and a speculative LLM - which should be the same when using eagle speculative decoding. Due to some variance - in the engine, it is possible for some outputs to differ, so we expect that at least - 6/10 output tokens match exactly, and that the GSM8k accuracy is above - a precomputed reference threshold for each model. + Compare the outputs of an original LLM and a speculative LLM + which should be the same when using eagle speculative decoding. """ if attn_backend == "TREE_ATTN": - # TODO: Fix this flaky test pytest.skip( "TREE_ATTN is flaky in the test disable for now until it can be " "resolved (see https://github.com/vllm-project/vllm/issues/22922)" @@ -484,17 +362,17 @@ def test_eagle_correctness( f"transformers>={required}, but got {installed}" ) - # Generate test prompts inside the function instead of using fixture test_prompts = get_test_prompts(mm_enabled) - # Determine attention config - # Scout requires default backend selection because vision encoder has - # head_dim 88 being incompatible with FLASH_ATTN and needs to fall back - # to Flex Attn + if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN": if current_platform.is_rocm(): - # TODO: Enable Flex Attn for spec_decode on ROCm - pytest.skip("Flex Attn for spec_decode not supported on ROCm currently") - attention_config = None # Let it fall back to default + print( + "FLASH_ATTN for spec_decode not supported on " + "ROCm currently. Changing to FLEX_ATTENTION backend." + ) + attention_config = {"backend": "FLEX_ATTENTION"} + else: + attention_config = None else: attention_config = {"backend": attn_backend} @@ -509,7 +387,9 @@ def test_eagle_correctness( if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): if "deepseek" in model_setup[1].lower(): - pytest.skip("ROCM_AITER_FA for deepseek not supported on ROCm platform") + m.setenv("VLLM_ROCM_USE_AITER", "1") + m.delenv("VLLM_MLA_DISABLE", raising=False) + attention_config = {"backend": "TRITON_MLA"} else: m.setenv("VLLM_ROCM_USE_AITER", "1") @@ -563,14 +443,235 @@ def test_eagle_correctness( print(f"ref_output: {ref_output.outputs[0].text}") print(f"spec_output: {spec_output.outputs[0].text}") - # Heuristic: expect at least 60% of the prompts to match exactly - # Upon failure, inspect the outputs to check for inaccuracy. assert matches > int(0.6 * len(ref_outputs)) del spec_llm torch.cuda.empty_cache() cleanup_dist_env_and_memory() +@single_gpu_only +@pytest.mark.parametrize( + [ + "model_setup", + "mm_enabled", + "enable_chunked_prefill", + "model_impl", + "expected_accuracy_threshold", + ], + [ + ( + ( + "eagle", + "eagle618/deepseek-v3-random", + "eagle618/eagle-deepseek-v3-random", + 1, + ), + False, + False, + "auto", + 0.0, + ), + ], + ids=["deepseek_eagle"], +) +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) +def test_eagle_correctness_light( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, str, int], + mm_enabled: bool, + expected_accuracy_threshold: float, + enable_chunked_prefill: bool, + model_impl: str, + attn_backend: str, +): + _run_eagle_correctness( + monkeypatch, + sampling_config, + model_setup, + mm_enabled, + expected_accuracy_threshold, + enable_chunked_prefill, + model_impl, + attn_backend, + ) + + +@single_gpu_only +@large_gpu_mark(min_gb=24) +@pytest.mark.parametrize( + [ + "model_setup", + "mm_enabled", + "enable_chunked_prefill", + "model_impl", + "expected_accuracy_threshold", + ], + [ + ( + ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), + False, + False, + "auto", + 0.8, + ), + ( + ("eagle3", "Qwen/Qwen3-8B", "AngelSlim/Qwen3-8B_eagle3", 1), + False, + False, + "transformers", + 0.8, + ), + pytest.param( + ( + "eagle3", + "Qwen/Qwen3-VL-8B-Instruct", + "taobao-mnn/Qwen3-VL-8B-Instruct-Eagle3", + 1, + ), + False, + False, + "auto", + 0.8, + marks=pytest.mark.skip( + reason="architecture of its eagle3 is LlamaForCausalLMEagle3" + ), + ), + pytest.param( + ( + "eagle3", + "Qwen/Qwen2.5-VL-7B-Instruct", + "Rayzl/qwen2.5-vl-7b-eagle3-sgl", + 1, + ), + False, + False, + "auto", + 0.7, + marks=pytest.mark.skip( + reason="Skipping due to its head_dim not being a multiple of 32" + ), + ), + ( + ( + "eagle3", + "meta-llama/Llama-3.1-8B-Instruct", + "yuhuili/EAGLE3-LLaMA3.1-Instruct-8B", + 1, + ), + False, + False, + "auto", + 0.7, + ), + ], + ids=[ + "qwen3_eagle3", + "qwen3_eagle3-transformers", + "qwen3_vl_eagle3", + "qwen2_5_vl_eagle3", + "llama3_eagle3", + ], +) +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) +def test_eagle_correctness_medium( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, str, int], + mm_enabled: bool, + expected_accuracy_threshold: float, + enable_chunked_prefill: bool, + model_impl: str, + attn_backend: str, +): + _run_eagle_correctness( + monkeypatch, + sampling_config, + model_setup, + mm_enabled, + expected_accuracy_threshold, + enable_chunked_prefill, + model_impl, + attn_backend, + ) + + +@pytest.mark.parametrize( + [ + "model_setup", + "mm_enabled", + "enable_chunked_prefill", + "model_impl", + "expected_accuracy_threshold", + ], + [ + pytest.param( + ( + "eagle", + "meta-llama/Llama-3.1-8B-Instruct", + "yuhuili/EAGLE-LLaMA3.1-Instruct-8B", + 1, + ), + False, + True, + "auto", + 0.7, + marks=large_gpu_mark(min_gb=40), + id="llama3_eagle", + ), + pytest.param( + ( + "eagle", + "meta-llama/Llama-4-Scout-17B-16E-Instruct", + "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", + 4, + ), + False, + False, + "auto", + 0.8, + marks=multi_gpu_marks(num_gpus=4), + id="llama4_eagle", + ), + pytest.param( + ( + "eagle", + "meta-llama/Llama-4-Scout-17B-16E-Instruct", + "morgendave/EAGLE-Llama-4-Scout-17B-16E-Instruct", + 4, + ), + True, + True, + "auto", + 0.8, + marks=[*multi_gpu_marks(num_gpus=4), large_gpu_mark(min_gb=80)], + id="llama4_eagle_mm", + ), + ], +) +@pytest.mark.parametrize("attn_backend", get_attn_backend_list_based_on_platform()) +def test_eagle_correctness_heavy( + monkeypatch: pytest.MonkeyPatch, + sampling_config: SamplingParams, + model_setup: tuple[str, str, str, int], + mm_enabled: bool, + expected_accuracy_threshold: float, + enable_chunked_prefill: bool, + model_impl: str, + attn_backend: str, +): + _run_eagle_correctness( + monkeypatch, + sampling_config, + model_setup, + mm_enabled, + expected_accuracy_threshold, + enable_chunked_prefill, + model_impl, + attn_backend, + ) + + @pytest.mark.parametrize( ["model_setup", "mm_enabled", "expected_accuracy_threshold"], [ @@ -579,6 +680,8 @@ def test_eagle_correctness( ], ids=["mimo", "deepseek"], ) +@single_gpu_only +@large_gpu_mark(min_gb=20) def test_mtp_correctness( monkeypatch: pytest.MonkeyPatch, sampling_config: SamplingParams, @@ -694,11 +797,13 @@ cases = [ @pytest.mark.parametrize("args", cases) @pytest.mark.parametrize("enforce_eager", [True, False]) +@single_gpu_only def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool): args.enforce_eager = enforce_eager assert_draft_model_correctness(args) +@single_gpu_only def test_draft_model_realistic_example(): args = ArgsTest( target_model="Qwen/Qwen3-1.7B", @@ -713,6 +818,7 @@ def test_draft_model_realistic_example(): assert_draft_model_correctness(args) +@single_gpu_only def test_draft_model_parallel_drafting(): args = ArgsTest( target_model="Qwen/Qwen3-1.7B", @@ -738,6 +844,7 @@ def test_draft_model_parallel_drafting(): ids=["target_quantized", "draft_quantized"], ) @pytest.mark.parametrize("enforce_eager", [True, False]) +@single_gpu_only def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): tgt_model, draft_model = models sd_case = ArgsTest( @@ -749,6 +856,7 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): assert_draft_model_correctness(sd_case) +@multi_gpu_only(num_gpus=2) def test_draft_model_tensor_parallelism(): """Ensure spec decode works when running with TP > 1.""" _skip_if_insufficient_gpus_for_tp(2) @@ -764,6 +872,7 @@ def test_draft_model_tensor_parallelism(): assert_draft_model_correctness(sd_case) +@multi_gpu_only(num_gpus=2) def test_draft_model_engine_args_tensor_parallelism(): """Ensure the vllm_config for the draft model is created correctly, and independently of the target model (quantization, TP, etc.)""" From ed42507f6d6e326663997da5cca6991da5d8a23f Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Wed, 25 Feb 2026 16:17:56 -0600 Subject: [PATCH 10/43] [ROCm][CI] Amending deletion of AMD mirror (#35322) Signed-off-by: Andreas Karatzas --- .buildkite/test_areas/entrypoints.yaml | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/.buildkite/test_areas/entrypoints.yaml b/.buildkite/test_areas/entrypoints.yaml index 5c58e97ef16..17201a07103 100644 --- a/.buildkite/test_areas/entrypoints.yaml +++ b/.buildkite/test_areas/entrypoints.yaml @@ -24,6 +24,11 @@ steps: - pytest -v -s entrypoints/llm --ignore=entrypoints/llm/test_generate.py --ignore=entrypoints/llm/test_collective_rpc.py - pytest -v -s entrypoints/llm/test_generate.py # it needs a clean process - pytest -v -s entrypoints/offline_mode # Needs to avoid interference with other tests + mirror: + amd: + device: mi325_1 + depends_on: + - image-build-amd - label: Entrypoints Integration (API Server 1) timeout_in_minutes: 130 From 6831650c40ac3a34f049e285d9ad6b87daddbe00 Mon Sep 17 00:00:00 2001 From: Ming Yang Date: Wed, 25 Feb 2026 17:20:59 -0800 Subject: [PATCH 11/43] [offloader] v2: Hide weight onloading latency via prefetching (#29941) Signed-off-by: Ming Yang Signed-off-by: Michael Goin Co-authored-by: Michael Goin --- .../deepseek_v2_lite_prefetch_offload.sh | 57 ++ .buildkite/test_areas/e2e_integration.yaml | 9 + .../test_prefetch_offload.py | 33 + vllm/compilation/cuda_graph.py | 14 + vllm/config/__init__.py | 11 + vllm/config/cache.py | 16 +- vllm/config/offload.py | 153 ++++ vllm/config/vllm.py | 7 + vllm/engine/arg_utils.py | 65 +- vllm/entrypoints/llm.py | 21 + vllm/model_executor/models/utils.py | 119 +-- vllm/model_executor/offloader/__init__.py | 23 + vllm/model_executor/offloader/base.py | 145 ++++ vllm/model_executor/offloader/prefetch.py | 704 ++++++++++++++++++ vllm/model_executor/offloader/prefetch_ops.py | 94 +++ vllm/model_executor/offloader/uva.py | 140 ++++ vllm/v1/worker/gpu/cudagraph_utils.py | 18 + .../worker/gpu/spec_decode/eagle/cudagraph.py | 17 + vllm/v1/worker/gpu_model_runner.py | 22 +- vllm/v1/worker/gpu_ubatch_wrapper.py | 13 + 20 files changed, 1550 insertions(+), 131 deletions(-) create mode 100755 .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh create mode 100644 tests/basic_correctness/test_prefetch_offload.py create mode 100644 vllm/config/offload.py create mode 100644 vllm/model_executor/offloader/__init__.py create mode 100644 vllm/model_executor/offloader/base.py create mode 100644 vllm/model_executor/offloader/prefetch.py create mode 100644 vllm/model_executor/offloader/prefetch_ops.py create mode 100644 vllm/model_executor/offloader/uva.py diff --git a/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh new file mode 100755 index 00000000000..dddf23f1f2f --- /dev/null +++ b/.buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh @@ -0,0 +1,57 @@ +#!/usr/bin/env bash +set -euxo pipefail + +# Nightly e2e test for prefetch offloading with a MoE model. +# Runs DeepSeek-V2-Lite with prefetch offloading of MoE expert weights +# and validates GSM8K accuracy matches baseline (no offloading). +# +# args: [THRESHOLD] [NUM_QUESTIONS] [START_PORT] +THRESHOLD=${1:-0.25} +NUM_Q=${2:-1319} +PORT=${3:-8030} +OUT_DIR=${OUT_DIR:-/tmp/vllm-scheduled} +mkdir -p "${OUT_DIR}" + +wait_for_server() { + local port=$1 + timeout 600 bash -c ' + until curl -sf "http://127.0.0.1:'"$port"'/health" > /dev/null; do + sleep 1 + done' +} + +MODEL="deepseek-ai/DeepSeek-V2-Lite" + +cleanup() { + if [[ -n "${SERVER_PID:-}" ]] && kill -0 "${SERVER_PID}" 2>/dev/null; then + kill "${SERVER_PID}" 2>/dev/null || true + for _ in {1..20}; do + kill -0 "${SERVER_PID}" 2>/dev/null || break + sleep 0.5 + done + kill -9 "${SERVER_PID}" 2>/dev/null || true + fi +} +trap cleanup EXIT + +vllm serve "$MODEL" \ + --max-model-len 2048 \ + --offload-group-size 8 \ + --offload-num-in-group 2 \ + --offload-prefetch-step 1 \ + --offload-params w13_weight w2_weight \ + --port "$PORT" & +SERVER_PID=$! +wait_for_server "$PORT" + +TAG=$(echo "$MODEL" | tr '/: \\n' '_____') +OUT="${OUT_DIR}/${TAG}_prefetch_offload.json" +python3 tests/evals/gsm8k/gsm8k_eval.py --host http://127.0.0.1 --port "$PORT" --num-questions "${NUM_Q}" --save-results "${OUT}" +python3 - <= ${THRESHOLD}, f"${MODEL} prefetch_offload accuracy {acc}" +PY + +cleanup +SERVER_PID= diff --git a/.buildkite/test_areas/e2e_integration.yaml b/.buildkite/test_areas/e2e_integration.yaml index d95b73073d6..5b7f96bc7a2 100644 --- a/.buildkite/test_areas/e2e_integration.yaml +++ b/.buildkite/test_areas/e2e_integration.yaml @@ -28,3 +28,12 @@ steps: working_dir: "/vllm-workspace" commands: - bash .buildkite/scripts/scheduled_integration_test/qwen30b_a3b_fp8_block_ep_eplb.sh 0.8 200 8020 2 1 + +- label: DeepSeek V2-Lite Prefetch Offload Accuracy (H100) + timeout_in_minutes: 60 + device: h100 + optional: true + num_devices: 1 + working_dir: "/vllm-workspace" + commands: + - bash .buildkite/scripts/scheduled_integration_test/deepseek_v2_lite_prefetch_offload.sh 0.25 200 8030 diff --git a/tests/basic_correctness/test_prefetch_offload.py b/tests/basic_correctness/test_prefetch_offload.py new file mode 100644 index 00000000000..498887024ee --- /dev/null +++ b/tests/basic_correctness/test_prefetch_offload.py @@ -0,0 +1,33 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Test prefetch offloading correctness with Llama model.""" + +from ..utils import compare_two_settings + + +def test_prefetch_offload_llama(): + """Test prefetch CPU offloading with Llama-3.2-1B-Instruct. + + Compares outputs between: + 1. Baseline (no offloading) + 2. Prefetch offloading (group_size=8, num_in_group=2, prefetch_step=1) + + This tests prefetching-based offloading on a dense model. + """ + compare_two_settings( + "meta-llama/Llama-3.2-1B-Instruct", + [ + # Prefetch offloading configuration + "--offload-group-size", + "8", + "--offload-num-in-group", + "2", + "--offload-prefetch-step", + "1", + # Selective offloading: only MLP weights + "--offload-params", + "gate_up_proj", + "down_proj", + ], + [], # Baseline: no offloading + ) diff --git a/vllm/compilation/cuda_graph.py b/vllm/compilation/cuda_graph.py index 7ffa74d0d7e..7bada5e7ca3 100644 --- a/vllm/compilation/cuda_graph.py +++ b/vllm/compilation/cuda_graph.py @@ -17,6 +17,7 @@ from vllm.config import CUDAGraphMode, VllmConfig from vllm.distributed.device_communicators.pynccl_allocator import set_graph_pool_id from vllm.forward_context import BatchDescriptor, get_forward_context from vllm.logger import init_logger +from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.utils.torch_utils import current_stream, weak_ref_tensors @@ -265,6 +266,11 @@ class CUDAGraphWrapper: set_graph_pool_id(self.graph_pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + get_offloader().sync_prev_onload() + # mind-exploding: carefully manage the reference and memory. with torch.cuda.graph( cudagraph, @@ -273,6 +279,11 @@ class CUDAGraphWrapper: ): # `output` is managed by pytorch's cudagraph pool output = self.runnable(*args, **kwargs) + # Join offloader's copy stream after forward to avoid + # unjoined stream error. The last layer's start_prefetch + # forks copy_stream, but wait_prefetch only happens in + # the next forward pass. + get_offloader().join_after_forward() if self.cudagraph_options.weak_ref_output: # by converting it to weak ref, # the original `output` will immediately be released @@ -305,5 +316,8 @@ class CUDAGraphWrapper: f"got {new_input_addresses}" ) + # Sync offloader before replay - ensures any external dependencies + # from pre-capture prefetches are satisfied. + get_offloader().sync_prev_onload() entry.cudagraph.replay() return entry.output diff --git a/vllm/config/__init__.py b/vllm/config/__init__.py index 5bcf9865c27..452fb046660 100644 --- a/vllm/config/__init__.py +++ b/vllm/config/__init__.py @@ -24,6 +24,12 @@ from vllm.config.model import ( ) from vllm.config.multimodal import MultiModalConfig from vllm.config.observability import ObservabilityConfig +from vllm.config.offload import ( + OffloadBackend, + OffloadConfig, + PrefetchOffloadConfig, + UVAOffloadConfig, +) from vllm.config.parallel import EPLBConfig, ParallelConfig from vllm.config.pooler import PoolerConfig from vllm.config.profiler import ProfilerConfig @@ -85,6 +91,11 @@ __all__ = [ "MultiModalConfig", # From vllm.config.observability "ObservabilityConfig", + # From vllm.config.offload + "OffloadBackend", + "OffloadConfig", + "PrefetchOffloadConfig", + "UVAOffloadConfig", # From vllm.config.parallel "EPLBConfig", "ParallelConfig", diff --git a/vllm/config/cache.py b/vllm/config/cache.py index daceaa6c2bb..39ceb39205f 100644 --- a/vllm/config/cache.py +++ b/vllm/config/cache.py @@ -100,17 +100,15 @@ class CacheConfig: load a 13B model with BF16 weight, which requires at least 26GB GPU memory. Note that this requires fast CPU-GPU interconnect, as part of the model is loaded from CPU memory to GPU memory on the fly in each model forward pass. + + DEPRECATED: This field is deprecated and will be removed in v0.16. + Please use OffloadConfig.uva.cpu_offload_gb instead. """ cpu_offload_params: set[str] = Field(default_factory=set) - """ The set of parameter name segments to target for CPU offloading. - Unmatched parameters are not offloaded. If this set is empty, parameters - are offloaded non-selectively until the memory limit defined by - `cpu_offload_gb` is reached. - Examples: - - For parameter name "mlp.experts.w2_weight": - - "experts" or "experts.w2_weight" will match. - - "expert" or "w2" will NOT match (must be exact segments). - This allows distinguishing parameters like "w2_weight" and "w2_weight_scale". + """The set of parameter name segments to target for CPU offloading. + + DEPRECATED: This field is deprecated and will be removed in v0.16. + Please use OffloadConfig.uva.cpu_offload_params instead. """ calculate_kv_scales: bool = False """This enables dynamic calculation of `k_scale` and `v_scale` when diff --git a/vllm/config/offload.py b/vllm/config/offload.py new file mode 100644 index 00000000000..ad65e8acf35 --- /dev/null +++ b/vllm/config/offload.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Configuration for model weight offloading.""" + +import warnings +from typing import Literal + +from pydantic import Field, model_validator + +from vllm.config.utils import config + +OffloadBackend = Literal["auto", "uva", "prefetch"] + + +@config +class UVAOffloadConfig: + """Configuration for UVA (Unified Virtual Addressing) CPU offloading. + + Uses zero-copy access from CPU-pinned memory. Simple but requires + fast CPU-GPU interconnect. + """ + + cpu_offload_gb: float = Field(default=0, ge=0) + """The space in GiB to offload to CPU, per GPU. Default is 0, which means + no offloading. Intuitively, this argument can be seen as a virtual way to + increase the GPU memory size. For example, if you have one 24 GB GPU and + set this to 10, virtually you can think of it as a 34 GB GPU. Then you can + load a 13B model with BF16 weight, which requires at least 26GB GPU memory. + Note that this requires fast CPU-GPU interconnect, as part of the model is + loaded from CPU memory to GPU memory on the fly in each model forward pass. + This uses UVA (Unified Virtual Addressing) for zero-copy access. + """ + + cpu_offload_params: set[str] = Field(default_factory=set) + """The set of parameter name segments to target for CPU offloading. + Unmatched parameters are not offloaded. If this set is empty, parameters + are offloaded non-selectively until the memory limit defined by + `cpu_offload_gb` is reached. + Examples: + - For parameter name "mlp.experts.w2_weight": + - "experts" or "experts.w2_weight" will match. + - "expert" or "w2" will NOT match (must be exact segments). + This allows distinguishing parameters like "w2_weight" and "w2_weight_scale". + """ + + +@config +class PrefetchOffloadConfig: + """Configuration for prefetch-based CPU offloading. + + Groups layers and uses async H2D prefetch to hide transfer latency. + """ + + offload_group_size: int = Field(default=0, ge=0) + """Group every N layers together. Offload last `offload_num_in_group` + layers of each group. Default is 0 (disabled). + Example: group_size=8, num_in_group=2 offloads layers 6,7,14,15,22,23,... + Unlike cpu_offload_gb, this uses explicit async prefetching to hide transfer + latency. + """ + + offload_num_in_group: int = Field(default=1, ge=1) + """Number of layers to offload per group. + Must be <= offload_group_size. Default is 1.""" + + offload_prefetch_step: int = Field(default=1, ge=0) + """Number of layers to prefetch ahead. + Higher values hide more latency but use more GPU memory. Default is 1.""" + + offload_params: set[str] = Field(default_factory=set) + """The set of parameter name segments to target for prefetch offloading. + Unmatched parameters are not offloaded. If this set is empty, ALL + parameters of each offloaded layer are offloaded. + Uses segment matching: "w13_weight" matches "mlp.experts.w13_weight" + but not "mlp.experts.w13_weight_scale". + """ + + +@config +class OffloadConfig: + """Configuration for model weight offloading to reduce GPU memory usage.""" + + offload_backend: OffloadBackend = "auto" + """The backend for weight offloading. Options: + - "auto": Selects based on which sub-config has non-default values + (prefetch if offload_group_size > 0, uva if cpu_offload_gb > 0). + - "uva": UVA (Unified Virtual Addressing) zero-copy offloading. + - "prefetch": Async prefetch with group-based layer offloading. + """ + + uva: UVAOffloadConfig = Field(default_factory=UVAOffloadConfig) + """Parameters for UVA offloading backend.""" + + prefetch: PrefetchOffloadConfig = Field(default_factory=PrefetchOffloadConfig) + """Parameters for prefetch offloading backend.""" + + @model_validator(mode="after") + def validate_offload_config(self) -> "OffloadConfig": + """Validate offload configuration constraints.""" + if self.offload_backend == "prefetch" or self.prefetch.offload_group_size > 0: + if self.prefetch.offload_num_in_group > self.prefetch.offload_group_size: + raise ValueError( + f"offload_num_in_group ({self.prefetch.offload_num_in_group})" + f" must be <= offload_group_size" + f" ({self.prefetch.offload_group_size})" + ) + if self.prefetch.offload_prefetch_step < 1: + raise ValueError( + f"offload_prefetch_step" + f" ({self.prefetch.offload_prefetch_step})" + f" must be >= 1 when prefetch offloading is enabled" + f" (offload_group_size > 0)" + ) + + # Warn if both backends have non-default values + uva_active = self.uva.cpu_offload_gb > 0 + prefetch_active = self.prefetch.offload_group_size > 0 + if self.offload_backend == "uva" and prefetch_active: + warnings.warn( + "Prefetch offload fields are set but offload_backend='uva'. " + "Prefetch settings will be ignored.", + stacklevel=2, + ) + elif self.offload_backend == "prefetch" and uva_active: + warnings.warn( + "UVA offload fields are set but offload_backend='prefetch'. " + "UVA settings will be ignored.", + stacklevel=2, + ) + elif self.offload_backend == "auto" and uva_active and prefetch_active: + warnings.warn( + "Both UVA and prefetch offload fields are set with " + "offload_backend='auto'. Prefetch backend will be selected. " + "Set offload_backend explicitly to suppress this warning.", + stacklevel=2, + ) + return self + + def compute_hash(self) -> str: + """ + Provide a hash that uniquely identifies all the offload configs. + + All fields are included because PrefetchOffloader patches module + forwards and inserts custom ops (wait_prefetch, start_prefetch) + into the computation graph. Changing any offload setting can + alter which layers are hooked and how prefetch indices are + computed, so the compilation cache must distinguish them. + """ + from vllm.config.utils import get_hash_factors, hash_factors + + factors = get_hash_factors(self, ignored_factors=set()) + hash_str = hash_factors(factors) + return hash_str diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index d7deadd501e..33d4862636c 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -37,6 +37,7 @@ from .load import LoadConfig from .lora import LoRAConfig from .model import ModelConfig from .observability import ObservabilityConfig +from .offload import OffloadConfig from .parallel import ParallelConfig from .profiler import ProfilerConfig from .scheduler import SchedulerConfig @@ -259,6 +260,8 @@ class VllmConfig: """Device configuration.""" load_config: LoadConfig = Field(default_factory=LoadConfig) """Load configuration.""" + offload_config: OffloadConfig = Field(default_factory=OffloadConfig) + """Model weight offloading configuration.""" attention_config: AttentionConfig = Field(default_factory=AttentionConfig) """Attention configuration.""" kernel_config: KernelConfig = Field(default_factory=KernelConfig) @@ -361,6 +364,10 @@ class VllmConfig: vllm_factors.append(self.load_config.compute_hash()) else: vllm_factors.append("None") + if self.offload_config: + vllm_factors.append(self.offload_config.compute_hash()) + else: + vllm_factors.append("None") if self.attention_config: vllm_factors.append(self.attention_config.compute_hash()) else: diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index a962baba2a9..15a662ba2e5 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -48,12 +48,15 @@ from vllm.config import ( ModelConfig, MultiModalConfig, ObservabilityConfig, + OffloadConfig, ParallelConfig, PoolerConfig, + PrefetchOffloadConfig, ProfilerConfig, SchedulerConfig, SpeculativeConfig, StructuredOutputsConfig, + UVAOffloadConfig, VllmConfig, WeightTransferConfig, get_attr_docs, @@ -439,8 +442,13 @@ class EngineArgs: disable_sliding_window: bool = ModelConfig.disable_sliding_window disable_cascade_attn: bool = ModelConfig.disable_cascade_attn swap_space: float = CacheConfig.swap_space - cpu_offload_gb: float = CacheConfig.cpu_offload_gb - cpu_offload_params: set[str] = get_field(CacheConfig, "cpu_offload_params") + offload_backend: str = OffloadConfig.offload_backend + cpu_offload_gb: float = UVAOffloadConfig.cpu_offload_gb + cpu_offload_params: set[str] = get_field(UVAOffloadConfig, "cpu_offload_params") + offload_group_size: int = PrefetchOffloadConfig.offload_group_size + offload_num_in_group: int = PrefetchOffloadConfig.offload_num_in_group + offload_prefetch_step: int = PrefetchOffloadConfig.offload_prefetch_step + offload_params: set[str] = get_field(PrefetchOffloadConfig, "offload_params") gpu_memory_utilization: float = CacheConfig.gpu_memory_utilization kv_cache_memory_bytes: int | None = CacheConfig.kv_cache_memory_bytes max_num_batched_tokens: int | None = None @@ -948,10 +956,6 @@ class EngineArgs: cache_group.add_argument( "--prefix-caching-hash-algo", **cache_kwargs["prefix_caching_hash_algo"] ) - cache_group.add_argument("--cpu-offload-gb", **cache_kwargs["cpu_offload_gb"]) - cache_group.add_argument( - "--cpu-offload-params", **cache_kwargs["cpu_offload_params"] - ) cache_group.add_argument( "--calculate-kv-scales", **cache_kwargs["calculate_kv_scales"] ) @@ -977,6 +981,37 @@ class EngineArgs: "--kv-offloading-backend", **cache_kwargs["kv_offloading_backend"] ) + # Model weight offload related configs + offload_kwargs = get_kwargs(OffloadConfig) + uva_kwargs = get_kwargs(UVAOffloadConfig) + prefetch_kwargs = get_kwargs(PrefetchOffloadConfig) + offload_group = parser.add_argument_group( + title="OffloadConfig", + description=OffloadConfig.__doc__, + ) + offload_group.add_argument( + "--offload-backend", **offload_kwargs["offload_backend"] + ) + offload_group.add_argument("--cpu-offload-gb", **uva_kwargs["cpu_offload_gb"]) + offload_group.add_argument( + "--cpu-offload-params", **uva_kwargs["cpu_offload_params"] + ) + offload_group.add_argument( + "--offload-group-size", + **prefetch_kwargs["offload_group_size"], + ) + offload_group.add_argument( + "--offload-num-in-group", + **prefetch_kwargs["offload_num_in_group"], + ) + offload_group.add_argument( + "--offload-prefetch-step", + **prefetch_kwargs["offload_prefetch_step"], + ) + offload_group.add_argument( + "--offload-params", **prefetch_kwargs["offload_params"] + ) + # Multimodal related configs multimodal_kwargs = get_kwargs(MultiModalConfig) multimodal_group = parser.add_argument_group( @@ -1466,8 +1501,6 @@ class EngineArgs: sliding_window=sliding_window, enable_prefix_caching=self.enable_prefix_caching, prefix_caching_hash_algo=self.prefix_caching_hash_algo, - cpu_offload_gb=self.cpu_offload_gb, - cpu_offload_params=self.cpu_offload_params, calculate_kv_scales=self.calculate_kv_scales, kv_sharing_fast_prefill=self.kv_sharing_fast_prefill, mamba_cache_dtype=self.mamba_cache_dtype, @@ -1825,6 +1858,21 @@ class EngineArgs: compilation_config.max_cudagraph_capture_size = ( self.max_cudagraph_capture_size ) + + offload_config = OffloadConfig( + offload_backend=self.offload_backend, + uva=UVAOffloadConfig( + cpu_offload_gb=self.cpu_offload_gb, + cpu_offload_params=self.cpu_offload_params, + ), + prefetch=PrefetchOffloadConfig( + offload_group_size=self.offload_group_size, + offload_num_in_group=self.offload_num_in_group, + offload_prefetch_step=self.offload_prefetch_step, + offload_params=self.offload_params, + ), + ) + config = VllmConfig( model_config=model_config, cache_config=cache_config, @@ -1832,6 +1880,7 @@ class EngineArgs: scheduler_config=scheduler_config, device_config=device_config, load_config=load_config, + offload_config=offload_config, attention_config=attention_config, kernel_config=kernel_config, lora_config=lora_config, diff --git a/vllm/entrypoints/llm.py b/vllm/entrypoints/llm.py index 2d925d0a992..ee78d4d4894 100644 --- a/vllm/entrypoints/llm.py +++ b/vllm/entrypoints/llm.py @@ -170,6 +170,19 @@ class LLM: the model weights. This virtually increases the GPU memory space you can use to hold the model weights, at the cost of CPU-GPU data transfer for every forward pass. + offload_group_size: Prefetch offloading: Group every N layers + together. Offload last `offload_num_in_group` layers of each group. + Default is 0 (disabled). + offload_num_in_group: Prefetch offloading: Number of layers to + offload per group. Default is 1. + offload_prefetch_step: Prefetch offloading: Number of layers to + prefetch ahead. Higher values hide more latency but use more GPU + memory. Default is 1. + offload_params: Prefetch offloading: Set of parameter name segments + to selectively offload. Only parameters whose names contain one of + these segments will be offloaded (e.g., {"gate_up_proj", "down_proj"} + for MLP weights, or {"w13_weight", "w2_weight"} for MoE expert + weights). If None or empty, all parameters are offloaded. enforce_eager: Whether to enforce eager execution. If True, we will disable CUDA graph and always execute the model in eager mode. If False, we will use CUDA graph and eager execution in hybrid. @@ -224,6 +237,10 @@ class LLM: gpu_memory_utilization: float = 0.9, swap_space: float = 4, cpu_offload_gb: float = 0, + offload_group_size: int = 0, + offload_num_in_group: int = 1, + offload_prefetch_step: int = 1, + offload_params: set[str] | None = None, enforce_eager: bool = False, enable_return_routed_experts: bool = False, disable_custom_all_reduce: bool = False, @@ -333,6 +350,10 @@ class LLM: kv_cache_memory_bytes=kv_cache_memory_bytes, swap_space=swap_space, cpu_offload_gb=cpu_offload_gb, + offload_group_size=offload_group_size, + offload_num_in_group=offload_num_in_group, + offload_prefetch_step=offload_prefetch_step, + offload_params=offload_params or set(), enforce_eager=enforce_eager, enable_return_routed_experts=enable_return_routed_experts, disable_custom_all_reduce=disable_custom_all_reduce, diff --git a/vllm/model_executor/models/utils.py b/vllm/model_executor/models/utils.py index 65874248977..c55693bcff9 100644 --- a/vllm/model_executor/models/utils.py +++ b/vllm/model_executor/models/utils.py @@ -9,11 +9,9 @@ from typing import Any, Literal, Protocol, overload import torch import torch.nn as nn -from torch.func import functional_call from torch.nn.modules.module import register_module_module_registration_hook from transformers import PretrainedConfig -import vllm.envs as envs from vllm.config import VllmConfig from vllm.distributed import ( get_tensor_model_parallel_rank, @@ -31,14 +29,11 @@ from vllm.model_executor.models.interfaces import supports_any_eagle from vllm.multimodal import NestedTensors from vllm.sequence import IntermediateTensors from vllm.utils.math_utils import cdiv -from vllm.utils.mem_utils import format_gib from vllm.utils.platform_utils import ( is_pin_memory_available, - is_uva_available, ) from vllm.utils.torch_utils import ( direct_register_custom_op, - get_accelerator_view_from_cpu_tensor, ) logger = init_logger(__name__) @@ -612,98 +607,6 @@ class PPMissingLayer(torch.nn.Identity): return args[0] if args else next(iter(kwargs.values())) -_CPU_OFFLOAD_BYTES = 0 -_CPU_OFFLOAD_MAX_BYTES = 0 -_CPU_OFFLOAD_PARAMS = set() - - -def set_cpu_offload_max_bytes(max_bytes: int) -> None: - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - _CPU_OFFLOAD_BYTES = 0 - _CPU_OFFLOAD_MAX_BYTES = max_bytes - - -def set_cpu_offload_params(params: set[str]) -> None: - global _CPU_OFFLOAD_PARAMS - _CPU_OFFLOAD_PARAMS = params - - -def maybe_offload_to_cpu(module: torch.nn.Module) -> torch.nn.Module: - if (params := next(module.parameters(), None)) is None: - return module - - device = params.device - - if device == torch.device("cpu"): - return module - - global _CPU_OFFLOAD_MAX_BYTES, _CPU_OFFLOAD_BYTES - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: - return module - - pin_memory = ( - is_pin_memory_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY - ) - uva_offloading = is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA - - # offload parameters to CPU - # use pin_memory if possible, which helps cudagraph capture speed - offloaded_parameters = False - for name, p in module.named_parameters(): - if _CPU_OFFLOAD_BYTES >= _CPU_OFFLOAD_MAX_BYTES: - # we use per-parameter offloading - # one module might have some parameters offloaded and some not - break - - if _CPU_OFFLOAD_PARAMS: - # Check if parameter belongs to the offloading set - # Add dots here to ensure we match full segments only - # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" but not - # "mlp.experts.w2_weight_scale" - should_offload = any( - f".{param}." in f".{name}." for param in _CPU_OFFLOAD_PARAMS - ) - if not should_offload: - continue - - cpu_data = p.data.to(device="cpu") - if pin_memory: - cpu_data = cpu_data.pin_memory() - - if not uva_offloading: - p.data = cpu_data - else: - p.data = get_accelerator_view_from_cpu_tensor(cpu_data) - p._vllm_is_uva_offloaded = True - - _CPU_OFFLOAD_BYTES += p.data.numel() * p.data.element_size() - offloaded_parameters = True - - if offloaded_parameters and not uva_offloading: - original_forward = module.forward - - def forward(*args, **kwargs): - module.forward = original_forward - device_state = { - # here we blindly call `to(device)` - # if the parameter is already on the device, it will be a no-op - k: v.to(device, non_blocking=True) - for k, v in module.state_dict().items() - } - - # set `tie_weights=False` as tied weights in original model - # become untied when calling .to(device) individually - output = functional_call( - module, device_state, args=args, kwargs=kwargs, tie_weights=False - ) - module.forward = forward - return output - - module.forward = forward - - return module - - def make_layers( num_hidden_layers: int, layer_fn: LayerFn, @@ -711,25 +614,31 @@ def make_layers( ) -> tuple[int, int, torch.nn.ModuleList]: """Make a list of layers with the given layer function, taking pipeline parallelism into account. + + Args: + num_hidden_layers: Total number of hidden layers in the model. + layer_fn: Function to create a layer given its index. + prefix: Prefix for layer names. + + Returns: + Tuple of (start_layer, end_layer, modules). """ from vllm.distributed.parallel_state import get_pp_group from vllm.distributed.utils import get_pp_indices + from vllm.model_executor.offloader import get_offloader start_layer, end_layer = get_pp_indices( num_hidden_layers, get_pp_group().rank_in_group, get_pp_group().world_size ) + modules = torch.nn.ModuleList( [PPMissingLayer() for _ in range(start_layer)] - + [ - maybe_offload_to_cpu(layer_fn(prefix=f"{prefix}.{idx}")) - for idx in range(start_layer, end_layer) - ] + + get_offloader().wrap_modules( + layer_fn(prefix=f"{prefix}.{idx}") for idx in range(start_layer, end_layer) + ) + [PPMissingLayer() for _ in range(end_layer, num_hidden_layers)] ) - if _CPU_OFFLOAD_MAX_BYTES > 0: - logger.info( - "Total CPU offloaded parameters: %s GBs", format_gib(_CPU_OFFLOAD_BYTES) - ) + return start_layer, end_layer, modules diff --git a/vllm/model_executor/offloader/__init__.py b/vllm/model_executor/offloader/__init__.py new file mode 100644 index 00000000000..a6522ff7c0a --- /dev/null +++ b/vllm/model_executor/offloader/__init__.py @@ -0,0 +1,23 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Model parameter offloading infrastructure.""" + +from vllm.model_executor.offloader.base import ( + BaseOffloader, + NoopOffloader, + create_offloader, + get_offloader, + set_offloader, +) +from vllm.model_executor.offloader.prefetch import PrefetchOffloader +from vllm.model_executor.offloader.uva import UVAOffloader + +__all__ = [ + "BaseOffloader", + "NoopOffloader", + "UVAOffloader", + "PrefetchOffloader", + "create_offloader", + "get_offloader", + "set_offloader", +] diff --git a/vllm/model_executor/offloader/base.py b/vllm/model_executor/offloader/base.py new file mode 100644 index 00000000000..7c61b318b88 --- /dev/null +++ b/vllm/model_executor/offloader/base.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py +"""Base classes for model parameter offloading.""" + +from abc import ABC, abstractmethod +from collections.abc import Generator +from typing import TYPE_CHECKING + +import torch.nn as nn + +from vllm.logger import init_logger + +if TYPE_CHECKING: + from vllm.config import OffloadConfig + +logger = init_logger(__name__) + + +""" +class relation: + +BaseOffloader (ABC) + * implemented by: UVAOffloader + * implemented by: PrefetchOffloader + * uses: _ModuleOffloader + * uses: _BaseParamOffloader (ABC) + * implemented by: _CpuParamOffloader +""" + + +class BaseOffloader(ABC): + """Base class for model parameter offloading strategies. + + Offloaders control how model parameters are stored and loaded during + inference. Different strategies trade memory for compute/transfer time. + """ + + @abstractmethod + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + ) -> list[nn.Module]: + """Wrap modules with offloading logic. + + Args: + modules_generator: Generator yielding modules to potentially offload. + + Returns: + List of modules, potentially with offloading hooks installed. + """ + pass + + def post_init(self): + """Called after model construction completes. + + Offloaders can use this to: + - Finalize parameter storage + - Start initial prefetching + - Allocate shared resources + """ + return + + def sync_prev_onload(self) -> None: # noqa: B027 + """Sync previous onload operations. Override in subclasses.""" + pass + + def join_after_forward(self) -> None: # noqa: B027 + """Join streams after forward. Override in subclasses.""" + pass + + def _wait_for_layer(self, layer_idx: int) -> None: # noqa: B027 + """Wait for layer prefetch. Override in subclasses.""" + pass + + def _start_prefetch(self, layer_idx: int) -> None: # noqa: B027 + """Start layer prefetch. Override in subclasses.""" + pass + + +class NoopOffloader(BaseOffloader): + """No-op offloader that returns modules as-is without any offloading.""" + + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + ) -> list[nn.Module]: + """Return modules unchanged.""" + return list(modules_generator) + + +# Global singleton offloader instance (defaults to no-op). +_instance: BaseOffloader = NoopOffloader() + + +def get_offloader() -> BaseOffloader: + """Get the global offloader instance.""" + return _instance + + +def set_offloader(instance: BaseOffloader) -> None: + """Set the global offloader instance.""" + global _instance + _instance = instance + logger.info("Offloader set to %s", type(instance).__name__) + + +def create_offloader(offload_config: "OffloadConfig") -> BaseOffloader: + """Create an offloader based on the offload configuration. + + Uses the explicit ``offload_backend`` selector. When set to ``"auto"``, + selects prefetch if ``offload_group_size > 0``, UVA if + ``cpu_offload_gb > 0``, otherwise noop. + """ + from vllm.model_executor.offloader.prefetch import PrefetchOffloader + from vllm.model_executor.offloader.uva import UVAOffloader + + backend = offload_config.offload_backend + uva = offload_config.uva + prefetch = offload_config.prefetch + + if backend == "auto": + if prefetch.offload_group_size > 0: + backend = "prefetch" + elif uva.cpu_offload_gb > 0: + backend = "uva" + else: + return NoopOffloader() + + if backend == "prefetch": + return PrefetchOffloader( + group_size=prefetch.offload_group_size, + num_in_group=prefetch.offload_num_in_group, + prefetch_step=prefetch.offload_prefetch_step, + offload_params=prefetch.offload_params, + mode="cpu", + ) + elif backend == "uva": + return UVAOffloader( + cpu_offload_max_bytes=int(uva.cpu_offload_gb * 1024**3), + cpu_offload_params=uva.cpu_offload_params, + ) + else: + return NoopOffloader() diff --git a/vllm/model_executor/offloader/prefetch.py b/vllm/model_executor/offloader/prefetch.py new file mode 100644 index 00000000000..b43cb8b7d87 --- /dev/null +++ b/vllm/model_executor/offloader/prefetch.py @@ -0,0 +1,704 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Adapted from +# https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/utils/offloader.py +"""Prefetch-based CPU offloading with async prefetching. + +Uses static buffers and event-based stream forking for torch.compile + +CUDA graph compatibility. Events allow the copy stream to join CUDA +graph captures, ensuring H2D copies are properly captured. +""" + +from abc import ABC, abstractmethod +from collections.abc import Generator +from dataclasses import dataclass +from typing import Any + +import torch +import torch.nn as nn + +# Import prefetch_ops to register custom ops at module load time +import vllm.model_executor.offloader.prefetch_ops # noqa: F401 +from vllm.logger import init_logger +from vllm.model_executor.offloader.base import BaseOffloader +from vllm.utils.platform_utils import is_pin_memory_available + +logger = init_logger(__name__) + + +@dataclass +class ParamInfo: + """Metadata about an offloaded parameter.""" + + name: str + shape: tuple[int, ...] + stride: tuple[int, ...] + dtype: torch.dtype + + @property + def key(self) -> tuple[str, tuple[int, ...], tuple[int, ...], torch.dtype]: + """Unique key for buffer pool grouping. + + Includes parameter name to prevent different parameters with the same + shape from sharing buffers within the same layer. Parameters with the + same name across different layers will share buffers (via slots). + + Includes stride because parameters with same shape but different + strides need separate buffers to preserve memory layout. + """ + return (self.name, self.shape, self.stride, self.dtype) + + @property + def num_bytes(self) -> int: + """Size in bytes.""" + numel = 1 + for dim in self.shape: + numel *= dim + return numel * torch.finfo(self.dtype).bits // 8 + + +class StaticBufferPool: + """Pre-allocated GPU buffer pool for offloaded parameters. + + Allocates slot_capacity copies of each unique parameter + (name, shape, stride, dtype), allowing for double/triple buffering + during prefetch. + + Buffer slots are reused circularly: layer N uses slot (N % slot_capacity). + + The key includes parameter name to prevent different parameters within + the same layer from sharing buffers. Parameters with the same name + across different layers share buffers via the slot mechanism. + """ + + def __init__( + self, + param_infos: list[ParamInfo], + slot_capacity: int, + device: torch.device, + ): + self.slot_capacity = slot_capacity + self.total_bytes = 0 + self._device = device + + # Group by (shape, stride, dtype) - only allocate unique combinations + unique_params: dict[tuple, ParamInfo] = {} + for info in param_infos: + if info.key not in unique_params: + unique_params[info.key] = info + + # Allocate buffers: key -> list of tensors (one per slot) + self._buffers: dict[tuple, list[torch.Tensor]] = {} + for key, info in unique_params.items(): + slot_tensors = [] + for _ in range(slot_capacity): + # Use empty_strided to preserve parameter's memory layout + buf = torch.empty_strided( + size=info.shape, + stride=info.stride, + dtype=info.dtype, + device=device, + ) + slot_tensors.append(buf) + self.total_bytes += info.num_bytes + self._buffers[key] = slot_tensors + + logger.debug( + "[StaticBufferPool] Allocated %d unique (name, shape, stride, dtype), " + "%d slots each, total %.4f GB", + len(unique_params), + slot_capacity, + self.total_bytes / 1e9, + ) + + def get_buffer( + self, + name: str, + shape: tuple[int, ...], + stride: tuple[int, ...], + dtype: torch.dtype, + slot_idx: int, + ) -> torch.Tensor: + """Get a static buffer for the given name/shape/stride/dtype/slot.""" + key = (name, shape, stride, dtype) + return self._buffers[key][slot_idx % self.slot_capacity] + + +class PrefetchOffloader(BaseOffloader): + """Prefetching-based offloader with group-based layer selection. + + Groups layers and uses async H2D prefetch to hide transfer latency. + Uses static buffers and stream synchronization for torch.compile and + CUDA graph compatibility. + + Args: + group_size: Group every N layers together. + num_in_group: Offload this many layers per group (last N of each group). + prefetch_step: Number of layers to prefetch ahead. + mode: Offload mode ("cpu" is currently supported). + """ + + def __init__( + self, + group_size: int, + num_in_group: int, + prefetch_step: int, + offload_params: set[str] | None = None, + mode: str = "cpu", + ): + self.group_size = group_size + self.num_in_group = num_in_group + self.prefetch_step = prefetch_step + self.offload_params = offload_params or set() + self.mode = mode + + # Copy stream for async H2D transfers + self.copy_stream = torch.cuda.Stream() + + # Module offloaders and buffer pool (populated in wrap_modules/post_init) + self.module_offloaders: list[_ModuleOffloader] = [] + self.buffer_pool: StaticBufferPool | None = None + self.total_offloaded_bytes = 0 + + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + ) -> list[nn.Module]: + """Wrap modules with prefetch offloading logic.""" + assert len(self.module_offloaders) == 0, ( + "wrap_modules should only be called once" + ) + + all_modules = [] + offload_modules = [] + + for module_index, module in enumerate(modules_generator): + all_modules.append(module) + + # Select layers to offload based on group pattern + # Offload last num_in_group layers of each group_size + if module_index % self.group_size >= self.group_size - self.num_in_group: + if self.offload_params: + whitelist = [ + name + for name, _ in module.named_parameters() + if any(f".{p}." in f".{name}." for p in self.offload_params) + ] + else: + whitelist = [name for name, _ in module.named_parameters()] + + if not whitelist: + continue # skip layers with no matching params + + offload_modules.append(module) + self.module_offloaders.append( + _ModuleOffloader( + mode=self.mode, + module=module, + copy_stream=self.copy_stream, + whitelist_param_names=whitelist, + layer_idx=len(self.module_offloaders), + ) + ) + + for index, module in enumerate(offload_modules): + self._hook_module_forward(index, module) + + return all_modules + + def _hook_module_forward(self, index: int, module: nn.Module): + """Hook module's forward with torch.compile-compatible sync.""" + original_forward = module.forward + + def forward(*args, **kwargs): + # Temporarily restore original forward to avoid recursion + module.forward = original_forward + + # Wait for this layer's prefetch to complete + # mutates_args on input_tensor creates data dependency for torch.compile + input_tensor = args[0] if args else kwargs.get("hidden_states") + torch.ops.vllm.wait_prefetch(input_tensor, index) + + # No parameter swapping needed - parameters already point to + # GPU static buffers (set in assign_static_buffer) + output = original_forward(*args, **kwargs) + + # Start prefetch for next layer (circular) + # mutates_args on output_tensor creates ordering dependency + next_index = (index + self.prefetch_step) % len(self.module_offloaders) + # Handle tuple output (e.g., (hidden_states, residual)) + if isinstance(output, tuple): + torch.ops.vllm.start_prefetch(output[0], next_index) + else: + torch.ops.vllm.start_prefetch(output, next_index) + + # No explicit offload needed - static buffers are reused implicitly + + # Restore hooked forward + module.forward = forward + return output + + module.forward = forward + + def _wait_for_layer(self, layer_idx: int): + """Called by custom op - wait for copy to complete. + + Synchronization strategy: + - During CUDA graph capture: use event-based wait (graph-compatible) + - Outside capture (warmup/eager): use wait_stream (more robust) + + During capture, we skip wait for pre-capture prefetches because: + 1. sync_before_graph_capture() ensures pre-capture work is complete + 2. We can't wait on pre-capture events during capture (isolation error) + """ + offloader = self.module_offloaders[layer_idx] + + if torch.cuda.is_current_stream_capturing(): + # During capture, skip wait for pre-capture prefetches. + # sync_before_graph_capture() ensures pre-capture work is complete. + if not offloader._prefetch_in_capture: + return + # Event-based wait for in-capture prefetches (graph-compatible) + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + # Mark that this prefetch has been waited on (joined). + offloader._prefetch_in_capture = False + else: + if offloader._event_valid_for_eager: + # Use per-layer event to only wait for THIS layer's copy, + # allowing other layers' prefetches to run concurrently. + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + else: + # Event not usable (unrecorded or recorded during capture). + # Fall back to wait_stream to drain all copy_stream work. + torch.cuda.current_stream().wait_stream(self.copy_stream) + + def sync_prev_onload(self): + """Sync previous onload operations. + + Ensures any H2D copies in flight on copy_stream complete before + the compute stream continues. Call this before CUDA graph + capture/replay or when synchronization is needed. + """ + torch.cuda.current_stream().wait_stream(self.copy_stream) + + def _start_prefetch(self, layer_idx: int): + """Called by custom op - start async copy to static buffer.""" + offloader = self.module_offloaders[layer_idx] + offloader.start_onload_to_static() + + def join_after_forward(self): + """Join copy_stream after model forward completes. + + Call this after the model forward pass but before CUDA graph capture + ends. This ensures copy_stream is rejoined for any prefetches started + during the forward pass. + + We join ALL layers that have _prefetch_in_capture=True, meaning their + prefetch was started during capture but not yet waited on (joined). + This handles both full and piecewise cudagraph modes correctly: + - Full mode: joins layers 0..prefetch_step-1 (prefetched by last layers) + - Piecewise mode: joins only layers prefetched by THIS subgraph's layers + """ + if not self.module_offloaders: + return + # Join all layers whose prefetch was started in capture but not waited on + for offloader in self.module_offloaders: + if offloader._prefetch_in_capture: + torch.cuda.current_stream().wait_event(offloader._copy_done_event) + offloader._prefetch_in_capture = False + + def post_init(self): + """Allocate static buffer pool and start initial prefetches. + + Note: Parameters have already been offloaded to CPU during wrap_modules() + (in _CpuParamOffloader.__init__), so GPU memory is available for the + static buffer pool. + """ + # Sync CPU storage with current param.data BEFORE collecting param info. + # This is needed because process_weights_after_loading may have: + # 1. Transformed weights (quantization, transpose, etc.) + # 2. Created new CPU tensors via device_loading_context + # Our _cpu_storage would be stale otherwise. + for offloader in self.module_offloaders: + offloader.sync_cpu_storage() + + # Collect parameter info (now using synced CPU storage) + param_infos: list[ParamInfo] = [] + device: torch.device | None = None + + for offloader in self.module_offloaders: + param_infos.extend(offloader.get_param_infos()) + if device is None: + device = offloader.device + + if device is None: + # No modules to offload + return + + # Allocate static buffer pool + self.buffer_pool = StaticBufferPool( + param_infos=param_infos, + slot_capacity=self.prefetch_step, + device=device, + ) + + # Assign buffer slots and point parameters to GPU buffers + for idx, offloader in enumerate(self.module_offloaders): + slot_idx = idx % self.prefetch_step + offloader.assign_buffer_slot(self.buffer_pool, slot_idx) + + # Collect offloaded bytes + for offloader in self.module_offloaders: + offloader.post_init() + self.total_offloaded_bytes += offloader.offloaded_bytes + + logger.info_once( + f"[PrefetchOffloader] Initialized {len(self.module_offloaders)} modules. " + f"Total GPU memory saved: {self.total_offloaded_bytes / 1e9:.4f} GB, " + f"Static buffer pool: {self.buffer_pool.total_bytes / 1e9:.4f} GB " + f"(group_size={self.group_size}, num_in_group={self.num_in_group}, " + f"prefetch_step={self.prefetch_step}, mode={self.mode})" + ) + + # Start initial prefetches + for i in range(min(self.prefetch_step, len(self.module_offloaders))): + self.module_offloaders[i].start_onload_to_static() + + +class _ModuleOffloader: + """Manages offloading for a single module. + + Uses static buffers from a shared pool instead of dynamic allocation. + """ + + def __init__( + self, + mode: str, + module: nn.Module, + copy_stream: torch.cuda.Stream, + whitelist_param_names: list[str], + layer_idx: int, + ): + self.mode = mode + self.module = module + self.device = next(module.parameters()).device + self.copy_stream = copy_stream + self.layer_idx = layer_idx + self.offloaded_bytes = 0 + + # Event to signal when H2D copy to static buffer is complete. + # Used for per-layer synchronization (both eager and capture modes). + self._copy_done_event = torch.cuda.Event() + + # Track whether _copy_done_event is valid for eager-mode wait_event. + # False when: (1) never recorded, or (2) last recorded during a + # cudagraph capture (events become invalid after capture ends). + # In these cases we fall back to wait_stream. + self._event_valid_for_eager = False + + # Track if last prefetch was started during CUDA graph capture. + # Used to skip wait_event during capture for pre-capture prefetches. + self._prefetch_in_capture = False + + assert self.device != torch.device("cpu"), ( + "Module parameters should not already be on CPU " + "(offloader handles CPU placement)" + ) + + # Buffer pool and slot (assigned in assign_buffer_slot) + self._buffer_pool: StaticBufferPool | None = None + self._buffer_slot_idx: int = 0 + + param_dict = dict(self.module.named_parameters()) + assert all(name in param_dict for name in whitelist_param_names), ( + f"Whitelist params {whitelist_param_names} not found in module params " + f"{list(param_dict.keys())}" + ) + + self._param_offloaders = { + name: _BaseParamOffloader.create(mode, module=module, param_name=name) + for name in whitelist_param_names + } + + def post_init(self): + """Collect total offloaded bytes (offloading already done in __init__).""" + for param_offloader in self._param_offloaders.values(): + param_offloader.post_init() + self.offloaded_bytes += param_offloader.offloaded_bytes + + def sync_cpu_storage(self): + """Sync CPU storage with current param.data. + + Called after process_weights_after_loading to ensure _cpu_storage + contains the final processed weights, not stale pre-loading data. + """ + for param_offloader in self._param_offloaders.values(): + param_offloader.sync_cpu_storage() + + def get_param_infos(self) -> list[ParamInfo]: + """Get parameter metadata for buffer pool allocation. + + Note: sync_cpu_storage() must be called before this method to ensure + _cpu_storage reflects the final processed weights (after quantization). + """ + infos = [] + for name, offloader in self._param_offloaders.items(): + cpu_storage = offloader._cpu_storage + assert cpu_storage is not None, "CPU storage not initialized" + infos.append( + ParamInfo( + name=name, + shape=tuple(cpu_storage.shape), + stride=tuple(cpu_storage.stride()), + dtype=cpu_storage.dtype, + ) + ) + return infos + + def assign_buffer_slot(self, pool: StaticBufferPool, slot_idx: int): + """Assign this module to a buffer slot in the pool. + + Also assigns static GPU buffers to each parameter offloader, + which moves the parameter data to point to the GPU buffer. + """ + self._buffer_pool = pool + self._buffer_slot_idx = slot_idx + + # Assign static buffers to parameters + # Use CPU storage shape/stride/dtype since param.data is now empty + for name, offloader in self._param_offloaders.items(): + cpu_storage = offloader._cpu_storage + assert cpu_storage is not None, "CPU storage not initialized" + buffer = pool.get_buffer( + name=name, + shape=tuple(cpu_storage.shape), + stride=tuple(cpu_storage.stride()), + dtype=cpu_storage.dtype, + slot_idx=slot_idx, + ) + offloader.assign_static_buffer(buffer) + + def start_onload_to_static(self): + """Start async copy from CPU storage to GPU buffer. + + Uses event-based forking to join copy_stream to CUDA graph capture. + This ensures H2D copies are properly captured when recording a graph. + + IMPORTANT: We must wait for the compute stream before copying, because + the previous layer's forward may still be using the buffer (GPU ops are + async). Without this sync, we could overwrite the buffer while it's + being read. + """ + assert self._buffer_pool is not None, "Buffer pool not assigned" + + # Track if this prefetch is being captured (for _wait_for_layer logic) + self._prefetch_in_capture = torch.cuda.is_current_stream_capturing() + + # Fork: record event on compute stream, copy_stream waits on it + # This joins copy_stream to any active CUDA graph capture + fork_event = torch.cuda.Event() + torch.cuda.current_stream().record_event(fork_event) + self.copy_stream.wait_event(fork_event) + + with torch.cuda.stream(self.copy_stream): + for name, offloader in self._param_offloaders.items(): + cpu_storage = offloader._cpu_storage + gpu_buffer = offloader._gpu_buffer + assert cpu_storage is not None, "CPU storage not initialized" + assert gpu_buffer is not None, "GPU buffer not assigned" + assert not is_pin_memory_available() or cpu_storage.is_pinned(), ( + f"CPU storage for {name} is not pinned! " + "non_blocking=True H2D copy from non-pinned memory " + "causes stream synchronization that breaks " + "event-based fork synchronization." + ) + gpu_buffer.copy_(cpu_storage, non_blocking=True) + + # Record completion event for _wait_for_layer to use + self._copy_done_event.record(self.copy_stream) + # Event is only valid for eager wait_event if recorded outside capture. + # Events recorded during capture become invalid after capture ends. + self._event_valid_for_eager = not torch.cuda.is_current_stream_capturing() + + +class _BaseParamOffloader(ABC): + """Base class for parameter offloading strategies.""" + + # CPU storage for offloaded parameters (set by subclasses) + _cpu_storage: torch.Tensor | None + # GPU buffer reference (set by subclasses when using static buffers) + _gpu_buffer: torch.Tensor | None + + @staticmethod + def create(mode: str, **kwargs) -> "_BaseParamOffloader": + """Factory method to create appropriate offloader for mode.""" + if mode == "cpu": + return _CpuParamOffloader(**kwargs) + else: + raise ValueError(f"Unknown offload mode: {mode}") + + def __init__(self, module: nn.Module, param_name: str): + self._module = module + self._param_name = param_name + self.offloaded_bytes = 0 + self._cpu_storage = None + self._gpu_buffer = None + + @property + def _param(self) -> nn.Parameter: + """Get the parameter being offloaded. + + Supports dotted names (e.g. 'self_attn.qkv_proj.weight') by + traversing the module hierarchy. + """ + obj: Any = self._module + for attr in self._param_name.split("."): + obj = getattr(obj, attr) + return obj + + def post_init(self): + """Initialize offloading (move parameter to storage).""" + return + + @abstractmethod + def sync_cpu_storage(self) -> None: + """Sync CPU storage with current param.data. + + Called after process_weights_after_loading to update _cpu_storage + with the final processed weights. + """ + pass + + @abstractmethod + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None: + """Point parameter data to GPU static buffer.""" + pass + + +class _CpuParamOffloader(_BaseParamOffloader): + """Offload parameter to pinned CPU memory. + + Uses GPU static buffers as the actual parameter, with CPU storage + kept separately. This ensures torch.compile sees GPU tensors at trace time. + + The offloading happens in two phases: + 1. __init__() - copies GPU data to CPU, frees GPU memory immediately + 2. assign_static_buffer() - points param.data to GPU static buffer + """ + + def __init__(self, module: nn.Module, param_name: str): + super().__init__(module, param_name) + self._cpu_storage: torch.Tensor | None = None + self._gpu_buffer: torch.Tensor | None = None # Store reference to GPU buffer + + # Offload to CPU immediately to free GPU memory during model loading + self._offload_to_cpu_internal() + + def _offload_to_cpu_internal(self): + """Copy parameter data to pinned CPU storage and free GPU memory. + + This replaces param.data with CPU storage, allowing weight loading + to continue writing to CPU memory. GPU memory is freed when the + original GPU tensor is garbage collected. + """ + param = self._param + pin_memory = is_pin_memory_available() + + # Create pinned CPU storage and copy current GPU data + self._cpu_storage = torch.empty_strided( + size=param.data.size(), + stride=param.data.stride(), + dtype=param.data.dtype, + layout=param.data.layout, + device="cpu", + pin_memory=pin_memory, + ) + self._cpu_storage.copy_(param.data) + + self.offloaded_bytes = ( + self._cpu_storage.numel() * self._cpu_storage.element_size() + ) + + # Point param.data to CPU storage - this allows weight loading to work + # and frees GPU memory when the original GPU tensor is garbage collected + param.data = self._cpu_storage + + def _update_cpu_storage_from_param(self) -> None: + """Update _cpu_storage from current param.data, ensuring pinned memory. + + After process_weights_after_loading, device_loading_context creates + non-pinned CPU tensors via `p.data = p.data.to("cpu")`. Using + non-pinned memory with `copy_(src, non_blocking=True)` causes CUDA to + perform a stream synchronization before the copy, breaking the + event-based fork synchronization and potentially allowing the copy + to overwrite the GPU buffer while the compute stream still reads it. + + This method ensures _cpu_storage always uses pinned memory when + available, re-pinning if necessary. + """ + param = self._param + + if param.data.device.type == "cpu": + if is_pin_memory_available() and not param.data.is_pinned(): + pinned = torch.empty_strided( + size=param.data.size(), + stride=param.data.stride(), + dtype=param.data.dtype, + layout=param.data.layout, + device="cpu", + pin_memory=True, + ) + pinned.copy_(param.data) + self._cpu_storage = pinned + else: + self._cpu_storage = param.data + else: + # param.data is on GPU - copy to existing CPU storage + assert self._cpu_storage is not None + self._cpu_storage.copy_(param.data) + + def assign_static_buffer(self, gpu_buffer: torch.Tensor) -> None: + """Point parameter data to GPU static buffer. + + This is called after weight loading AND process_weights_after_loading + complete. At this point: + - param.data may have been replaced by device_loading_context + (which creates new CPU tensors after quantization processing) + - We need to update _cpu_storage to point to current param.data + so that prefetch copies the processed weights, not stale data + - Then point param.data to the GPU buffer for torch.compile + """ + assert self._cpu_storage is not None, ( + "_offload_to_cpu_internal() must be called before assign_static_buffer()" + ) + + # Get current parameter (may have been replaced by + # process_weights_after_loading) + param = self._param + + # Update _cpu_storage to current param.data. This is critical because: + # 1. process_weights_after_loading may transform weights (quantization) + # 2. device_loading_context creates NEW CPU tensors when moving back + # 3. Our old _cpu_storage would have pre-processed or stale data + self._update_cpu_storage_from_param() + + # Store reference to GPU buffer for use in start_onload + self._gpu_buffer = gpu_buffer + + # Point parameter to static GPU buffer - this is what torch.compile sees + param.data = gpu_buffer + + def sync_cpu_storage(self) -> None: + """Sync CPU storage with current param.data. + + Called after process_weights_after_loading to update _cpu_storage + with the final processed weights. This is critical because: + 1. process_weights_after_loading may transform weights (quantization) + 2. device_loading_context creates NEW CPU tensors when moving back + 3. Our old _cpu_storage would have pre-processed or stale data + """ + self._update_cpu_storage_from_param() + + def post_init(self): + """No-op: offloading done in offload_to_cpu/assign_static_buffer.""" + pass diff --git a/vllm/model_executor/offloader/prefetch_ops.py b/vllm/model_executor/offloader/prefetch_ops.py new file mode 100644 index 00000000000..d1f59b67b4a --- /dev/null +++ b/vllm/model_executor/offloader/prefetch_ops.py @@ -0,0 +1,94 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Custom ops for prefetch offloader torch.compile + CUDA graph compatibility. + +These ops use mutates_args to create data dependencies that prevent +the compiler from reordering prefetch/sync operations. +""" + +from __future__ import annotations + +import torch + +from vllm.model_executor.offloader.base import get_offloader +from vllm.utils.torch_utils import direct_register_custom_op + +# --- wait_prefetch op --- + + +def _wait_prefetch_impl( + input_tensor: torch.Tensor, + layer_idx: int, +) -> None: + """Wait for prefetch of layer_idx to complete. + + Synchronizes the compute stream with the copy stream to ensure + the prefetched weights are ready for use. + + Args: + input_tensor: Input to the layer (e.g., hidden_states) - declared + as mutated to create data dependency for torch.compile. + layer_idx: Index of the layer to wait for. + """ + get_offloader()._wait_for_layer(layer_idx) + + +def _wait_prefetch_fake( + input_tensor: torch.Tensor, + layer_idx: int, +) -> None: + """Fake implementation for torch.compile tracing.""" + return + + +# --- start_prefetch op --- + + +def _start_prefetch_impl( + output_tensor: torch.Tensor, + layer_idx: int, +) -> None: + """Start async prefetch of layer_idx weights. + + Initiates H2D copy on the copy stream for the specified layer. + + Args: + output_tensor: Output from forward - declared as mutated to + prevent torch.compile from reordering this op before the + computation that produces output_tensor. + layer_idx: Index of the layer to prefetch. + """ + get_offloader()._start_prefetch(layer_idx) + + +def _start_prefetch_fake( + output_tensor: torch.Tensor, + layer_idx: int, +) -> None: + """Fake implementation for torch.compile tracing.""" + return + + +def register_prefetch_offloader_ops() -> None: + """Register custom ops for prefetch offloader. + + Must be called before the ops are used. This is typically done + at module import time. + """ + direct_register_custom_op( + op_name="wait_prefetch", + op_func=_wait_prefetch_impl, + mutates_args=["input_tensor"], + fake_impl=_wait_prefetch_fake, + ) + + direct_register_custom_op( + op_name="start_prefetch", + op_func=_start_prefetch_impl, + mutates_args=["output_tensor"], + fake_impl=_start_prefetch_fake, + ) + + +# Register ops at module import time +register_prefetch_offloader_ops() diff --git a/vllm/model_executor/offloader/uva.py b/vllm/model_executor/offloader/uva.py new file mode 100644 index 00000000000..c524e43cdda --- /dev/null +++ b/vllm/model_executor/offloader/uva.py @@ -0,0 +1,140 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""UVA-based CPU offloading using Unified Virtual Addressing.""" + +from collections.abc import Generator + +import torch +import torch.nn as nn +from torch.func import functional_call + +import vllm.envs as envs +from vllm.logger import init_logger +from vllm.model_executor.offloader.base import BaseOffloader +from vllm.utils.mem_utils import format_gib +from vllm.utils.platform_utils import is_pin_memory_available, is_uva_available +from vllm.utils.torch_utils import get_accelerator_view_from_cpu_tensor + +logger = init_logger(__name__) + + +class UVAOffloader(BaseOffloader): + """Offloader using Unified Virtual Addressing (UVA) for zero-copy access. + + This offloader moves parameters to pinned CPU memory and creates CUDA views + using UVA. The GPU can then directly access the CPU memory without explicit + transfers, at the cost of PCIe bandwidth (slower than GPU memory). + + When UVA is disabled via env var, falls back to a functional_call-based + approach that moves parameters on-demand. + + Args: + cpu_offload_max_bytes: Maximum bytes to offload to CPU. + cpu_offload_params: Set of parameter name segments to selectively + offload. If empty, all parameters are eligible up to the byte limit. + """ + + def __init__( + self, + cpu_offload_max_bytes: int, + cpu_offload_params: set[str] | None = None, + ): + self.cpu_offload_max_bytes = cpu_offload_max_bytes + self.cpu_offload_bytes = 0 + self.cpu_offload_params = cpu_offload_params or set() + + self.pin_memory = ( + is_pin_memory_available() + and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY + ) + self.uva_offloading = ( + is_uva_available() and not envs.VLLM_WEIGHT_OFFLOADING_DISABLE_UVA + ) + + def wrap_modules( + self, + modules_generator: Generator[nn.Module, None, None], + ) -> list[nn.Module]: + """Wrap modules with UVA offloading.""" + modules = [self._maybe_offload_to_cpu(module) for module in modules_generator] + if self.cpu_offload_bytes > 0: + logger.info( + "Total CPU offloaded parameters: %s", + format_gib(self.cpu_offload_bytes), + ) + return modules + + def _maybe_offload_to_cpu(self, module: nn.Module) -> nn.Module: + """Offload module parameters to CPU using UVA if budget allows.""" + if (params := next(module.parameters(), None)) is None: + return module + + device = params.device + + if device == torch.device("cpu"): + return module + + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: + return module + + # offload parameters to CPU + # use pin_memory if possible, which helps cudagraph capture speed + offloaded_parameters = False + for name, p in module.named_parameters(): + if self.cpu_offload_bytes >= self.cpu_offload_max_bytes: + # we use per-parameter offloading + # one module might have some parameters offloaded and some not + break + + if self.cpu_offload_params: + # Check if parameter belongs to the offloading set + # Add dots here to ensure we match full segments only + # e.g., "experts.w2_weight" matches "mlp.experts.w2_weight" + # but not "mlp.experts.w2_weight_scale" + should_offload = any( + f".{param}." in f".{name}." for param in self.cpu_offload_params + ) + if not should_offload: + continue + + cpu_data = p.data.to(device="cpu") + if self.pin_memory: + cpu_data = cpu_data.pin_memory() + + if not self.uva_offloading: + p.data = cpu_data + else: + p.data = get_accelerator_view_from_cpu_tensor(cpu_data) + p._vllm_is_uva_offloaded = True + + self.cpu_offload_bytes += p.data.numel() * p.data.element_size() + offloaded_parameters = True + + if offloaded_parameters and not self.uva_offloading: + original_forward = module.forward + + def forward(*args, **kwargs): + module.forward = original_forward + device_state = { + # here we blindly call `to(device)` + # if the parameter is already on the device, + # it will be a no-op + k: v.to(device, non_blocking=True) + for k, v in module.state_dict().items() + } + + # set `tie_weights=False` as tied weights in original model + # become untied when calling .to(device) individually + output = functional_call( + module, + device_state, + args=args, + kwargs=kwargs, + tie_weights=False, + ) + module.forward = forward + return output + + module.forward = forward + + return module diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 5665937a03e..d70a4c7ab18 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -12,6 +12,7 @@ from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode from vllm.distributed.parallel_state import graph_capture, is_global_first_rank from vllm.forward_context import BatchDescriptor, set_forward_context +from vllm.model_executor.offloader.base import get_offloader from vllm.utils.math_utils import cdiv from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.attn_utils import ( @@ -189,6 +190,11 @@ class CudaGraphManager: # Capture the graph. assert num_tokens not in self.graphs graph = torch.cuda.CUDAGraph() + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + get_offloader().sync_prev_onload() + with ( set_forward_context( attn_metadata=attn_metadata, @@ -205,6 +211,11 @@ class CudaGraphManager: positions=positions, inputs_embeds=inputs_embeds, ) + # Join offloader's copy stream after forward to avoid unjoined + # stream error. The last layer's start_prefetch forks copy_stream, + # but wait_prefetch only happens in the next forward pass. + get_offloader().join_after_forward() + if self.use_aux_hidden_state_outputs: hidden_states, aux_hidden_states = model_output else: @@ -329,6 +340,13 @@ class CudaGraphManager: self, num_tokens: int ) -> torch.Tensor | tuple[torch.Tensor, list[torch.Tensor]]: assert num_tokens in self.graphs, f"No cudagraph for {num_tokens} tokens" + # Sync offloader before replay - needed when transitioning from + # eager/piecewise to full cudagraph (e.g., prefill → decode). + # The previous eager iteration's start_prefetch may have queued + # H2D copies on copy_stream that the graph's captured events + # cannot see. Without this, replay could overwrite static buffers + # while those copies are still in flight. + get_offloader().sync_prev_onload() self.graphs[num_tokens].replay() assert self.hidden_states is not None hidden_states = self.hidden_states[:num_tokens] diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py index c489a172c5a..eda8c37d53f 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/cudagraph.py @@ -7,6 +7,7 @@ import torch from vllm.config import VllmConfig from vllm.config.compilation import CUDAGraphMode +from vllm.model_executor.offloader.base import get_offloader from vllm.v1.kv_cache_interface import KVCacheConfig from vllm.v1.worker.gpu.block_table import BlockTables from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -115,6 +116,11 @@ class EagleCudaGraphManager: ) -> None: assert num_tokens not in self.graphs graph = torch.cuda.CUDAGraph() + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + get_offloader().sync_prev_onload() + with torch.cuda.graph(graph, self.pool): generate_fn( num_reqs, @@ -124,6 +130,10 @@ class EagleCudaGraphManager: num_tokens_across_dp, CUDAGraphMode.NONE, ) + # Join offloader's copy stream after forward to avoid unjoined + # stream error. The last layer's start_prefetch forks copy_stream, + # but wait_prefetch only happens in the next forward pass. + get_offloader().join_after_forward() self.graphs[num_tokens] = graph def _capture_piecewise_graph( @@ -171,4 +181,11 @@ class EagleCudaGraphManager: def run_fullgraph(self, num_tokens: int) -> None: assert num_tokens in self.graphs + # Sync offloader before replay - needed when transitioning from + # eager/piecewise to full cudagraph (e.g., prefill → decode). + # The previous eager iteration's start_prefetch may have queued + # H2D copies on copy_stream that the graph's captured events + # cannot see. Without this, replay could overwrite static buffers + # while those copies are still in flight. + get_offloader().sync_prev_onload() self.graphs[num_tokens].replay() diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f711d1d791b..d82b83b8c49 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -81,6 +81,11 @@ from vllm.model_executor.models.interfaces_base import ( is_pooling_model, is_text_generation_model, ) +from vllm.model_executor.offloader import ( + create_offloader, + get_offloader, + set_offloader, +) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.encoder_budget import MultiModalBudget from vllm.multimodal.inputs import ( @@ -378,6 +383,7 @@ class GPUModelRunner( self.vllm_config = vllm_config self.model_config = vllm_config.model_config self.cache_config = vllm_config.cache_config + self.offload_config = vllm_config.offload_config self.compilation_config = vllm_config.compilation_config self.lora_config = vllm_config.lora_config self.load_config = vllm_config.load_config @@ -386,14 +392,6 @@ class GPUModelRunner( self.speculative_config = vllm_config.speculative_config self.observability_config = vllm_config.observability_config - from vllm.model_executor.models.utils import ( - set_cpu_offload_max_bytes, - set_cpu_offload_params, - ) - - set_cpu_offload_max_bytes(int(self.cache_config.cpu_offload_gb * 1024**3)) - set_cpu_offload_params(self.cache_config.cpu_offload_params) - model_config = self.model_config cache_config = self.cache_config scheduler_config = self.scheduler_config @@ -749,6 +747,10 @@ class GPUModelRunner( pin_memory=self.pin_memory, ) + # Model weight offloader + # Make sure this is called before any get_offloader call + set_offloader(create_offloader(self.offload_config)) + # Ephemeral state transferred between execute_model() and sample_tokens(). self.execute_model_state: ExecuteModelState | None = None self.kv_connector_output: KVConnectorOutput | None = None @@ -4342,6 +4344,8 @@ class GPUModelRunner( self.model, self.vllm_config, CUDAGraphMode.NONE, self.device ) + get_offloader().post_init() + def _get_eagle3_aux_layers_from_config(self) -> tuple[int, ...] | None: """Extract Eagle3 auxiliary layer indices from speculative config. @@ -5780,7 +5784,7 @@ class GPUModelRunner( if block_sizes != [self.cache_config.block_size] or kernel_block_sizes != [ self.cache_config.block_size ]: - assert self.cache_config.cpu_offload_gb == 0, ( + assert self.offload_config.uva.cpu_offload_gb == 0, ( "Cannot re-initialize the input batch when CPU weight " "offloading is enabled. See https://github.com/vllm-project/vllm/pull/18298 " # noqa: E501 "for more details." diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index edbf797b1f5..45ba1bef9f2 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -20,6 +20,7 @@ from vllm.forward_context import ( override_forward_context, ) from vllm.logger import init_logger +from vllm.model_executor.offloader.base import get_offloader from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.utils.import_utils import has_deep_gemm @@ -239,6 +240,11 @@ class UBatchWrapper: set_graph_pool_id(self.graph_pool) else: set_graph_pool_id(current_platform.graph_pool_handle()) + + # Sync offloader's copy stream before capture. + # Ensure any pre-capture prefetches from offloader are complete. + get_offloader().sync_prev_onload() + with torch.cuda.graph( cudagraph_metadata.cudagraph, stream=compute_stream, @@ -250,6 +256,10 @@ class UBatchWrapper: sorted_results = [value for position, value in sorted(results)] result = torch.cat(sorted_results, dim=0) cudagraph_metadata.outputs = result + # Join offloader's copy stream after forward to avoid unjoined + # stream error. The last layer's start_prefetch forks copy_stream, + # but wait_prefetch only happens in the next forward pass. + get_offloader().join_after_forward() self.cudagraphs[num_tokens] = cudagraph_metadata return cudagraph_metadata.outputs @@ -461,6 +471,9 @@ class UBatchWrapper: and cudagraph_runtime_mode is CUDAGraphMode.FULL ): cudagraph_metadata = self.cudagraphs[num_tokens] + # Sync offloader before replay - ensures any external dependencies + # from pre-capture prefetches are satisfied. + get_offloader().sync_prev_onload() cudagraph_metadata.cudagraph.replay() return cudagraph_metadata.outputs else: From cbf8f7028cc0d80de4eeaf789b4bd56afbb5aafd Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 25 Feb 2026 20:28:31 -0500 Subject: [PATCH 12/43] [UX] Add `--performance-mode {balanced,interactivity,throughput}` (#34936) Signed-off-by: mgoin --- vllm/config/vllm.py | 30 ++++++++++++++++++++++++++---- vllm/engine/arg_utils.py | 12 +++++++++++- 2 files changed, 37 insertions(+), 5 deletions(-) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 33d4862636c..ef71a05d393 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -14,7 +14,7 @@ from datetime import datetime from enum import IntEnum from functools import lru_cache from pathlib import Path -from typing import TYPE_CHECKING, Any, TypeVar, get_args +from typing import TYPE_CHECKING, Any, Literal, TypeVar, get_args import torch from pydantic import ConfigDict, Field, model_validator @@ -76,6 +76,8 @@ class OptimizationLevel(IntEnum): """O3: Currently the same as -O2s.""" +PerformanceMode = Literal["balanced", "interactivity", "throughput"] + IS_QUANTIZED = False IS_DENSE = False # The optimizations that depend on these properties currently set to False @@ -312,6 +314,13 @@ class VllmConfig: performance. -O2 is used by default. See OptimizationLevel for full description.""" + performance_mode: PerformanceMode = "balanced" + """Performance mode for runtime behavior, 'balanced' is the default. + 'interactivity' favors low end-to-end per-request latency at small batch + sizes (fine-grained CUDA graphs, latency-oriented kernels). + 'throughput' favors aggregate tokens/sec at high concurrency (larger CUDA + graphs, more aggressive batching, throughput-oriented kernels).""" + weight_transfer_config: WeightTransferConfig | None = None """The configurations for weight transfer during RL training.""" @@ -643,6 +652,11 @@ class VllmConfig: # To give each torch profile run a unique instance name. self.instance_id = f"{time.time_ns()}" + if self.performance_mode != "balanced": + logger.info_once( + "Performance mode set to '%s'.", self.performance_mode, scope="local" + ) + self.try_verify_and_update_config() if self.model_config is not None: @@ -1332,9 +1346,15 @@ class VllmConfig: # sort to make sure the sizes are in ascending order cudagraph_capture_sizes.sort() else: - cudagraph_capture_sizes = [ - i for i in [1, 2, 4] if i <= max_cudagraph_capture_size - ] + if self.performance_mode == "interactivity": + # Fine-grained CUDA graphs at small batch sizes + # for minimal padding overhead + interactivity_max = min(max_cudagraph_capture_size, 32) + cudagraph_capture_sizes = list(range(1, interactivity_max + 1)) + else: + cudagraph_capture_sizes = [ + i for i in [1, 2, 4] if i <= max_cudagraph_capture_size + ] if max_cudagraph_capture_size >= 8: # Step size 8 for small batch sizes, up to 256(not included) cudagraph_capture_sizes += list( @@ -1345,6 +1365,8 @@ class VllmConfig: cudagraph_capture_sizes += list( range(256, max_cudagraph_capture_size + 1, 16) ) + # de-duplicate and sort the sizes + cudagraph_capture_sizes = sorted(set(cudagraph_capture_sizes)) if ( self.parallel_config.tensor_parallel_size > 1 diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 15a662ba2e5..ca76454d6d1 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -89,7 +89,7 @@ from vllm.config.parallel import ( ) from vllm.config.scheduler import SchedulerPolicy from vllm.config.utils import get_field -from vllm.config.vllm import OptimizationLevel +from vllm.config.vllm import OptimizationLevel, PerformanceMode from vllm.logger import init_logger, suppress_logging from vllm.platforms import CpuArchEnum, current_platform from vllm.plugins import load_general_plugins @@ -596,6 +596,7 @@ class EngineArgs: kv_sharing_fast_prefill: bool = CacheConfig.kv_sharing_fast_prefill optimization_level: OptimizationLevel = VllmConfig.optimization_level + performance_mode: PerformanceMode = VllmConfig.performance_mode kv_offloading_size: float | None = CacheConfig.kv_offloading_size kv_offloading_backend: KVOffloadingBackend = CacheConfig.kv_offloading_backend @@ -1264,6 +1265,7 @@ class EngineArgs: vllm_group.add_argument( "--optimization-level", **vllm_kwargs["optimization_level"] ) + vllm_group.add_argument("--performance-mode", **vllm_kwargs["performance_mode"]) vllm_group.add_argument( "--weight-transfer-config", **vllm_kwargs["weight_transfer_config"] ) @@ -1894,6 +1896,7 @@ class EngineArgs: profiler_config=self.profiler_config, additional_config=self.additional_config, optimization_level=self.optimization_level, + performance_mode=self.performance_mode, weight_transfer_config=self.weight_transfer_config, ) @@ -2110,6 +2113,13 @@ class EngineArgs: SchedulerConfig.DEFAULT_MAX_NUM_SEQS, ) + # If throughput mode is set, double max_num_batched_tokens and max_num_seqs. + if self.performance_mode == "throughput": + if orig_max_num_batched_tokens is None: + self.max_num_batched_tokens *= 2 + if orig_max_num_seqs is None: + self.max_num_seqs *= 2 + if orig_max_num_batched_tokens is None: assert model_config.max_model_len is not None, ( "max_model_len must be set by this point" From 1976356ee69750630189eb127fc9eeaa6f8e0c9e Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Wed, 25 Feb 2026 17:32:39 -0800 Subject: [PATCH 13/43] [MoE Refactor] MXFP4 Cutlass Experts to MK (#34542) Signed-off-by: Yongye Zhu --- .buildkite/test_areas/lm_eval.yaml | 26 ++++ .buildkite/test_areas/misc.yaml | 27 ---- tests/evals/gpt_oss/README.md | 49 +++++++ .../gpt_oss/configs/gpt-oss-20b-baseline.yaml | 6 + .../gpt-oss-20b-flashinfer-mxfp4-bf16.yaml | 8 ++ ...ss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml | 8 ++ .../gpt_oss/configs/gpt-oss-20b-marlin.yaml | 8 ++ ...t-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml | 8 ++ tests/evals/gpt_oss/configs/models-b200.txt | 5 + tests/evals/gpt_oss/configs/models-h100.txt | 5 + tests/evals/gpt_oss/conftest.py | 60 +++++++- tests/evals/gpt_oss/test_gpqa_correctness.py | 110 ++++++++++---- .../model_executor/layers/fused_moe/config.py | 4 + .../fused_moe/flashinfer_cutlass_moe.py | 118 +++++++++++++-- .../layers/fused_moe/modular_kernel.py | 6 +- .../layers/fused_moe/trtllm_moe.py | 17 ++- vllm/model_executor/layers/fused_moe/utils.py | 11 +- .../layers/quantization/mxfp4.py | 134 ++++++------------ .../layers/quantization/utils/quant_utils.py | 13 ++ 19 files changed, 454 insertions(+), 169 deletions(-) create mode 100644 tests/evals/gpt_oss/README.md create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-baseline.yaml create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-bf16.yaml create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-marlin.yaml create mode 100644 tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml create mode 100644 tests/evals/gpt_oss/configs/models-b200.txt create mode 100644 tests/evals/gpt_oss/configs/models-h100.txt diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index 1ef29f36cec..f25eae2400c 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -73,3 +73,29 @@ steps: num_devices: 2 commands: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=evals/gsm8k/configs/moe-refactor-dp-ep/config-b200.txt + +- label: GPQA Eval (GPT-OSS) (H100) + timeout_in_minutes: 120 + device: h100 + optional: true + num_devices: 2 + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - tests/evals/gpt_oss/ + commands: + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-h100.txt + +- label: GPQA Eval (GPT-OSS) (B200) + timeout_in_minutes: 120 + device: b200 + optional: true + num_devices: 2 + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - tests/evals/gpt_oss/ + commands: + - uv pip install --system 'gpt-oss[eval]==0.0.5' + - pytest -s -v evals/gpt_oss/test_gpqa_correctness.py --config-list-file=configs/models-b200.txt diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 5c5a9dbcbb6..69390cd6d37 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -153,33 +153,6 @@ steps: - pytest -v -s transformers_utils - pytest -v -s config -- label: GPT-OSS Eval (H100) - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - device: h100 - optional: true - source_file_dependencies: - - tests/evals/gpt_oss - - vllm/model_executor/models/gpt_oss.py - - vllm/model_executor/layers/quantization/mxfp4.py - commands: - - uv pip install --system 'gpt-oss[eval]==0.0.5' - - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58 - -- label: GPT-OSS Eval (B200) - timeout_in_minutes: 60 - working_dir: "/vllm-workspace/" - device: b200 - optional: true - source_file_dependencies: - - tests/evals/gpt_oss - - vllm/model_executor/models/gpt_oss.py - - vllm/model_executor/layers/quantization/mxfp4.py - - vllm/v1/attention/backends/flashinfer.py - commands: - - uv pip install --system 'gpt-oss[eval]==0.0.5' - - pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py --model openai/gpt-oss-20b --metric 0.58 - - label: Batch Invariance (H100) timeout_in_minutes: 25 device: h100 diff --git a/tests/evals/gpt_oss/README.md b/tests/evals/gpt_oss/README.md new file mode 100644 index 00000000000..98c0098bbd2 --- /dev/null +++ b/tests/evals/gpt_oss/README.md @@ -0,0 +1,49 @@ +# GPQA Evaluation using GPT-OSS + +This directory contains GPQA evaluation tests using the GPT-OSS evaluation package and vLLM server. + +## Usage + +### Run tests with pytest (like buildkite) + +```bash +# H200 +pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \ + --config-list-file=configs/models-h200.txt + +# B200 +pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \ + --config-list-file=configs/models-b200.txt +``` + +## Configuration Format + +Model configs in `configs/` directory use this YAML format: + +```yaml +model_name: "openai/gpt-oss-20b" +metric_threshold: 0.568 # Minimum expected accuracy +reasoning_effort: "low" # Reasoning effort level (default: "low") +server_args: "--tensor-parallel-size 2" # Server arguments +startup_max_wait_seconds: 1800 # Max wait for server startup (default: 1800) +env: # Environment variables (optional) + SOME_VAR: "value" +``` + +The `server_args` field accepts any arguments that can be passed to `vllm serve`. + +The `env` field accepts a dictionary of environment variables to set for the server process. + +## Adding New Models + +1. Create a new YAML config file in the `configs/` directory +2. Add the filename to the appropriate `models-*.txt` file + +## Tiktoken Encoding Files + +The tiktoken encoding files required by the vLLM server are automatically downloaded from OpenAI's public blob storage on first run: + +- `cl100k_base.tiktoken` +- `o200k_base.tiktoken` + +Files are cached in the `data/` directory. The `TIKTOKEN_ENCODINGS_BASE` environment variable is automatically set to point to this directory when running evaluations. diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-baseline.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-baseline.yaml new file mode 100644 index 00000000000..1df1cc93e47 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-baseline.yaml @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: "openai/gpt-oss-20b" +metric_threshold: 0.568 +reasoning_effort: "low" +server_args: "--tensor-parallel-size 2" diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-bf16.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-bf16.yaml new file mode 100644 index 00000000000..952f7e87035 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-bf16.yaml @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: "openai/gpt-oss-20b" +metric_threshold: 0.568 +reasoning_effort: "low" +server_args: "--tensor-parallel-size 2" +env: + VLLM_USE_FLASHINFER_MOE_MXFP4_BF16: "1" diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml new file mode 100644 index 00000000000..23ec14819ef --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: "openai/gpt-oss-20b" +metric_threshold: 0.568 +reasoning_effort: "low" +server_args: "--tensor-parallel-size 2" +env: + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: "1" diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-marlin.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-marlin.yaml new file mode 100644 index 00000000000..97e97fd19a6 --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-marlin.yaml @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: "openai/gpt-oss-20b" +metric_threshold: 0.568 +reasoning_effort: "low" +server_args: "--tensor-parallel-size 2" +env: + VLLM_MXFP4_USE_MARLIN: "1" diff --git a/tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml b/tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml new file mode 100644 index 00000000000..4cea743490f --- /dev/null +++ b/tests/evals/gpt_oss/configs/gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml @@ -0,0 +1,8 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +model_name: "openai/gpt-oss-20b" +metric_threshold: 0.568 +reasoning_effort: "low" +server_args: "--tensor-parallel-size 2" +env: + VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8: "1" diff --git a/tests/evals/gpt_oss/configs/models-b200.txt b/tests/evals/gpt_oss/configs/models-b200.txt new file mode 100644 index 00000000000..8519109e192 --- /dev/null +++ b/tests/evals/gpt_oss/configs/models-b200.txt @@ -0,0 +1,5 @@ +# B200 model configurations for GPQA evaluation +# Tests different environment variable combinations +gpt-oss-20b-flashinfer-mxfp4-bf16.yaml +gpt-oss-20b-flashinfer-mxfp4-mxfp8-cutlass.yaml +gpt-oss-20b-sm100-fi-mxfp4-mxfp8-trtllm.yaml \ No newline at end of file diff --git a/tests/evals/gpt_oss/configs/models-h100.txt b/tests/evals/gpt_oss/configs/models-h100.txt new file mode 100644 index 00000000000..9577bac5f1d --- /dev/null +++ b/tests/evals/gpt_oss/configs/models-h100.txt @@ -0,0 +1,5 @@ +# H100 model configurations for GPQA evaluation +# Tests different environment variable combinations +gpt-oss-20b-baseline.yaml +gpt-oss-20b-flashinfer-mxfp4-bf16.yaml +gpt-oss-20b-marlin.yaml diff --git a/tests/evals/gpt_oss/conftest.py b/tests/evals/gpt_oss/conftest.py index 2f140ae2c8e..d35dec4831a 100644 --- a/tests/evals/gpt_oss/conftest.py +++ b/tests/evals/gpt_oss/conftest.py @@ -4,13 +4,61 @@ Pytest configuration for GPT-OSS evaluation tests. """ +from pathlib import Path + def pytest_addoption(parser): - """Add command line options for pytest.""" - parser.addoption("--model", action="store", help="Model name to evaluate") + """Add custom command line options.""" parser.addoption( - "--metric", action="store", type=float, help="Expected metric threshold" - ) - parser.addoption( - "--server-args", action="store", default="", help="Additional server arguments" + "--config-list-file", + required=True, + help="File containing list of config files to test", ) + + +def pytest_generate_tests(metafunc): + """Generate test parameters from config files.""" + if "config_filename" in metafunc.fixturenames: + config_list_file = metafunc.config.getoption("--config-list-file") + + # Handle both relative and absolute paths + config_list_path = Path(config_list_file) + if not config_list_path.is_absolute(): + # If relative, try relative to test directory first + test_dir_path = Path(__file__).parent / config_list_file + if test_dir_path.exists(): + config_list_path = test_dir_path + else: + # Try relative to current working directory + config_list_path = Path.cwd() / config_list_file + + print(f"Looking for config list at: {config_list_path}") + + config_files = [] + if config_list_path.exists(): + # Determine config directory (same directory as the list file) + config_dir = config_list_path.parent + + with open(config_list_path) as f: + for line in f: + line = line.strip() + if line and not line.startswith("#"): + config_path = config_dir / line + print(f"Checking config file: {config_path}") + if config_path.exists(): + config_files.append(config_path) + print(f" Found: {config_path}") + else: + print(f" Missing: {config_path}") + else: + print(f"Config list file not found: {config_list_path}") + + # Generate test parameters + if config_files: + metafunc.parametrize( + "config_filename", + config_files, + ids=[config_file.stem for config_file in config_files], + ) + else: + print("No config files found, test will be skipped") diff --git a/tests/evals/gpt_oss/test_gpqa_correctness.py b/tests/evals/gpt_oss/test_gpqa_correctness.py index 151deaa059f..63188ec4076 100644 --- a/tests/evals/gpt_oss/test_gpqa_correctness.py +++ b/tests/evals/gpt_oss/test_gpqa_correctness.py @@ -5,22 +5,48 @@ GPQA evaluation using vLLM server and GPT-OSS evaluation package. Usage: pytest -s -v tests/evals/gpt_oss/test_gpqa_correctness.py \ - --model openai/gpt-oss-20b \ - --metric 0.58 \ - --server-args "--tensor-parallel-size 2" + --config-list-file=configs/models-h200.txt """ +import os +import shlex import subprocess import sys +import urllib.request +from pathlib import Path import regex as re +import yaml from tests.utils import RemoteOpenAIServer TOL = 0.05 # Absolute tolerance for accuracy comparison +# Path to tiktoken encoding files +TIKTOKEN_DATA_DIR = Path(__file__).parent / "data" -def run_gpqa_eval(model_name: str, base_url: str) -> float: +# Tiktoken encoding files to download +TIKTOKEN_FILES = { + "cl100k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/cl100k_base.tiktoken", + "o200k_base.tiktoken": "https://openaipublic.blob.core.windows.net/encodings/o200k_base.tiktoken", +} + + +def ensure_tiktoken_files(): + """Download tiktoken encoding files if they don't exist.""" + TIKTOKEN_DATA_DIR.mkdir(parents=True, exist_ok=True) + + for filename, url in TIKTOKEN_FILES.items(): + filepath = TIKTOKEN_DATA_DIR / filename + if not filepath.exists(): + print(f"Downloading {filename} from {url}...") + urllib.request.urlretrieve(url, filepath) + print(f" Downloaded to {filepath}") + else: + print(f" {filename} already exists.") + + +def run_gpqa_eval(model_name: str, base_url: str, reasoning_effort: str) -> float: """Run GPQA evaluation using the gpt-oss evaluation package.""" # Build the command to run the evaluation @@ -33,7 +59,7 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float: "--model", model_name, "--reasoning-effort", - "low", + reasoning_effort, "--base-url", base_url, "--n-threads", @@ -41,16 +67,29 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float: ] try: + # Set up environment for the evaluation subprocess + # Inherit current environment and add required variables + eval_env = os.environ.copy() + eval_env["OPENAI_API_KEY"] = "dummy" + # Run the evaluation result = subprocess.run( cmd, text=True, capture_output=True, timeout=1800, # 30 minute timeout - env={"OPENAI_API_KEY": "dummy"}, + env=eval_env, ) - print("Evaluation process output:\n", result.stdout) + print("Evaluation process stdout:\n", result.stdout) + print("Evaluation process stderr:\n", result.stderr) + print(f"Evaluation process return code: {result.returncode}") + + if result.returncode != 0: + raise RuntimeError( + f"Evaluation failed with exit code {result.returncode}:\n" + f"stdout: {result.stdout}\nstderr: {result.stderr}" + ) # Parse the output to extract the score match = re.search(r"'metric':\s*([\d.]+)", result.stdout) @@ -64,47 +103,62 @@ def run_gpqa_eval(model_name: str, base_url: str) -> float: except subprocess.TimeoutExpired as e: raise RuntimeError("Evaluation timed out") from e - except subprocess.CalledProcessError as e: - raise RuntimeError( - f"Evaluation failed with exit code {e.returncode}:\n" - f"stdout: {e.stdout}\nstderr: {e.stderr}" - ) from e -def test_gpqa_correctness(request): - """Test GPQA correctness for GPT-OSS model.""" +def test_gpqa_correctness(config_filename): + """Test GPQA correctness for a given model configuration.""" + # Ensure tiktoken files are downloaded + ensure_tiktoken_files() - # Get command line arguments - model_name = request.config.getoption("--model") - expected_metric = request.config.getoption("--metric") - server_args_str = request.config.getoption("--server-args") + # Verify tiktoken files exist + for filename in TIKTOKEN_FILES: + filepath = TIKTOKEN_DATA_DIR / filename + assert filepath.exists(), f"Tiktoken file not found: {filepath}" - # Parse server arguments - server_args = [] - if server_args_str: - server_args = server_args_str.split() + eval_config = yaml.safe_load(config_filename.read_text(encoding="utf-8")) + + # Parse server arguments from config (use shlex to handle quoted strings) + server_args_str = eval_config.get("server_args", "") + server_args = shlex.split(server_args_str) if server_args_str else [] # Add standard server arguments server_args.extend( [ "--trust-remote-code", + "--enforce-eager", + "--disable-uvicorn-access-log", ] ) - print(f"Starting GPQA evaluation for model: {model_name}") - print(f"Expected metric threshold: {expected_metric}") + # Build server environment with tiktoken path and any config-specified vars + server_env = {"TIKTOKEN_ENCODINGS_BASE": str(TIKTOKEN_DATA_DIR)} + if eval_config.get("env"): + server_env.update(eval_config["env"]) + + reasoning_effort = eval_config.get("reasoning_effort", "low") + + print(f"Starting GPQA evaluation for model: {eval_config['model_name']}") + print(f"Expected metric threshold: {eval_config['metric_threshold']}") + print(f"Reasoning effort: {reasoning_effort}") print(f"Server args: {' '.join(server_args)}") + print(f"Server environment variables: {server_env}") # Launch server and run evaluation with RemoteOpenAIServer( - model_name, server_args, max_wait_seconds=1800 + eval_config["model_name"], + server_args, + env_dict=server_env, + max_wait_seconds=eval_config.get("startup_max_wait_seconds", 1800), ) as remote_server: base_url = remote_server.url_for("v1") print(f"Server started at: {base_url}") - measured_metric = run_gpqa_eval(model_name, base_url) + measured_metric = run_gpqa_eval( + eval_config["model_name"], base_url, reasoning_effort + ) + expected_metric = eval_config["metric_threshold"] - print(f"GPQA Results for {model_name}:") + print(f"GPQA Results for {eval_config['model_name']}:") print(f" Measured metric: {measured_metric:.4f}") print(f" Expected metric: {expected_metric:.4f}") print(f" Tolerance: {TOL:.4f}") @@ -115,4 +169,4 @@ def test_gpqa_correctness(request): f"{expected_metric:.4f} - {TOL:.4f} = {expected_metric - TOL:.4f}" ) - print(f"✅ GPQA test passed for {model_name}") + print(f"GPQA test passed for {eval_config['model_name']}") diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index b6b8a17aea6..22e71d39101 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -242,6 +242,10 @@ class FusedMoEQuantConfig: def quant_dtype(self) -> torch.dtype | str | None: return self._a1.dtype + @property + def weight_quant_dtype(self) -> torch.dtype | str | None: + return self._w1.dtype + @property def is_quantized(self) -> bool: return self.quant_dtype is not None diff --git a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py index 4ec76ee9820..b9566a3a921 100644 --- a/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py +++ b/vllm/model_executor/layers/fused_moe/flashinfer_cutlass_moe.py @@ -4,6 +4,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config import get_current_vllm_config from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.activation import MoEActivation from vllm.model_executor.layers.fused_moe.config import ( @@ -18,6 +19,8 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8Dynamic128Sym, kFp8Static128BlockSym, kFp8StaticTensorSym, + kMxfp4Static, + kMxfp8Dynamic, kNvfp4Dynamic, kNvfp4Static, ) @@ -64,10 +67,18 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): quant_config: FusedMoEQuantConfig, ): super().__init__(moe_config, quant_config) - assert quant_config.quant_dtype in ("nvfp4", torch.float8_e4m3fn, None), ( - "Only nvfp4, fp8, bfloat16 and" + + assert quant_config.weight_quant_dtype in ( + "mxfp4", + "nvfp4", + torch.float8_e4m3fn, + None, + ), ( + "Only mxfp4, nvfp4, fp8, bfloat16 and" " float16 quantization are currently supported." ) + self.device = moe_config.device + self.num_experts = moe_config.num_local_experts self.ep_rank = moe_config.moe_parallel_config.ep_rank self.ep_size = moe_config.moe_parallel_config.ep_size self.tp_rank = moe_config.moe_parallel_config.tp_rank @@ -78,6 +89,28 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): # - pass per-block weight scales to the kernel # - skip input activation quantization (kernel applies scaling) self.use_deepseek_fp8_block_scale = quant_config.is_block_quantized + self.max_capture_size = ( + get_current_vllm_config().compilation_config.max_cudagraph_capture_size + ) + + if quant_config.weight_quant_dtype == "mxfp4": + # This value is used specifically for gpt-oss, + # Need to revisit this for other models + self.gemm1_alpha = torch.tensor( + [1.702] * self.num_experts, dtype=torch.float32, device=self.device + ) + self.gemm1_beta = torch.tensor( + [1.0] * self.num_experts, dtype=torch.float32, device=self.device + ) + self.gemm1_clamp_limit = torch.tensor( + [7.0] * self.num_experts, dtype=torch.float32, device=self.device + ) + if quant_config.quant_dtype == "mxfp8": + self.fake_input_scale = torch.ones( + self.num_experts, + device=self.device, + dtype=torch.float32, + ) @property def expects_unquantized_inputs(self) -> bool: @@ -119,20 +152,33 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): ] and p.has_device_capability(90) ) - # fp8 block-scale on 9.0 + # fp8 block-scale, wmxfp4a16 on 9.0 or ( - scheme == (kFp8Static128BlockSym, kFp8Dynamic128Sym) + scheme + in [ + (kMxfp4Static, None), + (kFp8Static128BlockSym, kFp8Dynamic128Sym), + ] and p.is_device_capability(90) ) - # nvfp4 on 10.0+ + # nvfp4, wmxfp4amxfp8 on 10.0+ or ( - scheme == (kNvfp4Static, kNvfp4Dynamic) and p.has_device_capability(100) + scheme + in [ + (kMxfp4Static, kMxfp8Dynamic), + (kNvfp4Static, kNvfp4Dynamic), + ] + and p.has_device_capability(100) ) ) @staticmethod def _supports_activation(activation: MoEActivation) -> bool: - return activation in [MoEActivation.SILU, MoEActivation.RELU2_NO_MUL] + return activation in [ + MoEActivation.SILU, + MoEActivation.RELU2_NO_MUL, + MoEActivation.SWIGLUOAI, + ] @staticmethod def _supports_parallel_config(moe_parallel_config: FusedMoEParallelConfig) -> bool: @@ -216,12 +262,23 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): activation_str_to_value_map = { MoEActivation.SILU: ActivationType.Swiglu, # This is the default + MoEActivation.SWIGLUOAI: ActivationType.Swiglu, # gpt-oss alias MoEActivation.RELU2_NO_MUL: ActivationType.Relu2, } assert activation in activation_str_to_value_map, ( f"{activation=} missing from {activation_str_to_value_map.keys()=}" ) + quant_scales = None + fc1_expert_weights = None + fc2_expert_weights = None + fc1_expert_biases = None + fc2_expert_biases = None + swiglu_alpha = None + swiglu_beta = None + swiglu_limit = None + use_mxfp8_act_scaling = False + use_w4_group_scaling = False # Select quantization metadata based on FP8 format/path if ( self.quant_dtype == torch.float8_e4m3fn @@ -256,6 +313,43 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): # FlashInfer API requires weight to be long for nvfp4 fc1_expert_weights = w1.view(torch.long) fc2_expert_weights = w2.view(torch.long) + elif self.weight_quant_dtype == "mxfp4": + assert self.w1_scale is not None and self.w2_scale is not None + assert w1.is_contiguous() and w2.is_contiguous() + assert self.gemm1_alpha is not None + assert self.gemm1_beta is not None + assert self.gemm1_clamp_limit is not None + assert topk_ids.is_contiguous() + + fc1_expert_biases = self.w1_bias + fc2_expert_biases = self.w2_bias + swiglu_alpha = self.gemm1_alpha + swiglu_beta = self.gemm1_beta + swiglu_limit = self.gemm1_clamp_limit + + if self.quant_dtype == "mxfp8": + assert self.fake_input_scale is not None + fc1_expert_weights = w1.view(torch.long) + fc2_expert_weights = w2.view(torch.long) + + quant_scales = [ + self.w1_scale.view(torch.int32), + self.fake_input_scale, + self.w2_scale.view(torch.int32), + self.fake_input_scale, + ] + use_mxfp8_act_scaling = True + else: + assert hidden_states.dtype == torch.bfloat16 + fc1_expert_weights = w1 + fc2_expert_weights = w2 + quant_scales = [ + self.w1_scale, + self.w2_scale, + ] + a1q_scale = None + use_w4_group_scaling = True + elif self.use_deepseek_fp8_block_scale: # FP8 block-scale path: provide block-scale weights, omit a1q_scale quant_scales = [ @@ -277,6 +371,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): token_final_scales=topk_weights, fc1_expert_weights=fc1_expert_weights, fc2_expert_weights=fc2_expert_weights, + fc1_expert_biases=fc1_expert_biases, + fc2_expert_biases=fc2_expert_biases, + swiglu_alpha=swiglu_alpha, + swiglu_beta=swiglu_beta, + swiglu_limit=swiglu_limit, + output=output, output_dtype=self.out_dtype, quant_scales=quant_scales, input_sf=a1q_scale, @@ -284,10 +384,12 @@ class FlashInferExperts(mk.FusedMoEPermuteExpertsUnpermute): tp_rank=self.tp_rank, ep_size=self.ep_size, ep_rank=self.ep_rank, - output=output, activation_type=activation_str_to_value_map[activation], # Informs FlashInfer to use the block-scale decoding path when True use_deepseek_fp8_block_scale=self.use_deepseek_fp8_block_scale, + use_mxfp8_act_scaling=use_mxfp8_act_scaling, + use_w4_group_scaling=use_w4_group_scaling, + tune_max_num_tokens=max(self.max_capture_size, 1), ) def moe_sum(self, input: torch.Tensor, output: torch.Tensor) -> None: diff --git a/vllm/model_executor/layers/fused_moe/modular_kernel.py b/vllm/model_executor/layers/fused_moe/modular_kernel.py index b4ceaa379f0..c2c0e809d70 100644 --- a/vllm/model_executor/layers/fused_moe/modular_kernel.py +++ b/vllm/model_executor/layers/fused_moe/modular_kernel.py @@ -564,9 +564,13 @@ class FusedMoEPermuteExpertsUnpermute(ABC): # @property - def quant_dtype(self) -> torch.dtype | None: + def quant_dtype(self) -> torch.dtype | str | None: return self.quant_config.quant_dtype + @property + def weight_quant_dtype(self) -> torch.dtype | str | None: + return self.quant_config.weight_quant_dtype + @property def block_shape(self) -> list[int] | None: return self.quant_config.block_shape diff --git a/vllm/model_executor/layers/fused_moe/trtllm_moe.py b/vllm/model_executor/layers/fused_moe/trtllm_moe.py index 61e06fa603d..2bd4cd79e03 100644 --- a/vllm/model_executor/layers/fused_moe/trtllm_moe.py +++ b/vllm/model_executor/layers/fused_moe/trtllm_moe.py @@ -25,15 +25,20 @@ class TrtLlmGenExperts(mk.FusedMoEPermuteExpertsUnpermute): self, moe_config: FusedMoEConfig, quant_config: FusedMoEQuantConfig, - gemm1_alpha, - gemm1_beta, - gemm1_clamp_limit, max_capture_size, ): super().__init__(moe_config, quant_config) - self.gemm1_alpha = gemm1_alpha - self.gemm1_beta = gemm1_beta - self.gemm1_clamp_limit = gemm1_clamp_limit + self.device = torch.cuda.current_device() + self.num_experts = moe_config.num_local_experts + self.gemm1_alpha = torch.tensor( + [1.702] * self.num_experts, dtype=torch.float32, device=self.device + ) + self.gemm1_beta = torch.tensor( + [1.0] * self.num_experts, dtype=torch.float32, device=self.device + ) + self.gemm1_clamp_limit = torch.tensor( + [7.0] * self.num_experts, dtype=torch.float32, device=self.device + ) self.max_capture_size = max_capture_size @staticmethod diff --git a/vllm/model_executor/layers/fused_moe/utils.py b/vllm/model_executor/layers/fused_moe/utils.py index ad32abf582c..019e408c195 100644 --- a/vllm/model_executor/layers/fused_moe/utils.py +++ b/vllm/model_executor/layers/fused_moe/utils.py @@ -195,11 +195,12 @@ def _mxfp8_e4m3_quantize( A_scale: torch.Tensor | None, per_act_token_quant: bool, block_shape: list[int] | None = None, + is_sf_swizzled_layout: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: assert A_scale is None assert not per_act_token_quant assert block_shape is None - return mxfp8_e4m3_quantize(A) + return mxfp8_e4m3_quantize(A, is_sf_swizzled_layout) def _mxfp6_e3m2_quantize( @@ -275,7 +276,13 @@ def moe_kernel_quantize_input( elif quant_dtype == "mxfp8": # TODO: `quant_dtype == "mxfp8"` is ambiguous, # should be fp8_e4m3. OCP MX also defines `fp8_e5m2`. - return _mxfp8_e4m3_quantize(A, A_scale, per_act_token_quant, block_shape) + return _mxfp8_e4m3_quantize( + A, + A_scale, + per_act_token_quant, + block_shape, + is_sf_swizzled_layout=is_fp4_scale_swizzled, + ) elif quant_dtype == "mxfp6_e3m2": return _mxfp6_e3m2_quantize(A, A_scale, per_act_token_quant, block_shape) elif quant_dtype == "mxfp6_e2m3": diff --git a/vllm/model_executor/layers/quantization/mxfp4.py b/vllm/model_executor/layers/quantization/mxfp4.py index 492963855e8..d81f0f80d2e 100644 --- a/vllm/model_executor/layers/quantization/mxfp4.py +++ b/vllm/model_executor/layers/quantization/mxfp4.py @@ -256,6 +256,7 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): "Please check your environment and try again." ) self._cache_permute_indices: dict[torch.Size, torch.Tensor] = {} + # Initialized in process_weights_after_loading for CUTLASS/SM90 backends self.moe_mk: mk.FusedMoEModularKernel | None = None def create_weights( @@ -648,19 +649,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16 ): - layer.gemm1_alpha = Parameter( - torch.tensor([1.702] * self.num_experts, dtype=torch.float32).cuda(), - requires_grad=False, - ) - layer.gemm1_beta = Parameter( - torch.tensor([1.0] * self.num_experts, dtype=torch.float32).cuda(), - requires_grad=False, - ) - layer.gemm1_clamp_limit = Parameter( - torch.tensor([7.0] * self.num_experts, dtype=torch.float32).cuda(), - requires_grad=False, - ) - sf_block_size = 32 # mxfp4 block size # Common shape assertions @@ -772,6 +760,30 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): layer.w2_weight_scale = torch.nn.Parameter( w2_scales_interleaved, requires_grad=False ) + + # theses two kernels go through the `flashinfer_cutlass_fused_moe` path + from vllm.model_executor.layers.fused_moe.flashinfer_cutlass_moe import ( + FlashInferExperts, + ) + + self.moe_quant_config = self.get_fused_moe_quant_config(layer) + assert self.moe_quant_config is not None + prepare_finalize = maybe_make_prepare_finalize( + moe=self.moe, + quant_config=self.moe_quant_config, + routing_tables=layer._maybe_init_expert_routing_tables(), + allow_new_interface=True, + ) + assert prepare_finalize is not None + + self.moe_mk = mk.FusedMoEModularKernel( + prepare_finalize, + FlashInferExperts( + moe_config=self.moe, + quant_config=self.moe_quant_config, + ), + shared_experts=None, + ) elif self.mxfp4_backend == Mxfp4Backend.TRITON: from triton_kernels.matmul_ogs import FlexCtx, PrecisionConfig @@ -847,7 +859,10 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): w1_scale=layer.w13_weight_scale, w2_scale=layer.w2_weight_scale, ) - elif self.mxfp4_backend in [Mxfp4Backend.SM100_FI_MXFP4_BF16]: + elif self.mxfp4_backend in [ + Mxfp4Backend.SM100_FI_MXFP4_BF16, + Mxfp4Backend.SM90_FI_MXFP4_BF16, + ]: return mxfp4_w4a16_moe_quant_config( w1_bias=layer.w13_bias, w2_bias=layer.w2_bias, @@ -897,9 +912,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): ): # B200 code-path kwargs = { - "gemm1_alpha": layer.gemm1_alpha, - "gemm1_beta": layer.gemm1_beta, - "gemm1_clamp_limit": layer.gemm1_clamp_limit, # TODO(bnell): part of quant_config "max_capture_size": self.max_capture_size, } @@ -935,20 +947,6 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): if layer.enable_eplb: raise NotImplementedError("EPLB is not supported for mxfp4") - if self.mxfp4_backend == Mxfp4Backend.MARLIN: - assert self.moe_mk is not None - - return self.moe_mk( - hidden_states=x, - w1=layer.w13_weight, - w2=layer.w2_weight, - topk_weights=topk_weights, - topk_ids=topk_ids, - activation=layer.activation, - global_num_experts=layer.global_num_experts, - expert_map=layer.expert_map, - apply_router_weight_on_input=layer.apply_router_weight_on_input, - ) assert _can_support_mxfp4( layer.use_grouped_topk, layer.topk_group, @@ -967,68 +965,22 @@ class Mxfp4MoEMethod(FusedMoEMethodBase): assert ( self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS or self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16 - ) - from vllm.utils.flashinfer import flashinfer_cutlass_fused_moe - - # Backend-specific preparation - if self.mxfp4_backend == Mxfp4Backend.SM100_FI_MXFP4_MXFP8_CUTLASS: - from flashinfer import mxfp8_quantize - - x_quant, x_scale = mxfp8_quantize(x, True, 32) - - fake_input_scale = torch.ones(self.num_experts, device=x.device) - quant_scales = [ - layer.w13_weight_scale.contiguous().view(torch.int32), - fake_input_scale, - layer.w2_weight_scale.contiguous().view(torch.int32), - fake_input_scale, - ] - - fi_input = x_quant - extra_kwargs = dict( - use_mxfp8_act_scaling=True, - input_sf=x_scale, - fc1_expert_weights=layer.w13_weight.contiguous().view(torch.long), - fc2_expert_weights=layer.w2_weight.contiguous().view(torch.long), - ) - elif self.mxfp4_backend == Mxfp4Backend.SM90_FI_MXFP4_BF16: - assert x.dtype == torch.bfloat16 - - quant_scales = [ - layer.w13_weight_scale, - layer.w2_weight_scale, - ] - - fi_input = x - extra_kwargs = dict( - use_w4_group_scaling=True, - fc1_expert_weights=layer.w13_weight, - fc2_expert_weights=layer.w2_weight, - ) - - output = torch.empty_like(x, dtype=torch.bfloat16) - - flashinfer_cutlass_fused_moe( - input=fi_input, - token_selected_experts=topk_ids.to(torch.int).contiguous(), - token_final_scales=topk_weights, - output_dtype=torch.bfloat16, - output=output, - quant_scales=quant_scales, - fc1_expert_biases=layer.w13_bias, - fc2_expert_biases=layer.w2_bias, - swiglu_alpha=layer.gemm1_alpha, - swiglu_beta=layer.gemm1_beta, - swiglu_limit=layer.gemm1_clamp_limit, - tp_size=self.moe.tp_size, - tp_rank=self.moe.tp_rank, - ep_size=self.moe.ep_size, - ep_rank=self.moe.ep_rank, - tune_max_num_tokens=max(self.max_capture_size, 1), - **extra_kwargs, + or self.mxfp4_backend == Mxfp4Backend.MARLIN ) - return output + assert self.moe_mk is not None + return self.moe_mk( + hidden_states=x, + w1=layer.w13_weight, + w2=layer.w2_weight, + topk_weights=topk_weights, + topk_ids=topk_ids, + activation=layer.activation, + global_num_experts=layer.global_num_experts, + apply_router_weight_on_input=layer.apply_router_weight_on_input, + expert_map=layer.expert_map, + shared_experts_input=shared_experts_input, + ) def apply_monolithic( self, diff --git a/vllm/model_executor/layers/quantization/utils/quant_utils.py b/vllm/model_executor/layers/quantization/utils/quant_utils.py index e42868e4176..12a1799d157 100644 --- a/vllm/model_executor/layers/quantization/utils/quant_utils.py +++ b/vllm/model_executor/layers/quantization/utils/quant_utils.py @@ -19,6 +19,7 @@ if TYPE_CHECKING: FP8_DTYPE = current_platform.fp8_dtype() FP4_DTYPE = torch.uint8 +MXFP_SCALE_DTYPE = torch.uint8 def get_fp8_min_max() -> tuple[float, float]: @@ -151,6 +152,18 @@ kFp8Static128BlockSym = QuantKey(FP8_DTYPE, kStatic128BlockScale, symmetric=True kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64)) kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True) +# TODO (zyongye): Convert all the torch.dtype to scale_dtype +# Changing that requires changing torch compile fused AR+Quant Quant key +# to avoid assertion error +kMxfp4DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp4Dynamic = QuantKey(FP4_DTYPE, scale=kMxfp4DynamicGroupScale, symmetric=True) + +kMxfp8DynamicGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, False, GroupShape(1, 32)) +kMxfp8Dynamic = QuantKey(FP8_DTYPE, scale=kMxfp8DynamicGroupScale, symmetric=True) + +kMxfp4StaticGroupScale = ScaleDesc(MXFP_SCALE_DTYPE, True, GroupShape(1, 32)) +kMxfp4Static = QuantKey(FP4_DTYPE, scale=kMxfp4StaticGroupScale, symmetric=True) + # Normalize the group_shape to the full extent for any dims that are -1 def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape): From de527e1cec820686f2bead759e4e99a20b172589 Mon Sep 17 00:00:00 2001 From: Michael Goin Date: Wed, 25 Feb 2026 20:44:44 -0500 Subject: [PATCH 14/43] [UX] Add `--moe-backend` arg for explicit kernel selection (#33807) Signed-off-by: mgoin Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- .../configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml | 3 +- .../gsm8k/configs/Qwen3-Next-FP8-EP2.yaml | 3 +- .../Llama-4-Scout-Fp8-ModelOpt-triton.yaml | 3 +- ...30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml | 5 +- .../Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml | 5 +- ...B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml | 5 +- ...en3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml | 5 +- ...wen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml | 5 +- .../Llama-4-Scout-BF16-fi-cutlass.yaml | 6 +- ...Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml | 5 +- .../Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml | 5 +- .../Llama-4-Scout-Fp8-ModelOpt-triton.yaml | 4 +- .../Mixtral-8x7B-BF16-fi-cutlass.yaml | 5 +- .../Mixtral-8x7B-Fp8-AutoFp8-fi-cutlass.yaml | 5 +- ...otron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml | 5 +- ...on-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml | 5 +- .../Qwen3-30B-A3B-BF16-fi-cutlass.yaml | 4 +- .../Qwen3-30B-A3B-Fp8-AutoFp8-fi-cutlass.yaml | 5 +- .../Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml | 5 +- .../Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml | 3 +- ...Qwen3-30B-A3B-Fp8-CT-Block-fi-cutlass.yaml | 5 +- .../Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml | 3 +- .../Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml | 5 +- .../Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml | 5 +- .../Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml | 4 +- ...en3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml | 5 +- ...wen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml | 5 +- ...3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml | 4 +- tests/quantization/test_blackwell_moe.py | 70 +++++++++++-------- vllm/config/kernel.py | 34 ++++++++- vllm/engine/arg_utils.py | 7 ++ .../model_executor/layers/fused_moe/config.py | 2 +- .../fused_moe/deepep_ll_prepare_finalize.py | 1 - vllm/model_executor/layers/fused_moe/layer.py | 1 + .../layers/fused_moe/oracle/fp8.py | 59 ++++++++++++++++ .../layers/fused_moe/oracle/nvfp4.py | 35 ++++++++++ .../layers/fused_moe/oracle/unquantized.py | 64 ++++++++++++++--- 37 files changed, 260 insertions(+), 140 deletions(-) diff --git a/tests/evals/gsm8k/configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml b/tests/evals/gsm8k/configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml index 673b473f817..7f2f096fd27 100644 --- a/tests/evals/gsm8k/configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml +++ b/tests/evals/gsm8k/configs/Qwen3-Next-80B-A3B-NVFP4-EP2.yaml @@ -8,5 +8,4 @@ server_args: >- --tensor-parallel-size 2 --enable-expert-parallel --speculative-config '{"method":"qwen3_next_mtp","num_speculative_tokens":1}' -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" + --moe-backend=flashinfer_trtllm diff --git a/tests/evals/gsm8k/configs/Qwen3-Next-FP8-EP2.yaml b/tests/evals/gsm8k/configs/Qwen3-Next-FP8-EP2.yaml index 9fae32734d7..abcb784a71e 100644 --- a/tests/evals/gsm8k/configs/Qwen3-Next-FP8-EP2.yaml +++ b/tests/evals/gsm8k/configs/Qwen3-Next-FP8-EP2.yaml @@ -7,5 +7,4 @@ server_args: >- --tensor-parallel-size 2 --enable-expert-parallel --async-scheduling -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" + --moe-backend=flashinfer_trtllm diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Llama-4-Scout-Fp8-ModelOpt-triton.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Llama-4-Scout-Fp8-ModelOpt-triton.yaml index 9e13797bb9a..fda02c367a3 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Llama-4-Scout-Fp8-ModelOpt-triton.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Llama-4-Scout-Fp8-ModelOpt-triton.yaml @@ -2,7 +2,6 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8" accuracy_threshold: 0.92 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel" +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=triton" env: - VLLM_USE_FLASHINFER_MOE_FP8: "0" VLLM_USE_DEEP_GEMM: "0" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml index 1328fdedf0c..6624cea1ef2 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutedsl-deepep-ll.yaml @@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "masked_gemm" +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml index 53fd62bac83..90265a12afc 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml index 87fac0e708c..f2d4588e3ae 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutedsl-deepep-ll.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "masked_gemm" +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --all2all-backend deepep_low_latency --moe-backend=flashinfer_cutedsl" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml index 44f8700e4b4..49be54e26b1 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml index 91a220c4f21..23d29e06f8c 100644 --- a/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor-dp-ep/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "latency" +server_args: "--enforce-eager --max-model-len 8192 --data-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_trtllm" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml index 5416d9232cd..e19500fd369 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-BF16-fi-cutlass.yaml @@ -2,8 +2,4 @@ model_name: "meta-llama/Llama-4-Scout-17B-16E-Instruct" accuracy_threshold: 0.92 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel" -env: - VLLM_USE_FLASHINFER_MOE_FP16: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" - +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml index 4c9a01274d9..217ee5e6034 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8" accuracy_threshold: 0.92 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml index 17f067215eb..7e9300d9fc7 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-fi-trtllm.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8" accuracy_threshold: 0.92 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "latency" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-triton.yaml b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-triton.yaml index ae6bf67556e..87f960afec2 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-triton.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Llama-4-Scout-Fp8-ModelOpt-triton.yaml @@ -2,6 +2,4 @@ model_name: "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8" accuracy_threshold: 0.92 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "0" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml index cc8df6292cf..1c5865974f7 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-BF16-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "mistralai/Mixtral-8x7B-v0.1" accuracy_threshold: 0.58 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel" -env: - VLLM_USE_FLASHINFER_MOE_FP16: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-Fp8-AutoFp8-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-Fp8-AutoFp8-fi-cutlass.yaml index b9c6a1997dc..f836a503803 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-Fp8-AutoFp8-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Mixtral-8x7B-Fp8-AutoFp8-fi-cutlass.yaml @@ -3,7 +3,4 @@ # accuracy_threshold: 0.62 # num_questions: 1319 # num_fewshot: 5 -# server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -# env: -# VLLM_USE_FLASHINFER_MOE_FP8: "1" -# VLLM_FLASHINFER_MOE_BACKEND: "throughput" +# server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml index 570569def1e..a06c93dcc87 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-Fp8-ModelOpt-fi-trtllm.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-FP8" accuracy_threshold: 0.29 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "latency" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml index d802ac3f31c..b5a8676d765 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Nemotron-Nano-30B-NvFp4-ModelOpt-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/NVIDIA-Nemotron-3-Nano-30B-A3B-NVFP4" accuracy_threshold: 0.29 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-fi-cutlass.yaml index b15126a4521..92b9c071e18 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-BF16-fi-cutlass.yaml @@ -2,6 +2,4 @@ model_name: "Qwen/Qwen3-30B-A3B" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel" -env: - VLLM_USE_FLASHINFER_MOE_FP16: "1" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --enable-expert-parallel --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-cutlass.yaml index 74820cd2834..b392f92453f 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml index d745c9b5b2b..4fd2f8d261b 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-fi-trtllm.yaml @@ -2,7 +2,4 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "latency" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml index 1b2d7216051..0dd401d2d56 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-AutoFp8-triton.yaml @@ -2,7 +2,6 @@ model_name: "Qwen/Qwen3-Coder-30B-A3B-Instruct-FP8" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton" env: - VLLM_USE_FLASHINFER_MOE_FP8: "0" VLLM_USE_DEEP_GEMM: "0" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-fi-cutlass.yaml index 48ab58c4611..fb52d3600eb 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block" accuracy_threshold: 0.85 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP8: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml index 3e30d4d154a..5bd907c0509 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-Fp8-CT-Block-triton.yaml @@ -2,7 +2,6 @@ model_name: "RedHatAI/Qwen3-30B-A3B-FP8-block" accuracy_threshold: 0.85 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=triton" env: - VLLM_USE_FLASHINFER_MOE_FP8: "0" VLLM_USE_DEEP_GEMM: "0" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml index 6edacc32975..3c1b20c242a 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml index 8e0b155fa70..094ec92f1e7 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-fi-trtllm.yaml @@ -2,7 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "latency" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml index 0d7884928ef..c38bc162eb2 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-CT-vllm-cutlass.yaml @@ -2,6 +2,4 @@ model_name: "RedHatAI/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "0" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml index 09e76e21ab4..0ebc68ad3ef 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-cutlass.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "throughput" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_cutlass" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml index a98afafbcde..491b3c82faf 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-fi-trtllm.yaml @@ -2,7 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "1" - VLLM_FLASHINFER_MOE_BACKEND: "latency" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=flashinfer_trtllm" diff --git a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml index a340b6fdae4..242c6ff529a 100644 --- a/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml +++ b/tests/evals/gsm8k/configs/moe-refactor/Qwen3-30B-A3B-NvFp4-ModelOpt-vllm-cutlass.yaml @@ -2,6 +2,4 @@ model_name: "nvidia/Qwen3-30B-A3B-NVFP4" accuracy_threshold: 0.88 num_questions: 1319 num_fewshot: 5 -server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2" -env: - VLLM_USE_FLASHINFER_MOE_FP4: "0" +server_args: "--enforce-eager --max-model-len 8192 --tensor-parallel-size 2 --moe-backend=cutlass" diff --git a/tests/quantization/test_blackwell_moe.py b/tests/quantization/test_blackwell_moe.py index 07da2b454e6..3a44ff4236a 100644 --- a/tests/quantization/test_blackwell_moe.py +++ b/tests/quantization/test_blackwell_moe.py @@ -85,34 +85,34 @@ def can_initialize( ) ) def test_llama4_fp8_tensor_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", + hf_overrides=HF_OVERRIDE_MM, + extra_args=["--moe-backend=flashinfer_cutlass"], ) def test_llama4_fp8_tensor_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", hf_overrides=HF_OVERRIDE_MM + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP8", + hf_overrides=HF_OVERRIDE_MM, + extra_args=["--moe-backend=flashinfer_trtllm"], ) def test_llama4_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", hf_overrides=HF_OVERRIDE_MM + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", + hf_overrides=HF_OVERRIDE_MM, + extra_args=["--moe-backend=flashinfer_cutlass"], ) def test_llama4_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") can_initialize( - "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", hf_overrides=HF_OVERRIDE_MM + "nvidia/Llama-4-Scout-17B-16E-Instruct-FP4", + hf_overrides=HF_OVERRIDE_MM, + extra_args=["--moe-backend=flashinfer_trtllm"], ) @@ -120,8 +120,11 @@ def test_llama4_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_DEEP_GEMM", "1") - can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "deepseek-ai/DeepSeek-V3.1", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=deep_gemm"], + ) @pytest.mark.skip( @@ -131,27 +134,35 @@ def test_deepseek_fp8_block_moe_deep_gemm(monkeypatch: pytest.MonkeyPatch): ) ) def test_deepseek_fp8_block_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") - can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "deepseek-ai/DeepSeek-V3.1", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_cutlass"], + ) def test_deepseek_fp8_block_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP8", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") - can_initialize("deepseek-ai/DeepSeek-V3.1", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "deepseek-ai/DeepSeek-V3.1", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_trtllm"], + ) def test_deepseek_nvfp4_moe_flashinfer_cutlass(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "throughput") - can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "nvidia/DeepSeek-R1-0528-FP4-v2", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_cutlass"], + ) def test_deepseek_nvfp4_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP4", "1") - monkeypatch.setenv("VLLM_FLASHINFER_MOE_BACKEND", "latency") - can_initialize("nvidia/DeepSeek-R1-0528-FP4-v2", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "nvidia/DeepSeek-R1-0528-FP4-v2", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_trtllm"], + ) ## GPT-OSS ## @@ -184,5 +195,8 @@ def test_gptoss_eager(monkeypatch: pytest.MonkeyPatch): def test_qwen3_next_bf16_moe_flashinfer_trtllm(monkeypatch: pytest.MonkeyPatch): - monkeypatch.setenv("VLLM_USE_FLASHINFER_MOE_FP16", "1") - can_initialize("Qwen/Qwen3-Next-80B-A3B-Instruct", hf_overrides=HF_OVERRIDE_TEXT) + can_initialize( + "Qwen/Qwen3-Next-80B-A3B-Instruct", + hf_overrides=HF_OVERRIDE_TEXT, + extra_args=["--moe-backend=flashinfer_trtllm"], + ) diff --git a/vllm/config/kernel.py b/vllm/config/kernel.py index 0730e464927..3c08ef88201 100644 --- a/vllm/config/kernel.py +++ b/vllm/config/kernel.py @@ -2,13 +2,25 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project from collections.abc import Callable -from typing import Any +from typing import Any, Literal from pydantic import Field, field_validator from vllm.config.utils import config from vllm.utils.hashing import safe_hash +MoEBackend = Literal[ + "auto", + "triton", + "deep_gemm", + "cutlass", + "flashinfer_trtllm", + "flashinfer_cutlass", + "flashinfer_cutedsl", + "marlin", + "aiter", +] + @config class KernelConfig: @@ -17,6 +29,26 @@ class KernelConfig: enable_flashinfer_autotune: bool = Field(default=None) """If True, run FlashInfer autotuning during kernel warmup.""" + moe_backend: MoEBackend = "auto" + """Backend for MoE expert computation kernels. Available options: + + - "auto": Automatically select the best backend based on model and hardware\n + - "triton": Use Triton-based fused MoE kernels\n + - "deep_gemm": Use DeepGEMM kernels (FP8 block-quantized only)\n + - "cutlass": Use vLLM CUTLASS kernels\n + - "flashinfer_trtllm": Use FlashInfer with TRTLLM-GEN kernels\n + - "flashinfer_cutlass": Use FlashInfer with CUTLASS kernels\n + - "flashinfer_cutedsl": Use FlashInfer with CuteDSL kernels (FP4 only)\n + - "marlin": Use Marlin kernels (weight-only quantization)\n + - "aiter": Use AMD AITer kernels (ROCm only)""" + + @field_validator("moe_backend", mode="before") + @classmethod + def _normalize_moe_backend(cls, value: Any) -> Any: + if isinstance(value, str): + return value.lower().replace("-", "_") + return value + def compute_hash(self) -> str: """ WARNING: Whenever a new field is added to this config, diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index ca76454d6d1..036178887db 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -70,6 +70,7 @@ from vllm.config.cache import ( PrefixCachingHashAlgo, ) from vllm.config.device import Device +from vllm.config.kernel import MoEBackend from vllm.config.lora import MaxLoRARanks from vllm.config.model import ( ConvertOption, @@ -416,6 +417,7 @@ class EngineArgs: data_parallel_external_lb: bool = False data_parallel_backend: DataParallelBackend = ParallelConfig.data_parallel_backend enable_expert_parallel: bool = ParallelConfig.enable_expert_parallel + moe_backend: MoEBackend = KernelConfig.moe_backend all2all_backend: All2AllBackend = ParallelConfig.all2all_backend enable_dbo: bool = ParallelConfig.enable_dbo ubatch_size: int = ParallelConfig.ubatch_size @@ -1227,6 +1229,9 @@ class EngineArgs: "--enable-flashinfer-autotune", **kernel_kwargs["enable_flashinfer_autotune"], ) + moe_backend_kwargs = kernel_kwargs["moe_backend"] + moe_backend_kwargs["type"] = lambda s: s.lower().replace("-", "_") + kernel_group.add_argument("--moe-backend", **moe_backend_kwargs) # vLLM arguments vllm_kwargs = get_kwargs(VllmConfig) @@ -1817,6 +1822,8 @@ class EngineArgs: "are mutually exclusive" ) kernel_config.enable_flashinfer_autotune = self.enable_flashinfer_autotune + if self.moe_backend != "auto": + kernel_config.moe_backend = self.moe_backend load_config = self.create_load_config() diff --git a/vllm/model_executor/layers/fused_moe/config.py b/vllm/model_executor/layers/fused_moe/config.py index 22e71d39101..87e1e244b3a 100644 --- a/vllm/model_executor/layers/fused_moe/config.py +++ b/vllm/model_executor/layers/fused_moe/config.py @@ -1066,7 +1066,6 @@ class FusedMoEParallelConfig: - Comment: There are 2 engine instances and the experts are split between the 4 devices. """ - use_ep = ( dp_size_ * pcp_size_ * tp_size_ > 1 and vllm_parallel_config.enable_expert_parallel @@ -1155,6 +1154,7 @@ class FusedMoEConfig: # Defaults to in_dtype if not specified. router_logits_dtype: torch.dtype | None = None + moe_backend: str = "auto" max_num_tokens: int = envs.VLLM_MOE_DP_CHUNK_SIZE has_bias: bool = False is_act_and_mul: bool = True diff --git a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py index f5a3da43878..a4cee76f716 100644 --- a/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py +++ b/vllm/model_executor/layers/fused_moe/deepep_ll_prepare_finalize.py @@ -198,7 +198,6 @@ class DeepEPLLPrepareAndFinalize(mk.FusedMoEPrepareAndFinalize): x = x[0].permute(2, 0, 1) num_experts, max_tokens, hidden_dim_by_2 = x.shape hidden_dim = hidden_dim_by_2 * 2 - assert envs.VLLM_FLASHINFER_MOE_BACKEND == "masked_gemm" logger.info_once( "Quantization is fused with DeepEP nvfp4 dispatch for " "FlashInfer CUTEDSL as VLLM_DEEPEPLL_NVFP4_DISPATCH==1" diff --git a/vllm/model_executor/layers/fused_moe/layer.py b/vllm/model_executor/layers/fused_moe/layer.py index 6cb3dae2673..679b79ce971 100644 --- a/vllm/model_executor/layers/fused_moe/layer.py +++ b/vllm/model_executor/layers/fused_moe/layer.py @@ -550,6 +550,7 @@ class FusedMoE(CustomOp): num_logical_experts=self.logical_num_experts, moe_parallel_config=self.moe_parallel_config, in_dtype=moe_in_dtype, + moe_backend=vllm_config.kernel_config.moe_backend, router_logits_dtype=router_logits_dtype, max_num_tokens=envs.VLLM_MOE_DP_CHUNK_SIZE, has_bias=has_bias, diff --git a/vllm/model_executor/layers/fused_moe/oracle/fp8.py b/vllm/model_executor/layers/fused_moe/oracle/fp8.py index 2432209899e..6f961df07d8 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/fp8.py +++ b/vllm/model_executor/layers/fused_moe/oracle/fp8.py @@ -7,6 +7,7 @@ import torch import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm import envs from vllm._aiter_ops import rocm_aiter_ops +from vllm.config.kernel import MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, @@ -180,6 +181,25 @@ def backend_to_kernel_cls( raise ValueError(f"Unknown FP8 MoE backend: {backend.value}") +def map_fp8_backend(runner_backend: MoEBackend) -> Fp8MoeBackend: + """Map user's MoEBackend to Fp8MoeBackend.""" + mapping = { + "triton": Fp8MoeBackend.TRITON, + "deep_gemm": Fp8MoeBackend.DEEPGEMM, + "cutlass": Fp8MoeBackend.VLLM_CUTLASS, + "flashinfer_trtllm": Fp8MoeBackend.FLASHINFER_TRTLLM, + "flashinfer_cutlass": Fp8MoeBackend.FLASHINFER_CUTLASS, + "marlin": Fp8MoeBackend.MARLIN, + "aiter": Fp8MoeBackend.AITER, + } + if backend := mapping.get(runner_backend): + return backend + raise ValueError( + f"moe_backend='{runner_backend}' is not supported for FP8 MoE. " + f"Expected one of {list(mapping.keys())}." + ) + + def select_fp8_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, @@ -242,6 +262,45 @@ def select_fp8_moe_backend( return backend, k_cls raise ValueError(_make_log_unsupported(backend, reason)) + # Handle explicit moe_backend from user. + runner_backend = config.moe_backend + if runner_backend != "auto": + requested_backend = map_fp8_backend(runner_backend) + # For batched activation format, use batched variants if available. + if activation_format == mk.FusedMoEActivationFormat.BatchedExperts: + if requested_backend == Fp8MoeBackend.DEEPGEMM: + requested_backend = Fp8MoeBackend.BATCHED_DEEPGEMM + elif requested_backend == Fp8MoeBackend.TRITON: + requested_backend = Fp8MoeBackend.BATCHED_TRITON + elif requested_backend == Fp8MoeBackend.VLLM_CUTLASS: + requested_backend = Fp8MoeBackend.BATCHED_VLLM_CUTLASS + + if ( + requested_backend + in [ + Fp8MoeBackend.VLLM_CUTLASS, + Fp8MoeBackend.BATCHED_VLLM_CUTLASS, + ] + and not allow_vllm_cutlass + ): + raise ValueError( + "vLLM CUTLASS FP8 MoE backend is disabled for this configuration." + ) + + # Handle FLASHINFER_TRTLLM specially (no kernel class). + if requested_backend == Fp8MoeBackend.FLASHINFER_TRTLLM: + supported, reason = is_supported_config_trtllm_fp8( + config, weight_key, activation_key, activation_format + ) + if supported: + logger.info_once(_make_log_backend(requested_backend)) + return requested_backend, None + raise ValueError(_make_log_unsupported(requested_backend, reason)) + + return _return_or_raise( + requested_backend, config, weight_key, activation_key, activation_format + ) + # Handle explicit FlashInfer FP8 configuration. if envs.is_set("VLLM_USE_FLASHINFER_MOE_FP8"): if not envs.VLLM_USE_FLASHINFER_MOE_FP8: diff --git a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py index dc3ac61ad14..ee7db88ccbd 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py +++ b/vllm/model_executor/layers/fused_moe/oracle/nvfp4.py @@ -6,6 +6,7 @@ import torch import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk +from vllm.config.kernel import MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.all2all_utils import ( maybe_make_prepare_finalize, @@ -103,6 +104,23 @@ def backend_to_kernel_cls( raise ValueError(f"Unknown NvFP4 MoE backend: {backend.value}") +def map_nvfp4_backend(runner_backend: MoEBackend) -> NvFp4MoeBackend: + """Map user's MoEBackend to NvFp4MoeBackend.""" + mapping = { + "cutlass": NvFp4MoeBackend.VLLM_CUTLASS, + "flashinfer_trtllm": NvFp4MoeBackend.FLASHINFER_TRTLLM, + "flashinfer_cutlass": NvFp4MoeBackend.FLASHINFER_CUTLASS, + "flashinfer_cutedsl": NvFp4MoeBackend.FLASHINFER_CUTEDSL, + "marlin": NvFp4MoeBackend.MARLIN, + } + if backend := mapping.get(runner_backend): + return backend + raise ValueError( + f"moe_backend='{runner_backend}' is not supported for NvFP4 MoE. " + f"Expected one of {list(mapping.keys())}." + ) + + def select_nvfp4_moe_backend( config: FusedMoEConfig, weight_key: QuantKey | None, @@ -170,6 +188,23 @@ def select_nvfp4_moe_backend( return backend, k_cls raise ValueError(_make_log_unsupported(backend, reason)) + # Handle explicit moe_backend from user. + runner_backend = config.moe_backend + if runner_backend != "auto": + requested_backend = map_nvfp4_backend(runner_backend) + if requested_backend == NvFp4MoeBackend.FLASHINFER_TRTLLM: + supported, reason = is_supported_config_trtllm( + config, weight_key, activation_key, activation_format + ) + if supported: + logger.info_once(_make_log_backend(requested_backend)) + return requested_backend, None + raise ValueError(_make_log_unsupported(requested_backend, reason)) + + return _return_or_raise( + requested_backend, config, weight_key, activation_key, activation_format + ) + if envs.is_set("VLLM_USE_FLASHINFER_MOE_FP4"): if not envs.VLLM_USE_FLASHINFER_MOE_FP4: # If the user rejects FlashInfer remove those backends. diff --git a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py index 61aaa692777..1c582bcdc53 100644 --- a/vllm/model_executor/layers/fused_moe/oracle/unquantized.py +++ b/vllm/model_executor/layers/fused_moe/oracle/unquantized.py @@ -9,6 +9,7 @@ from torch.nn import Module import vllm.envs as envs import vllm.model_executor.layers.fused_moe.modular_kernel as mk from vllm._aiter_ops import rocm_aiter_ops +from vllm.config.kernel import MoEBackend from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe.config import ( FusedMoEConfig, @@ -51,6 +52,22 @@ UNSUPPORTED_BACKEND = [ ] +def map_unquantized_backend(runner_backend: MoEBackend) -> UnquantizedMoeBackend: + """Map user's MoEBackend to UnquantizedMoeBackend.""" + mapping = { + "triton": UnquantizedMoeBackend.TRITON, + "flashinfer_trtllm": UnquantizedMoeBackend.FLASHINFER_TRTLLM, + "flashinfer_cutlass": UnquantizedMoeBackend.FLASHINFER_CUTLASS, + "aiter": UnquantizedMoeBackend.AITER, + } + if backend := mapping.get(runner_backend): + return backend + raise ValueError( + f"moe_backend='{runner_backend}' is not supported for unquantized MoE. " + f"Expected one of {list(mapping.keys())}." + ) + + def select_unquantized_moe_backend( moe_config: FusedMoEConfig, use_ep: bool, @@ -64,8 +81,6 @@ def select_unquantized_moe_backend( def _make_log_backend(backend: UnquantizedMoeBackend): return f"Using {backend.value} backend for Unquantized MoE" - rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() - activation_format = ( mk.FusedMoEActivationFormat.BatchedExperts if moe_config.moe_parallel_config.use_batched_activation_format @@ -77,20 +92,49 @@ def select_unquantized_moe_backend( moe_config=moe_config, activation_format=activation_format, ) - flashinfer_trtllm_moe_enabled = ( - has_flashinfer() - and envs.VLLM_USE_FLASHINFER_MOE_FP16 - and trtllm_supported - and envs.VLLM_FLASHINFER_MOE_BACKEND == "latency" - ) + flashinfer_trtllm_available = has_flashinfer() and trtllm_supported # FlashInfer CUTLASS MoE is only supported on Hopper and later GPUS - flashinfer_cutlass_moe_enabled = ( + flashinfer_cutlass_available = ( has_flashinfer_cutlass_fused_moe() - and envs.VLLM_USE_FLASHINFER_MOE_FP16 and use_ep and (not use_dp) and current_platform.has_device_capability(90) ) + flashinfer_trtllm_moe_enabled = ( + flashinfer_trtllm_available + and envs.VLLM_USE_FLASHINFER_MOE_FP16 + and envs.VLLM_FLASHINFER_MOE_BACKEND == "latency" + ) + flashinfer_cutlass_moe_enabled = ( + flashinfer_cutlass_available and envs.VLLM_USE_FLASHINFER_MOE_FP16 + ) + rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled() + + # Handle explicit moe_backend from user. + runner_backend = moe_config.moe_backend + if runner_backend != "auto": + requested_backend = map_unquantized_backend(runner_backend) + if requested_backend == UnquantizedMoeBackend.FLASHINFER_TRTLLM: + if not flashinfer_trtllm_available: + raise ValueError( + "FlashInfer TRTLLM MoE backend is not available for this " + "configuration." + ) + elif requested_backend == UnquantizedMoeBackend.FLASHINFER_CUTLASS: + if not flashinfer_cutlass_available: + raise ValueError( + "FlashInfer CUTLASS MoE backend is not available for this " + "configuration." + ) + elif requested_backend == UnquantizedMoeBackend.AITER and not ( + current_platform.is_rocm() and rocm_aiter_moe_enabled + ): + raise ValueError( + "ROCm AITer MoE backend is not available for this configuration." + ) + logger.info_once(_make_log_backend(requested_backend), scope="local") + return requested_backend + if current_platform.is_rocm(): if rocm_aiter_moe_enabled: backend = UnquantizedMoeBackend.AITER From 9511a3f8eec6992d8834ad85af855683bd74ba29 Mon Sep 17 00:00:00 2001 From: Lucas Wilkinson Date: Wed, 25 Feb 2026 21:01:10 -0500 Subject: [PATCH 15/43] [Bugfix] Fix AttributeError in SMControlContextManager (#35338) Signed-off-by: Lucas Wilkinson --- vllm/v1/worker/gpu_ubatch_wrapper.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu_ubatch_wrapper.py b/vllm/v1/worker/gpu_ubatch_wrapper.py index 45ba1bef9f2..754f2981c9f 100644 --- a/vllm/v1/worker/gpu_ubatch_wrapper.py +++ b/vllm/v1/worker/gpu_ubatch_wrapper.py @@ -74,7 +74,7 @@ class SMControlContextManager: "SM control is currently only supported on CUDA" ) - total_sms = num_compute_units(torch.cuda.current_device().index) + total_sms = num_compute_units(torch.cuda.current_device()) assert comm_sms < total_sms self.total_sms = total_sms From 160424a937d373101818876103227cc986887b55 Mon Sep 17 00:00:00 2001 From: Seungmin Kim <8457324+ehfd@users.noreply.github.com> Date: Thu, 26 Feb 2026 11:15:51 +0900 Subject: [PATCH 16/43] [Bugfix] Fix CUDA compatibility path setting for both datacenter and consumer NVIDIA GPUs (#33992) Signed-off-by: Seungmin Kim <8457324+ehfd@users.noreply.github.com> Signed-off-by: Andrew Mello <19512127+88plug@users.noreply.github.com> Co-authored-by: 88plug <19512127+88plug@users.noreply.github.com> Co-authored-by: Michael Goin --- docker/Dockerfile | 12 +- .../installation/gpu.cuda.inc.md | 17 ++ docs/usage/troubleshooting.md | 27 ++- tests/cuda/test_cuda_compatibility_path.py | 187 ++++++++++++++++++ vllm/env_override.py | 82 ++++++++ vllm/envs.py | 14 ++ 6 files changed, 334 insertions(+), 5 deletions(-) create mode 100644 tests/cuda/test_cuda_compatibility_path.py diff --git a/docker/Dockerfile b/docker/Dockerfile index cc2ccc11cdc..717f27b6b23 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -132,8 +132,10 @@ ENV UV_LINK_MODE=copy # Verify GCC version RUN gcc --version -# Ensure CUDA compatibility library is loaded -RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig +# Enable CUDA forward compatibility by setting '-e VLLM_ENABLE_CUDA_COMPATIBILITY=1' +# Only needed for datacenter/professional GPUs with older drivers. +# See: https://docs.nvidia.com/deploy/cuda-compatibility/ +ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0 # ============================================================ # SLOW-CHANGING DEPENDENCIES BELOW @@ -560,8 +562,10 @@ ENV UV_HTTP_TIMEOUT=500 ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE=copy -# Ensure CUDA compatibility library is loaded -RUN echo "/usr/local/cuda-$(echo "$CUDA_VERSION" | cut -d. -f1,2)/compat/" > /etc/ld.so.conf.d/cuda-compat.conf && ldconfig +# Enable CUDA forward compatibility by setting '-e VLLM_ENABLE_CUDA_COMPATIBILITY=1' +# Only needed for datacenter/professional GPUs with older drivers. +# See: https://docs.nvidia.com/deploy/cuda-compatibility/ +ENV VLLM_ENABLE_CUDA_COMPATIBILITY=0 # ============================================================ # SLOW-CHANGING DEPENDENCIES BELOW diff --git a/docs/getting_started/installation/gpu.cuda.inc.md b/docs/getting_started/installation/gpu.cuda.inc.md index 661e0934eef..da8b7d3fa1d 100644 --- a/docs/getting_started/installation/gpu.cuda.inc.md +++ b/docs/getting_started/installation/gpu.cuda.inc.md @@ -297,6 +297,23 @@ You can add any other [engine-args](https://docs.vllm.ai/en/latest/configuration RUN uv pip install --system git+https://github.com/huggingface/transformers.git ``` +#### Running on Systems with Older CUDA Drivers + +vLLM's Docker image comes with [CUDA compatibility libraries](https://docs.nvidia.com/deploy/cuda-compatibility/index.html) pre-installed. This allows you to run vLLM on systems with NVIDIA drivers that are older than the CUDA Toolkit version used in the image, but only supports select professional and datacenter NVIDIA GPUs. + +To enable this feature, set the `VLLM_ENABLE_CUDA_COMPATIBILITY` environment variable to `1` or `true` when running the container: + +```bash +docker run --runtime nvidia --gpus all \ + -v ~/.cache/huggingface:/root/.cache/huggingface \ + -p 8000:8000 \ + --env "HF_TOKEN=" \ + --env "VLLM_ENABLE_CUDA_COMPATIBILITY=1" \ + vllm/vllm-openai +``` + +This will automatically configure `LD_LIBRARY_PATH` to point to the compatibility libraries before loading PyTorch and other dependencies. + # --8<-- [end:pre-built-images] # --8<-- [start:build-image-from-source] diff --git a/docs/usage/troubleshooting.md b/docs/usage/troubleshooting.md index 128c36b784d..814b03c1e38 100644 --- a/docs/usage/troubleshooting.md +++ b/docs/usage/troubleshooting.md @@ -318,7 +318,32 @@ This indicates vLLM failed to initialize the NCCL communicator, possibly due to ## CUDA error: the provided PTX was compiled with an unsupported toolchain -If you see an error like `RuntimeError: CUDA error: the provided PTX was compiled with an unsupported toolchain.`, it means that the CUDA PTX in vLLM's wheels was compiled with a toolchain unsupported by your system. The released vLLM wheels have to be compiled with a specific version of CUDA toolkit, and the compiled code might fail to run on lower versions of CUDA drivers. Read [cuda compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/) for more details. The solution is to install `cuda-compat` package from your package manager. For example, on Ubuntu, you can run `sudo apt-get install cuda-compat-12-9`, and then add `export LD_LIBRARY_PATH=/usr/local/cuda-12.9/compat:$LD_LIBRARY_PATH` to your `.bashrc` file. When successfully installed, you should see that the output of `nvidia-smi` will show `CUDA Version: 12.9`. Note that we use CUDA 12.9 as an example here, you may want to install a higher version of cuda-compat package in case vLLM's default CUDA version goes higher. +If you see an error like `RuntimeError: CUDA error: the provided PTX was compiled with an unsupported toolchain`, it means that the CUDA PTX in vLLM's wheels was compiled with a toolchain unsupported by your system. This section also applies if you get the error `RuntimeError: The NVIDIA driver on your system is too old`. + +The released vLLM wheels are compiled with a specific version of CUDA toolkit, and the compiled code might fail to run on lower versions of CUDA drivers. Read [CUDA compatibility](https://docs.nvidia.com/deploy/cuda-compatibility/) for more details. **This is only supported on select professional and datacenter NVIDIA GPUs.** + +If you are using the vLLM official Docker image, you can solve this by adding `-e VLLM_ENABLE_CUDA_COMPATIBILITY=1` to your `docker run` command. This will enable the pre-installed CUDA forward compatibility libraries. + +If you are running vLLM outside of Docker, the solution is to install the `cuda-compat` package from your package manager with the [CUDA repository](https://docs.nvidia.com/cuda/cuda-installation-guide-linux/) enabled. For example, on Ubuntu, you can run `sudo apt-get install cuda-compat-12-9`, and then set `export VLLM_ENABLE_CUDA_COMPATIBILITY=1` and `export VLLM_CUDA_COMPATIBILITY_PATH="/usr/local/cuda-12.9/compat"`. + +On Conda, you can install the `conda-forge::cuda-compat` package (e.g., `conda install -c conda-forge cuda-compat=12.9`), then after activating the environment, set `export VLLM_ENABLE_CUDA_COMPATIBILITY=1` and `export VLLM_CUDA_COMPATIBILITY_PATH="${CONDA_PREFIX}/cuda-compat"`. + +You can verify the configuration works by running a minimal Python script that initializes CUDA via vLLM: + +```bash +export VLLM_ENABLE_CUDA_COMPATIBILITY=1 +export VLLM_CUDA_COMPATIBILITY_PATH="/usr/local/cuda-12.9/compat" + +python3 - << 'EOF' +import vllm +import torch + +print(f"CUDA available: {torch.cuda.is_available()}") +print(f"CUDA device count: {torch.cuda.device_count()}") +EOF +``` + +Note that we use CUDA 12.9 as an example here, and you may want to install a higher version of cuda-compat package in case vLLM's default CUDA version goes higher. ## ptxas fatal: Value 'sm_110a' is not defined for option 'gpu-name' diff --git a/tests/cuda/test_cuda_compatibility_path.py b/tests/cuda/test_cuda_compatibility_path.py new file mode 100644 index 00000000000..837d2c49cfb --- /dev/null +++ b/tests/cuda/test_cuda_compatibility_path.py @@ -0,0 +1,187 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for CUDA forward compatibility path logic in env_override.py. + +Verifies the opt-in LD_LIBRARY_PATH manipulation for CUDA compat libs, +including env var parsing, path detection, and deduplication. +""" + +import os +from unittest.mock import patch + +import pytest + +# Import the functions directly (they're module-level in env_override) +# We must import them without triggering the module-level side effects, +# so we import the functions by name after the module is already loaded. +from vllm.env_override import ( + _get_torch_cuda_version, + _maybe_set_cuda_compatibility_path, +) + + +class TestCudaCompatibilityEnvParsing: + """Test VLLM_ENABLE_CUDA_COMPATIBILITY env var parsing.""" + + def test_disabled_by_default(self, monkeypatch): + """Compat path is NOT set when env var is absent.""" + monkeypatch.delenv("VLLM_ENABLE_CUDA_COMPATIBILITY", raising=False) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + _maybe_set_cuda_compatibility_path() + assert ( + "LD_LIBRARY_PATH" not in os.environ + or os.environ.get("LD_LIBRARY_PATH", "") == "" + ) + + @pytest.mark.parametrize("value", ["0", "false", "False", "no", ""]) + def test_disabled_values(self, monkeypatch, value): + """Various falsy values should not activate compat path.""" + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", value) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + _maybe_set_cuda_compatibility_path() + # LD_LIBRARY_PATH should not be set (or remain empty) + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + assert "compat" not in ld_path + + @pytest.mark.parametrize("value", ["1", "true", "True", " 1 ", " TRUE "]) + def test_enabled_values_with_valid_path(self, monkeypatch, tmp_path, value): + """Truthy values activate compat path when a valid path exists.""" + compat_dir = tmp_path / "compat" + compat_dir.mkdir() + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", value) + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir)) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + _maybe_set_cuda_compatibility_path() + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + assert str(compat_dir) in ld_path + + +class TestCudaCompatibilityPathDetection: + """Test path detection: custom override, conda, default.""" + + def test_custom_path_override(self, monkeypatch, tmp_path): + """VLLM_CUDA_COMPATIBILITY_PATH takes highest priority.""" + custom_dir = tmp_path / "my-compat" + custom_dir.mkdir() + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(custom_dir)) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + _maybe_set_cuda_compatibility_path() + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + assert ld_path.startswith(str(custom_dir)) + + def test_conda_prefix_fallback(self, monkeypatch, tmp_path): + """Falls back to $CONDA_PREFIX/cuda-compat if custom not set.""" + conda_dir = tmp_path / "conda-env" + compat_dir = conda_dir / "cuda-compat" + compat_dir.mkdir(parents=True) + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.delenv("VLLM_CUDA_COMPATIBILITY_PATH", raising=False) + monkeypatch.setenv("CONDA_PREFIX", str(conda_dir)) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + _maybe_set_cuda_compatibility_path() + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + assert str(compat_dir) in ld_path + + def test_no_valid_path_does_nothing(self, monkeypatch): + """When enabled but no valid path exists, LD_LIBRARY_PATH unchanged.""" + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", "/nonexistent/path") + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + with patch("vllm.env_override._get_torch_cuda_version", return_value=None): + _maybe_set_cuda_compatibility_path() + assert os.environ.get("LD_LIBRARY_PATH", "") == "" + + def test_default_cuda_path_fallback(self, monkeypatch, tmp_path): + """Falls back to /usr/local/cuda-{ver}/compat via torch version.""" + fake_cuda = tmp_path / "cuda-12.8" / "compat" + fake_cuda.mkdir(parents=True) + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.delenv("VLLM_CUDA_COMPATIBILITY_PATH", raising=False) + monkeypatch.delenv("CONDA_PREFIX", raising=False) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + with ( + patch("vllm.env_override._get_torch_cuda_version", return_value="12.8"), + patch( + "vllm.env_override.os.path.isdir", + side_effect=lambda p: p == "/usr/local/cuda-12.8/compat" + or os.path.isdir(p), + ), + ): + _maybe_set_cuda_compatibility_path() + ld_path = os.environ.get("LD_LIBRARY_PATH", "") + assert "/usr/local/cuda-12.8/compat" in ld_path + + +class TestCudaCompatibilityLdPathManipulation: + """Test LD_LIBRARY_PATH prepend and deduplication logic.""" + + def test_prepends_to_empty_ld_path(self, monkeypatch, tmp_path): + """Compat path is set when LD_LIBRARY_PATH is empty.""" + compat_dir = tmp_path / "compat" + compat_dir.mkdir() + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir)) + monkeypatch.delenv("LD_LIBRARY_PATH", raising=False) + _maybe_set_cuda_compatibility_path() + assert os.environ["LD_LIBRARY_PATH"] == str(compat_dir) + + def test_prepends_to_existing_ld_path(self, monkeypatch, tmp_path): + """Compat path is prepended before existing entries.""" + compat_dir = tmp_path / "compat" + compat_dir.mkdir() + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir)) + monkeypatch.setenv("LD_LIBRARY_PATH", "/usr/lib:/other/lib") + _maybe_set_cuda_compatibility_path() + ld_path = os.environ["LD_LIBRARY_PATH"] + parts = ld_path.split(os.pathsep) + assert parts[0] == str(compat_dir) + assert "/usr/lib" in parts + assert "/other/lib" in parts + + def test_deduplicates_existing_compat_path(self, monkeypatch, tmp_path): + """If compat path already in LD_LIBRARY_PATH, move to front.""" + compat_dir = tmp_path / "compat" + compat_dir.mkdir() + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir)) + monkeypatch.setenv( + "LD_LIBRARY_PATH", + f"/usr/lib:{compat_dir}:/other/lib", + ) + _maybe_set_cuda_compatibility_path() + ld_path = os.environ["LD_LIBRARY_PATH"] + parts = ld_path.split(os.pathsep) + assert parts[0] == str(compat_dir) + assert parts.count(str(compat_dir)) == 1 + + def test_already_at_front_is_noop(self, monkeypatch, tmp_path): + """If compat path is already first, don't modify LD_LIBRARY_PATH.""" + compat_dir = tmp_path / "compat" + compat_dir.mkdir() + original = f"{compat_dir}:/usr/lib" + monkeypatch.setenv("VLLM_ENABLE_CUDA_COMPATIBILITY", "1") + monkeypatch.setenv("VLLM_CUDA_COMPATIBILITY_PATH", str(compat_dir)) + monkeypatch.setenv("LD_LIBRARY_PATH", original) + _maybe_set_cuda_compatibility_path() + assert os.environ["LD_LIBRARY_PATH"] == original + + +class TestGetTorchCudaVersion: + """Test _get_torch_cuda_version() helper.""" + + def test_returns_string_when_torch_available(self): + """Should return a CUDA version string like '12.8'.""" + version = _get_torch_cuda_version() + # torch is installed in vllm's environment + assert version is None or isinstance(version, str) + + def test_returns_none_when_torch_missing(self): + """Should return None when torch is not importable.""" + with patch( + "vllm.env_override.importlib.util.find_spec", + return_value=None, + ): + assert _get_torch_cuda_version() is None diff --git a/vllm/env_override.py b/vllm/env_override.py index e5a40dc3cd8..181d000a68a 100644 --- a/vllm/env_override.py +++ b/vllm/env_override.py @@ -1,7 +1,89 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# ruff: noqa: E402 +import importlib.util import os + +def _get_torch_cuda_version(): + """Peripheral function to _maybe_set_cuda_compatibility_path(). + PyTorch version must not be determined by importing directly + because it will trigger the CUDA initialization, losing the + chance to set the LD_LIBRARY_PATH beforehand. + """ + try: + spec = importlib.util.find_spec("torch") + if not spec: + return None + if spec.origin: + torch_root = os.path.dirname(spec.origin) + elif spec.submodule_search_locations: + torch_root = spec.submodule_search_locations[0] + else: + return None + version_path = os.path.join(torch_root, "version.py") + if not os.path.exists(version_path): + return None + # Load the version module without importing torch + ver_spec = importlib.util.spec_from_file_location("torch.version", version_path) + if not ver_spec or not ver_spec.loader: + return None + module = importlib.util.module_from_spec(ver_spec) + # Avoid registering in sys.modules to not confuse future imports + ver_spec.loader.exec_module(module) + return getattr(module, "cuda", None) + except Exception: + return None + + +def _maybe_set_cuda_compatibility_path(): + """Set LD_LIBRARY_PATH for CUDA forward compatibility if enabled. + + Must run before 'import torch' since torch loads CUDA shared libraries + at import time and the dynamic linker only consults LD_LIBRARY_PATH when + a library is first loaded. + + CUDA forward compatibility is only supported on select professional and + datacenter NVIDIA GPUs. Consumer GPUs (GeForce, RTX) do not support it + and will get Error 803 if compat libs are loaded. + """ + enable = os.environ.get("VLLM_ENABLE_CUDA_COMPATIBILITY", "0").strip().lower() in ( + "1", + "true", + ) + if not enable: + return + + cuda_compat_path = os.environ.get("VLLM_CUDA_COMPATIBILITY_PATH", "") + if not cuda_compat_path or not os.path.isdir(cuda_compat_path): + conda_prefix = os.environ.get("CONDA_PREFIX", "") + conda_compat = os.path.join(conda_prefix, "cuda-compat") + if conda_prefix and os.path.isdir(conda_compat): + cuda_compat_path = conda_compat + if not cuda_compat_path or not os.path.isdir(cuda_compat_path): + torch_cuda_version = _get_torch_cuda_version() + if torch_cuda_version: + default_path = f"/usr/local/cuda-{torch_cuda_version}/compat" + if os.path.isdir(default_path): + cuda_compat_path = default_path + if not cuda_compat_path or not os.path.isdir(cuda_compat_path): + return + + norm_path = os.path.normpath(cuda_compat_path) + existing = os.environ.get("LD_LIBRARY_PATH", "") + ld_paths = existing.split(os.pathsep) if existing else [] + + if ld_paths and ld_paths[0] and os.path.normpath(ld_paths[0]) == norm_path: + return # Already at the front + + new_paths = [norm_path] + [ + p for p in ld_paths if not p or os.path.normpath(p) != norm_path + ] + os.environ["LD_LIBRARY_PATH"] = os.pathsep.join(new_paths) + + +_maybe_set_cuda_compatibility_path() + import torch from vllm.logger import init_logger diff --git a/vllm/envs.py b/vllm/envs.py index 0d8cf021e4f..d62438d5735 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -239,6 +239,8 @@ if TYPE_CHECKING: VLLM_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False VLLM_DISABLE_LOG_LOGO: bool = False VLLM_LORA_DISABLE_PDL: bool = False + VLLM_ENABLE_CUDA_COMPATIBILITY: bool = False + VLLM_CUDA_COMPATIBILITY_PATH: str | None = None def get_default_cache_root(): @@ -1591,6 +1593,16 @@ environment_variables: dict[str, Callable[[], Any]] = { # Disable PDL for LoRA, as enabling PDL with LoRA on SM100 causes # Triton compilation to fail. "VLLM_LORA_DISABLE_PDL": lambda: bool(int(os.getenv("VLLM_LORA_DISABLE_PDL", "0"))), + # Enable CUDA compatibility mode for datacenter GPUs with older + # driver versions than the CUDA toolkit major version of vLLM. + "VLLM_ENABLE_CUDA_COMPATIBILITY": lambda: ( + os.environ.get("VLLM_ENABLE_CUDA_COMPATIBILITY", "0").strip().lower() + in ("1", "true") + ), + # Path to the CUDA compatibility libraries when CUDA compatibility is enabled. + "VLLM_CUDA_COMPATIBILITY_PATH": lambda: os.environ.get( + "VLLM_CUDA_COMPATIBILITY_PATH", None + ), } @@ -1731,6 +1743,8 @@ def compile_factors() -> dict[str, object]: "VLLM_CPU_MOE_PREPACK", "VLLM_CPU_SGL_KERNEL", "VLLM_TEST_FORCE_LOAD_FORMAT", + "VLLM_ENABLE_CUDA_COMPATIBILITY", + "VLLM_CUDA_COMPATIBILITY_PATH", "LOCAL_RANK", "CUDA_VISIBLE_DEVICES", "NO_COLOR", From 86c3b5a808506e325fd7e59d86d83170fc98c93c Mon Sep 17 00:00:00 2001 From: "Roberto L. Castro" <38211239+LopezCastroRoberto@users.noreply.github.com> Date: Thu, 26 Feb 2026 03:32:50 +0100 Subject: [PATCH 17/43] [BugFix] Fix fp4 quant kernel on CUDA 12.8 (#35210) Signed-off-by: LopezCastroRoberto --- .../fp4/activation_nvfp4_quant_fusion_kernels.cu | 6 ++++-- csrc/quantization/fp4/nvfp4_quant_kernels.cu | 12 +++++++----- 2 files changed, 11 insertions(+), 7 deletions(-) diff --git a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu index d0264c4d154..8583b79fd58 100644 --- a/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu +++ b/csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu @@ -107,7 +107,9 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) (uint64_t(out_val.hi) << 32) | uint64_t(out_val.lo); reinterpret_cast(out)[outOffset >> 1] = packed64; } else { - out[inOffset] = out_val; + int64_t outOffset = + rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; + out[outOffset] = out_val; } } } @@ -140,7 +142,7 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output, // [..., d] int const numBlocksPerSM = vllm_runtime_blocks_per_sm(static_cast(block.x)); - int sf_n_unpadded = int(n / CVT_FP4_SF_VEC_SIZE); + int sf_n_unpadded = int(n / CVT_FP4_ELTS_PER_THREAD); int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast(block.x)); int grid_x = std::min( diff --git a/csrc/quantization/fp4/nvfp4_quant_kernels.cu b/csrc/quantization/fp4/nvfp4_quant_kernels.cu index c27fb69d44b..b521b4707a4 100644 --- a/csrc/quantization/fp4/nvfp4_quant_kernels.cu +++ b/csrc/quantization/fp4/nvfp4_quant_kernels.cu @@ -109,7 +109,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) template __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) cvt_fp16_to_fp4_sf_major(int32_t numRows, int32_t numCols, - int32_t sf_n_unpadded, Type const* __restrict__ in, + int32_t sf_n_unpadded, int32_t num_packed_cols, + Type const* __restrict__ in, float const* __restrict__ SFScale, uint32_t* __restrict__ out, uint32_t* __restrict__ SFout) { @@ -131,7 +132,7 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512)) // Iterate over all rows and cols including padded ones - // ensures we visit every single scale factor address to initialize it. for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) { - if (colIdx < sf_n_unpadded) { + if (colIdx < num_packed_cols) { PackedVec in_vec; int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx; @@ -222,7 +223,8 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, reinterpret_cast(sf_out)); }); } else { - int grid_y = vllm::div_round_up(sf_n_unpadded, static_cast(block.x)); + int num_packed_cols = n / CVT_FP4_ELTS_PER_THREAD; + int grid_y = vllm::div_round_up(num_packed_cols, static_cast(block.x)); int grid_x = std::min( m, std::max(1, (multiProcessorCount * numBlocksPerSM) / grid_y)); dim3 grid(grid_x, grid_y); @@ -232,8 +234,8 @@ void scaled_fp4_quant_sm1xxa(torch::Tensor const& output, auto input_ptr = static_cast(input.data_ptr()); // NOTE: We don't support e8m0 scales at this moment. vllm::cvt_fp16_to_fp4_sf_major - <<>>(m, n, sf_n_unpadded, input_ptr, - input_sf_ptr, + <<>>(m, n, sf_n_unpadded, num_packed_cols, + input_ptr, input_sf_ptr, reinterpret_cast(output_ptr), reinterpret_cast(sf_out)); }); From 2aa414040243bc24447aa5a4f244f4104064d539 Mon Sep 17 00:00:00 2001 From: hujiaxin0 <524446785@qq.com> Date: Thu, 26 Feb 2026 11:08:09 +0800 Subject: [PATCH 18/43] openpangu-vl support video input (#34134) Signed-off-by: hujiaxin <524446785@qq.com> Signed-off-by: Emilie1001 <79921183+Emilie1001@users.noreply.github.com> Co-authored-by: Emilie1001 <79921183+Emilie1001@users.noreply.github.com> Co-authored-by: Isotr0py --- vllm/multimodal/video.py | 87 ++++++++++++++++++++++++++++++++++++++++ 1 file changed, 87 insertions(+) diff --git a/vllm/multimodal/video.py b/vllm/multimodal/video.py index f123799ca90..fb4e19fa674 100644 --- a/vllm/multimodal/video.py +++ b/vllm/multimodal/video.py @@ -747,3 +747,90 @@ class Molmo2VideoBackend(VideoLoader): **kwargs, ) return out + + +@VIDEO_LOADER_REGISTRY.register("openpangu") +class OpenCVDynamicOpenPanguVideoBackend(OpenCVVideoBackend): + @classmethod + def load_bytes( + cls, + data: bytes, + num_frames: int = 32, + fps: int = 1, + max_duration: int = 300, + frame_recovery: bool = False, + **kwargs, + ) -> tuple[npt.NDArray, dict[str, Any]]: + """ + Load video frames with dynamic sampling based on duration. + Assume that total_num_frames = 10 and fps = 1. + The timestamp of frame 0 is 0.0. + The timestamp of frame 1 is 1.0.… + The timestamp of frame 9 (the last frame) should be 9.0, that is, + (total_frames_num – 1) / original_fps. + + Args: + data: Raw video bytes + num_frames: Not used in dynamic backend + fps: Target FPS for sampling (default: 1) + + Returns: + Tuple of (frames_array, metadata_dict) + """ + import cv2 + + backend = cls().get_cv2_video_api() + cap = cv2.VideoCapture(BytesIO(data), backend, []) + if not cap.isOpened(): + raise ValueError("Could not open video stream") + + total_frames_num = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + original_fps = float(cap.get(cv2.CAP_PROP_FPS)) + # The timestamp of the rightmost frame, cannot be used to calculate frame 0. + if total_frames_num >= 1 and original_fps > 0: + total_duration = (total_frames_num - 1) / original_fps + else: + total_duration = 0 + + # `fps` is the FPS parameter passed in for sampling, + # -1 indicates that sampling can be performed directly without FPS limitation. + if fps > 0: + # Num_frames is the maximum number of frames to sample. + # If fewer frames are sampled at this sample_fps, the update duration will be longer. # noqa: E501 + if num_frames >= int(total_duration * fps) + 1: + num_frames = int(total_duration * fps) + 1 + # Under the new maximum frame rate, the video duration of the rightmost frame, # noqa: E501 + # cannot be calculated for frame 0. + total_duration = min(total_duration, (num_frames - 1) / fps) + elif fps != -1: + raise ValueError( + f"requires dataset fps is -1 or greater than 0 but got {fps}" + ) + + sample_frame_timestamps = np.linspace( + 0, total_duration, num_frames, dtype=float + ) + frames_indices = [ + min(total_frames_num - 1, round(t * original_fps)) + for t in sample_frame_timestamps + ] + + frames, valid_frame_indices, recovered_map = cls._read_frames_with_recovery( + cap, frames_indices, total_frames_num + ) + + if recovered_map: + logger.info( + "Frame recovery: %d frames recovered using forward scan.", + len(recovered_map), + ) + + metadata = { + "total_num_frames": total_frames_num, + "fps": original_fps, + "duration": total_duration, + "video_backend": "opencv_dynamic_openpangu", + "frames_indices": valid_frame_indices, + "do_sample_frames": False, + } + return frames, metadata From 71dfce6aa6cc14d016154b4e3fd8cc40c05415f9 Mon Sep 17 00:00:00 2001 From: Hanjie Qiu <50634613+hjjq@users.noreply.github.com> Date: Wed, 25 Feb 2026 19:17:20 -0800 Subject: [PATCH 19/43] [Kernel] Refactor FlashInfer allreduce for mnnvl backend (#34109) Signed-off-by: hjjq <50634613+hjjq@users.noreply.github.com> Signed-off-by: wzhao18 Co-authored-by: wzhao18 Co-authored-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> --- .../kernels/benchmark_device_communicators.py | 113 ++++++-- .../kernels/benchmark_fused_collective.py | 210 ++++++++------- .../distributed/test_fusion_all_reduce.py | 10 +- .../passes/fusion/allreduce_rms_fusion.py | 146 ++++++---- .../device_communicators/cuda_communicator.py | 28 +- .../flashinfer_all_reduce.py | 252 ++++++++++++++++++ vllm/envs.py | 14 + 7 files changed, 593 insertions(+), 180 deletions(-) create mode 100644 vllm/distributed/device_communicators/flashinfer_all_reduce.py diff --git a/benchmarks/kernels/benchmark_device_communicators.py b/benchmarks/kernels/benchmark_device_communicators.py index 7b453fe7b68..d1005461ab9 100644 --- a/benchmarks/kernels/benchmark_device_communicators.py +++ b/benchmarks/kernels/benchmark_device_communicators.py @@ -30,6 +30,9 @@ import torch.distributed as dist from torch.distributed import ProcessGroup from vllm.distributed.device_communicators.custom_all_reduce import CustomAllreduce +from vllm.distributed.device_communicators.flashinfer_all_reduce import ( + FlashInferAllReduce, +) from vllm.distributed.device_communicators.pynccl import ( PyNcclCommunicator, register_nccl_symmetric_ops, @@ -44,7 +47,7 @@ from vllm.utils.argparse_utils import FlexibleArgumentParser logger = init_logger(__name__) # Default sequence lengths to benchmark -DEFAULT_SEQUENCE_LENGTHS = [128, 512, 1024, 2048, 4096, 8192] +DEFAULT_SEQUENCE_LENGTHS = [16, 64, 128, 512, 1024, 2048, 4096, 8192] # Fixed hidden size and dtype for all benchmarks HIDDEN_SIZE = 8192 @@ -81,6 +84,7 @@ class CommunicatorBenchmark: self.symm_mem_comm = None self.symm_mem_comm_multimem = None self.symm_mem_comm_two_shot = None + self.fi_ar_comm = None self._init_communicators() @@ -161,6 +165,22 @@ class CommunicatorBenchmark: ) self.symm_mem_comm_two_shot = None + try: + self.fi_ar_comm = FlashInferAllReduce( + group=self.cpu_group, + device=self.device, + ) + if not self.fi_ar_comm.disabled: + logger.info("Rank %s: FlashInferAllReduce initialized", self.rank) + else: + logger.info("Rank %s: FlashInferAllReduce disabled", self.rank) + self.fi_ar_comm = None + except Exception as e: + logger.warning( + "Rank %s: Failed to initialize FlashInferAllReduce: %s", self.rank, e + ) + self.fi_ar_comm = None + def benchmark_allreduce( self, sequence_length: int, num_warmup: int, num_trials: int ) -> dict[str, float]: @@ -180,7 +200,8 @@ class CommunicatorBenchmark: lambda t, c=comm: c.custom_all_reduce(t), lambda t, c=comm: c.should_custom_ar(t), comm.capture(), - "1stage", # env variable value + {"VLLM_CUSTOM_ALLREDUCE_ALGO": "1stage"}, + None, # no destroy function ) ) # CustomAllreduce two-shot @@ -190,7 +211,8 @@ class CommunicatorBenchmark: lambda t, c=comm: c.custom_all_reduce(t), lambda t, c=comm: c.should_custom_ar(t), comm.capture(), - "2stage", # env variable value + {"VLLM_CUSTOM_ALLREDUCE_ALGO": "2stage"}, + None, # no destroy function ) ) @@ -202,7 +224,8 @@ class CommunicatorBenchmark: lambda t, c=comm: c.all_reduce(t), lambda t: True, # Always available if initialized nullcontext(), - None, # no env variable needed + {}, # no env variable needed + None, # no destroy function ) ) communicators.append( @@ -211,7 +234,8 @@ class CommunicatorBenchmark: lambda t: torch.ops.vllm.all_reduce_symmetric_with_copy(t), lambda t: True, # Always available if initialized nullcontext(), - None, # no env variable needed + {}, # no env variable needed + None, # no destroy function ) ) @@ -223,7 +247,8 @@ class CommunicatorBenchmark: lambda t, c=comm: c.all_reduce(t), lambda t, c=comm: c.should_use_symm_mem(t), nullcontext(), - None, # no env variable needed + {}, # no env variable needed + None, # no destroy function ) ) @@ -235,29 +260,67 @@ class CommunicatorBenchmark: lambda t, c=comm: c.all_reduce(t), lambda t, c=comm: c.should_use_symm_mem(t), nullcontext(), - None, # no env variable needed + {}, # no env variable needed + None, # no destroy function needed + ) + ) + + if self.fi_ar_comm is not None: + comm = self.fi_ar_comm + communicators.append( + ( + "flashinfer_trtllm", + lambda t, c=comm: c.all_reduce(t), + lambda t, c=comm: c.should_use_fi_ar(t), + nullcontext(), + {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "trtllm"}, + lambda c=comm: c.destroy(), + ) + ) + communicators.append( + ( + "flashinfer_mnnvl", + lambda t, c=comm: c.all_reduce(t), + lambda t, c=comm: c.should_use_fi_ar(t), + nullcontext(), + {"VLLM_FLASHINFER_ALLREDUCE_BACKEND": "mnnvl"}, + lambda c=comm: c.destroy(), ) ) # Benchmark each communicator - for name, allreduce_fn, should_use_fn, context, env_var in communicators: - # Set environment variable if needed - if env_var is not None: - os.environ["VLLM_CUSTOM_ALLREDUCE_ALGO"] = env_var - else: - # Clear the environment variable to avoid interference - os.environ.pop("VLLM_CUSTOM_ALLREDUCE_ALGO", None) - - latency = self.benchmark_allreduce_single( - sequence_length, - allreduce_fn, - should_use_fn, - context, - num_warmup, - num_trials, - ) - if latency is not None: - results[name] = latency + for ( + name, + allreduce_fn, + should_use_fn, + context, + env_dict, + destroy_fn, + ) in communicators: + # Save original values and apply new environment variables + saved_env = {key: os.environ.get(key) for key in env_dict} + for key, value in env_dict.items(): + os.environ[key] = value + try: + latency = self.benchmark_allreduce_single( + sequence_length, + allreduce_fn, + should_use_fn, + context, + num_warmup, + num_trials, + ) + if latency is not None: + results[name] = latency + finally: + if destroy_fn is not None: + destroy_fn() + # Restore environment variables to their original state + for key, original_value in saved_env.items(): + if original_value is None: + os.environ.pop(key, None) + else: + os.environ[key] = original_value return results diff --git a/benchmarks/kernels/benchmark_fused_collective.py b/benchmarks/kernels/benchmark_fused_collective.py index 633529edf16..e18f6a7580f 100644 --- a/benchmarks/kernels/benchmark_fused_collective.py +++ b/benchmarks/kernels/benchmark_fused_collective.py @@ -5,8 +5,11 @@ Benchmark for FlashInfer fused collective operations vs standard operations. This benchmark compares: -1. FlashInfer's allreduce_fusion (fused allreduce + rmsnorm + optional quant) -2. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations +1. FlashInfer's allreduce_fusion with trtllm backend + (fused allreduce + rmsnorm + optional FP8/FP4 quant) +2. FlashInfer's allreduce_fusion with mnnvl backend + (fused allreduce + rmsnorm only, no quantization support) +3. Standard tensor_model_parallel_all_reduce + separate rmsnorm/quant operations Usage with torchrun: torchrun --nproc_per_node=2 benchmark_fused_collective.py @@ -48,8 +51,12 @@ SCALED_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant logger = init_logger(__name__) # Try to import FlashInfer +TorchDistBackend = None try: import flashinfer.comm as flashinfer_comm # type: ignore + from flashinfer.comm.mnnvl import ( # type: ignore + TorchDistBackend, + ) if not ( hasattr(flashinfer_comm, "allreduce_fusion") @@ -74,11 +81,15 @@ _FI_MAX_SIZES = { 8: 64 * MiB, # 64MB } -# Global workspace tensor for FlashInfer -_FI_WORKSPACE = None +# Global workspace tensors for FlashInfer (keyed by backend name) +_FI_WORKSPACES: dict = {} + +# Backends to benchmark +FLASHINFER_BACKENDS = ["trtllm", "mnnvl"] def setup_flashinfer_workspace( + backend: str, world_size: int, rank: int, hidden_dim: int, @@ -86,41 +97,54 @@ def setup_flashinfer_workspace( dtype: torch.dtype, ): """Setup FlashInfer workspace for fused allreduce operations.""" - global _FI_WORKSPACE + global FI_WORKSPACES if flashinfer_comm is None: - return None, None + return None if world_size not in _FI_MAX_SIZES: logger.warning("FlashInfer not supported for world size %s", world_size) - return None, None + return None try: + kwargs = {} + if TorchDistBackend is not None: + kwargs["comm_backend"] = TorchDistBackend(group=dist.group.WORLD) + workspace = flashinfer_comm.create_allreduce_fusion_workspace( - backend="trtllm", + backend=backend, world_size=world_size, rank=rank, max_token_num=max_token_num, hidden_dim=hidden_dim, dtype=dtype, + **kwargs, ) - _FI_WORKSPACE = workspace + _FI_WORKSPACES[backend] = workspace return workspace except Exception as e: - logger.error("Failed to setup FlashInfer workspace: %s", e) + logger.error( + "Failed to setup FlashInfer workspace (backend=%s): %s", backend, e + ) return None -def cleanup_flashinfer_workspace(workspace): - """Cleanup FlashInfer workspace.""" - if flashinfer_comm is None or workspace is None: +def cleanup_flashinfer_workspaces(): + """Cleanup all FlashInfer workspaces.""" + if flashinfer_comm is None: return - try: - workspace.destroy() - except Exception as e: - logger.error("Failed to cleanup FlashInfer workspace: %s", e) + for backend, workspace in _FI_WORKSPACES.items(): + try: + workspace.destroy() + except Exception as e: + logger.error( + "Failed to cleanup FlashInfer workspace (backend=%s): %s", + backend, + e, + ) + _FI_WORKSPACES.clear() class FlashInferFusedAllReduceParams: @@ -134,7 +158,7 @@ class FlashInferFusedAllReduceParams: self.fp32_acc = True self.max_token_num = max_token_num - def get_trtllm_fused_allreduce_kwargs(self): + def get_flashinfer_fused_allreduce_kwargs(self): return { "launch_with_pdl": self.launch_with_pdl, "fp32_acc": self.fp32_acc, @@ -147,11 +171,12 @@ def flashinfer_fused_allreduce_rmsnorm( rms_gamma: torch.Tensor, rms_eps: float, allreduce_params: "FlashInferFusedAllReduceParams", + workspace: object, use_oneshot: bool, norm_out: torch.Tensor | None = None, ): """FlashInfer fused allreduce + rmsnorm operation.""" - if flashinfer_comm is None or _FI_WORKSPACE is None: + if flashinfer_comm is None or workspace is None: raise RuntimeError("FlashInfer not available or workspace not initialized") if norm_out is None: @@ -160,9 +185,13 @@ def flashinfer_fused_allreduce_rmsnorm( else: residual_out = input_tensor + layout_code = None + if workspace.backend == "trtllm": + layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4 + flashinfer_comm.allreduce_fusion( input=input_tensor, - workspace=_FI_WORKSPACE, + workspace=workspace, pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNorm, residual_in=residual, residual_out=residual_out, @@ -171,10 +200,10 @@ def flashinfer_fused_allreduce_rmsnorm( rms_eps=rms_eps, quant_out=None, scale_out=None, - layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4, + layout_code=layout_code, scale_factor=None, use_oneshot=use_oneshot, - **allreduce_params.get_trtllm_fused_allreduce_kwargs(), + **allreduce_params.get_flashinfer_fused_allreduce_kwargs(), ) @@ -185,12 +214,16 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant( rms_eps: float, scale_factor: torch.Tensor, allreduce_params: FlashInferFusedAllReduceParams, + workspace: object, use_oneshot: bool = True, norm_out: torch.Tensor | None = None, quant_out: torch.Tensor | None = None, ): - """FlashInfer fused allreduce + rmsnorm + FP8 quantization.""" - if flashinfer_comm is None or _FI_WORKSPACE is None: + """FlashInfer fused allreduce + rmsnorm + FP8 quantization. + + Note: Only supported by the trtllm backend. + """ + if flashinfer_comm is None or workspace is None: raise RuntimeError("FlashInfer not available or workspace not initialized") if norm_out is None: @@ -201,7 +234,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant( flashinfer_comm.allreduce_fusion( input=input_tensor, - workspace=_FI_WORKSPACE, + workspace=workspace, pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP8Quant, residual_in=residual, residual_out=residual_out, @@ -213,7 +246,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp8_quant( layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4, scale_factor=scale_factor, use_oneshot=use_oneshot, - **allreduce_params.get_trtllm_fused_allreduce_kwargs(), + **allreduce_params.get_flashinfer_fused_allreduce_kwargs(), ) @@ -224,13 +257,17 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant( rms_eps: float, input_global_scale: torch.Tensor, allreduce_params: FlashInferFusedAllReduceParams, + workspace: object, quant_out: torch.Tensor, use_oneshot: bool, output_scale: torch.Tensor, norm_out: torch.Tensor | None = None, ): - """FlashInfer fused allreduce + rmsnorm + FP4 quantization.""" - if flashinfer_comm is None or _FI_WORKSPACE is None: + """FlashInfer fused allreduce + rmsnorm + FP4 quantization. + + Note: Only supported by the trtllm backend. + """ + if flashinfer_comm is None or workspace is None: raise RuntimeError("FlashInfer not available or workspace not initialized") if norm_out is None: @@ -241,7 +278,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant( flashinfer_comm.allreduce_fusion( input=input_tensor, - workspace=_FI_WORKSPACE, + workspace=workspace, pattern=flashinfer_comm.AllReduceFusionPattern.kARResidualRMSNormFP4Quant, residual_in=residual, residual_out=residual_out, @@ -253,7 +290,7 @@ def flashinfer_fused_allreduce_rmsnorm_fp4_quant( layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4, scale_factor=input_global_scale, use_oneshot=use_oneshot, - **allreduce_params.get_trtllm_fused_allreduce_kwargs(), + **allreduce_params.get_flashinfer_fused_allreduce_kwargs(), ) @@ -386,13 +423,16 @@ def run_benchmarks( dtype: torch.dtype, use_residual: bool, allreduce_params: FlashInferFusedAllReduceParams | None, + workspaces: dict, quant_modes: set[str], no_oneshot: bool, ): """Run all benchmarks for given configuration. Args: - quant_mode: "none", "fp8_only", "fp4_only", or "all" + allreduce_params: Shared parameters for FlashInfer fused allreduce. + workspaces: Dict mapping backend name ("trtllm", "mnnvl") to workspace. + quant_modes: Set of quantization modes: "none", "fp8", "fp4". """ ( input_tensor, @@ -454,10 +494,11 @@ def run_benchmarks( logger.error("Standard AllReduce+RMSNorm Native Compiled failed: %s", e) results["standard_allreduce_rmsnorm_native_compiled"] = float("inf") - # FlashInfer Fused AllReduce + RMSNorm Oneshot/Twoshot - if flashinfer_comm is not None and allreduce_params is not None: + # FlashInfer Fused AllReduce + RMSNorm (all backends) + for backend, workspace in workspaces.items(): for use_oneshot in use_oneshot_options: suffix = "_oneshot" if use_oneshot else "_twoshot" + key = f"flashinfer_{backend}_fused_allreduce_rmsnorm{suffix}" try: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm, @@ -467,14 +508,17 @@ def run_benchmarks( rms_gamma=rms_gamma, rms_eps=rms_eps, allreduce_params=allreduce_params, + workspace=workspace, use_oneshot=use_oneshot, ) - results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = time_ms + results[key] = time_ms except Exception as e: - logger.error("FlashInfer Fused AllReduce+RMSNorm failed: %s", e) - results[f"flashinfer_fused_allreduce_rmsnorm{suffix}"] = float( - "inf" + logger.error( + "FlashInfer (%s) Fused AllReduce+RMSNorm failed: %s", + backend, + e, ) + results[key] = float("inf") if "fp8" in quant_modes: # Standard AllReduce + RMSNorm + FP8 Quant @@ -540,10 +584,12 @@ def run_benchmarks( "inf" ) - # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant Oneshot - if flashinfer_comm is not None and allreduce_params is not None: + # FlashInfer Fused AllReduce + RMSNorm + FP8 Quant (trtllm only) + if "trtllm" in workspaces: + trtllm_ws = workspaces["trtllm"] for use_oneshot in use_oneshot_options: suffix = "_oneshot" if use_oneshot else "_twoshot" + key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp8_quant{suffix}" try: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm_fp8_quant, @@ -555,19 +601,16 @@ def run_benchmarks( scale_factor=scale_fp8, quant_out=quant_out_fp8, allreduce_params=allreduce_params, + workspace=trtllm_ws, use_oneshot=use_oneshot, ) - results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = ( - time_ms - ) + results[key] = time_ms except Exception as e: logger.error( - "FlashInfer Fused AllReduce+RMSNorm+FP8 Oneshot failed: %s", + "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP8 failed: %s", e, ) - results[f"flashinfer_fused_allreduce_rmsnorm_fp8_quant{suffix}"] = ( - float("inf") - ) + results[key] = float("inf") if "fp4" in quant_modes and current_platform.has_device_capability(100): # Standard AllReduce + RMSNorm + FP4 Quant @@ -627,10 +670,12 @@ def run_benchmarks( "inf" ) - # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Oneshot - if flashinfer_comm is not None and allreduce_params is not None: + # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant (trtllm only) + if "trtllm" in workspaces: + trtllm_ws = workspaces["trtllm"] for use_oneshot in use_oneshot_options: suffix = "_oneshot" if use_oneshot else "_twoshot" + key = f"flashinfer_trtllm_fused_allreduce_rmsnorm_fp4_quant{suffix}" try: time_ms = benchmark_operation( flashinfer_fused_allreduce_rmsnorm_fp4_quant, @@ -641,49 +686,18 @@ def run_benchmarks( rms_eps=rms_eps, input_global_scale=scale_fp4, allreduce_params=allreduce_params, + workspace=trtllm_ws, quant_out=fp4_quant_out, output_scale=fp4_output_scale, use_oneshot=use_oneshot, ) - results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = ( - time_ms - ) + results[key] = time_ms except Exception as e: logger.error( - "FlashInfer Fused AllReduce+RMSNorm+FP4 Oneshot failed: %s", + "FlashInfer (trtllm) Fused AllReduce+RMSNorm+FP4 failed: %s", e, ) - results[f"flashinfer_fused_allreduce_rmsnorm_fp4_quant{suffix}"] = ( - float("inf") - ) - - # FlashInfer Fused AllReduce + RMSNorm + FP4 Quant Two-shot - if flashinfer_comm is not None and allreduce_params is not None: - try: - time_ms = benchmark_operation( - flashinfer_fused_allreduce_rmsnorm_fp4_quant, - input_tensor, - residual=residual, - norm_out=norm_out, - rms_gamma=rms_gamma, - rms_eps=rms_eps, - input_global_scale=scale_fp4, - allreduce_params=allreduce_params, - quant_out=fp4_quant_out, - output_scale=fp4_output_scale, - use_oneshot=False, - ) - results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = ( - time_ms - ) - except Exception as e: - logger.error( - "FlashInfer Fused AllReduce+RMSNorm+FP4 Two-shot failed: %s", - e, - ) - results["flashinfer_fused_allreduce_rmsnorm_fp4_quant_twoshot"] = float( - "inf" - ) + results[key] = float("inf") return results @@ -1021,8 +1035,7 @@ def main(): configs = list(itertools.product(args.num_tokens, dtypes, residual_options)) - # Setup FlashInfer workspace if available - workspace = None + # Setup FlashInfer workspaces for all backends allreduce_params = None if flashinfer_comm is not None: @@ -1037,15 +1050,17 @@ def main(): args.hidden_dim * max_element_size ) - workspace = setup_flashinfer_workspace( - world_size, - rank, - args.hidden_dim, - max_num_token, - dtype=workspace_dtype, - ) + for backend in FLASHINFER_BACKENDS: + setup_flashinfer_workspace( + backend=backend, + world_size=world_size, + rank=rank, + hidden_dim=args.hidden_dim, + max_token_num=max_num_token, + dtype=workspace_dtype, + ) - if workspace is not None: + if _FI_WORKSPACES: allreduce_params = FlashInferFusedAllReduceParams( max_token_num=max_num_token, ) @@ -1071,6 +1086,7 @@ def main(): dtype, use_residual, allreduce_params, + workspaces=_FI_WORKSPACES, quant_modes=quant_modes, no_oneshot=args.no_oneshot, ) @@ -1109,11 +1125,13 @@ def main(): finally: # Cleanup - if workspace is not None: - cleanup_flashinfer_workspace(workspace) + cleanup_flashinfer_workspaces() dist.barrier() if __name__ == "__main__": - main() + from vllm.config import VllmConfig, set_current_vllm_config + + with set_current_vllm_config(VllmConfig()): + main() diff --git a/tests/compile/passes/distributed/test_fusion_all_reduce.py b/tests/compile/passes/distributed/test_fusion_all_reduce.py index d48f2297031..6d5113b1e84 100644 --- a/tests/compile/passes/distributed/test_fusion_all_reduce.py +++ b/tests/compile/passes/distributed/test_fusion_all_reduce.py @@ -142,7 +142,6 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module): *(scaled_fp4_quant(w, wg) for w, wg in zip(self.w, wgscale)) ) self.wq, self.wscale = list(wq_gen), list(wscale_gen) - print(f"{self.wq=}, {self.wscale=}") def forward(self, hidden_states): # avoid having graph input be an arg to a pattern directly @@ -199,6 +198,7 @@ class TestAllReduceFusedAddRMSNormStaticQuantFP4Model(torch.nn.Module): @pytest.mark.parametrize("hidden_size", [64]) @pytest.mark.parametrize("dtype", [torch.bfloat16]) @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False]) +@pytest.mark.parametrize("flashinfer_allreduce_backend", ["trtllm", "mnnvl"]) @pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda"], reason="Only test on CUDA") @pytest.mark.skipif( not find_spec("flashinfer") @@ -215,6 +215,7 @@ def test_all_reduce_fusion_pass_replace( dtype: torch.dtype, enable_rms_norm_custom_op, enable_quant_fp8_custom_op, + flashinfer_allreduce_backend, ): num_processes = 2 if ( @@ -238,6 +239,7 @@ def test_all_reduce_fusion_pass_replace( dtype, enable_rms_norm_custom_op, enable_quant_fp8_custom_op, + flashinfer_allreduce_backend, ), nprocs=nprocs, ) @@ -255,6 +257,7 @@ def all_reduce_fusion_pass_on_test_model( dtype: torch.dtype, enable_rms_norm_custom_op, enable_quant_fp8_custom_op, + flashinfer_allreduce_backend, ): set_random_seed(0) @@ -270,6 +273,7 @@ def all_reduce_fusion_pass_on_test_model( "WORLD_SIZE": str(world_size), "MASTER_ADDR": "localhost", "MASTER_PORT": "12345", + "VLLM_FLASHINFER_ALLREDUCE_BACKEND": flashinfer_allreduce_backend, } ) @@ -317,6 +321,10 @@ def all_reduce_fusion_pass_on_test_model( compiled_model = torch.compile(model, backend=backend) compiled_model(hidden_states) + results_unfused = model(hidden_states) + results_fused = compiled_model(hidden_states) + torch.testing.assert_close(results_unfused, results_fused, atol=1e-2, rtol=1e-2) + assert all_reduce_fusion_pass.matched_count == 4, ( f"{all_reduce_fusion_pass.matched_count=}" ) diff --git a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py index b6a1314af9e..44dc3d67bb9 100644 --- a/vllm/compilation/passes/fusion/allreduce_rms_fusion.py +++ b/vllm/compilation/passes/fusion/allreduce_rms_fusion.py @@ -22,7 +22,9 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) from vllm.platforms import current_platform -from vllm.utils.torch_utils import direct_register_custom_op +from vllm.utils.torch_utils import ( + direct_register_custom_op, +) from ..inductor_pass import enable_fake_mode from ..vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass @@ -44,8 +46,6 @@ if find_spec("flashinfer"): except ImportError: pass -logger = init_logger(__name__) - if hasattr(torch.ops._C, "scaled_fp4_quant"): STATIC_FP4_QUANT_OP = torch.ops._C.scaled_fp4_quant.default @@ -82,7 +82,16 @@ _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB: dict[int, dict[int, float]] = { if flashinfer_comm is not None: - _FI_WORKSPACE = None + from vllm.distributed.device_communicators.flashinfer_all_reduce import ( + destroy_fi_ar_workspace, + get_fi_ar_quant_workspace, + get_fi_ar_workspace, + initialize_fi_ar_quant_workspace, + initialize_fi_ar_workspace, + ) + + ar_fusion_patterns = flashinfer_comm.AllReduceFusionPattern + MiB = 1024 * 1024 def call_trtllm_fused_allreduce_norm( @@ -122,9 +131,19 @@ if flashinfer_comm is not None: max_one_shot_size is None or current_tensor_size <= max_one_shot_size * MiB ) - assert _FI_WORKSPACE is not None, ( - "Flashinfer must be enabled when using flashinfer" + # Select workspace based on pattern: quant patterns use the + # trtllm quant workspace, non-quant patterns use the primary workspace. + if pattern_code in ( + ar_fusion_patterns.kARResidualRMSNormFP8Quant, + ar_fusion_patterns.kARResidualRMSNormFP4Quant, + ): + workspace = get_fi_ar_quant_workspace() + else: + workspace = get_fi_ar_workspace() + assert workspace is not None, ( + "Flashinfer workspace must be initialized when using flashinfer" ) + assert flashinfer_comm is not None if norm_out is None: norm_out = allreduce_in residual_out = residual @@ -133,25 +152,30 @@ if flashinfer_comm is not None: # as flashinfer does not support rms_norm # and allreduce_out together residual_out = allreduce_in - # For the sizes that are smaller than the max size, - # we only use flashinfer one shot allreduce + + layout_code = None + # layout_code only supported by trtllm backend + if workspace.backend == "trtllm": + # in vllm we only support swizzled layout + layout_code = flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4 + flashinfer_comm.allreduce_fusion( input=allreduce_in, - workspace=_FI_WORKSPACE, + workspace=workspace, pattern=pattern_code, - residual_in=residual, + launch_with_pdl=launch_with_pdl, + output=None, residual_out=residual_out, norm_out=norm_out, - rms_gamma=rms_gamma, - rms_eps=rms_eps, - launch_with_pdl=launch_with_pdl, - use_oneshot=use_oneshot, - fp32_acc=fp32_acc, quant_out=quant_out, scale_out=scale_out, - # in vllm we only support swizzled layout - layout_code=flashinfer_comm.QuantizationSFLayout.SWIZZLED_128x4, + residual_in=residual, + rms_gamma=rms_gamma, + rms_eps=rms_eps, scale_factor=scale_factor, + layout_code=layout_code, + use_oneshot=use_oneshot, + fp32_acc=fp32_acc, ) def call_trtllm_fused_allreduce_norm_fake( @@ -729,29 +753,36 @@ class AllReduceFusionPass(VllmPatternMatcherPass): scope="global", ) - try: - self.workspace = flashinfer_comm.create_allreduce_fusion_workspace( - backend="trtllm", - world_size=self.tp_size, - rank=rank, - max_token_num=self.max_token_num, - hidden_dim=self.hidden_dim, - dtype=self.model_dtype, - ) - except RuntimeError as e: - if "multicast" not in str(e).lower(): - raise - logger.warning_once( - "AllReduce fusion pass is disabled: flashinfer workspace " - "creation failed: %s. This is expected on GPUs without " - "NVSwitch (e.g., NVLink bridge-only or PCIe topologies). " - "Falling back to non-fused allreduce.", - str(e), - ) - return + for workspace_init_fn in [ + initialize_fi_ar_workspace, + initialize_fi_ar_quant_workspace, + ]: + try: + workspace_init_fn( + world_size=self.tp_size, + rank=rank, + max_token_num=self.max_token_num, + hidden_dim=self.hidden_dim, + dtype=self.model_dtype, + group=self.group, + ) + except Exception as e: + if "multicast" in str(e).lower(): + logger.warning( + "AllReduce fusion pass is disabled: flashinfer workspace " + "creation failed: %s. This is expected on GPUs without " + "NVSwitch (e.g., NVLink bridge-only or PCIe topologies). " + "Falling back to non-fused allreduce.", + str(e), + ) + else: + logger.warning( + "Failed to initialize FlashInfer All Reduce workspace: %s. " + "AllReduce fusion pass will be disabled.", + e, + ) + return - global _FI_WORKSPACE - _FI_WORKSPACE = self.workspace self.allreduce_params = FlashInferFusedAllReduceParams( world_size=self.tp_size, max_token_num=self.max_token_num, @@ -762,32 +793,34 @@ class AllReduceFusionPass(VllmPatternMatcherPass): @enable_fake_mode def register_patterns(self) -> None: + supports_quantization = get_fi_ar_quant_workspace() is not None for epsilon in [1e-5, 1e-6]: - AllReduceFusedRMSNormStaticQuantFP8Pattern( - epsilon, - self.model_dtype, - self.device, - self.allreduce_params, - ).register(self.patterns) - AllReduceFusedAddRMSNormStaticQuantFP8Pattern( - epsilon, - self.model_dtype, - self.device, - self.allreduce_params, - ).register(self.patterns) - if current_platform.has_device_capability(100): - AllReduceFusedRMSNormStaticQuantNVFP4Pattern( + if supports_quantization: + AllReduceFusedRMSNormStaticQuantFP8Pattern( epsilon, self.model_dtype, self.device, self.allreduce_params, ).register(self.patterns) - AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern( + AllReduceFusedAddRMSNormStaticQuantFP8Pattern( epsilon, self.model_dtype, self.device, self.allreduce_params, ).register(self.patterns) + if current_platform.has_device_capability(100): + AllReduceFusedRMSNormStaticQuantNVFP4Pattern( + epsilon, + self.model_dtype, + self.device, + self.allreduce_params, + ).register(self.patterns) + AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern( + epsilon, + self.model_dtype, + self.device, + self.allreduce_params, + ).register(self.patterns) AllReduceRMSNormPattern( epsilon, self.model_dtype, @@ -825,6 +858,5 @@ class AllReduceFusionPass(VllmPatternMatcherPass): def __del__(self) -> None: if getattr(self, "disabled", True): return - if getattr(self, "workspace", None) is not None: - with contextlib.suppress(Exception): - self.workspace.destroy() + with contextlib.suppress(Exception): + destroy_fi_ar_workspace() diff --git a/vllm/distributed/device_communicators/cuda_communicator.py b/vllm/distributed/device_communicators/cuda_communicator.py index 4c78871e1fd..62e2b90377f 100644 --- a/vllm/distributed/device_communicators/cuda_communicator.py +++ b/vllm/distributed/device_communicators/cuda_communicator.py @@ -34,19 +34,25 @@ class CudaCommunicator(DeviceCommunicatorBase): # custom allreduce or torch symm mem can be used only by tp use_custom_allreduce = False use_torch_symm_mem = False + use_flashinfer_allreduce = False else: from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM + use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER self.use_custom_allreduce = use_custom_allreduce self.use_torch_symm_mem = use_torch_symm_mem + self.use_flashinfer_allreduce = use_flashinfer_allreduce # lazy import to avoid documentation build error from vllm.distributed.device_communicators.custom_all_reduce import ( CustomAllreduce, ) + from vllm.distributed.device_communicators.flashinfer_all_reduce import ( + FlashInferAllReduce, + ) from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator from vllm.distributed.device_communicators.quick_all_reduce import ( QuickAllReduce, @@ -65,12 +71,20 @@ class CudaCommunicator(DeviceCommunicatorBase): self.ca_comm: CustomAllreduce | None = None self.qr_comm: QuickAllReduce | None = None self.symm_mem_comm: SymmMemCommunicator | None = None + self.fi_ar_comm: FlashInferAllReduce | None = None + if use_torch_symm_mem and current_platform.is_cuda(): self.symm_mem_comm = SymmMemCommunicator( group=self.cpu_group, device=self.device, ) + if self.use_flashinfer_allreduce and self.world_size > 1: + self.fi_ar_comm = FlashInferAllReduce( + group=self.cpu_group, + device=self.device, + ) + if use_custom_allreduce and self.world_size > 1: # Initialize a custom fast all-reduce implementation. self.ca_comm = CustomAllreduce( @@ -136,7 +150,7 @@ class CudaCommunicator(DeviceCommunicatorBase): out = torch.ops.vllm.all_reduce_symmetric_with_copy(input_) if out is not None: return out - # always try quick reduce first, then custom allreduce, + # always try quick reduce first, then flashinfer, then custom allreduce, # and then pynccl. (quick reduce just for ROCM MI3*) qr_comm = self.qr_comm if ( @@ -147,6 +161,15 @@ class CudaCommunicator(DeviceCommunicatorBase): out = qr_comm.quick_all_reduce(input_) assert out is not None return out + fi_ar_comm = self.fi_ar_comm + if ( + fi_ar_comm is not None + and not fi_ar_comm.disabled + and fi_ar_comm.should_use_fi_ar(input_) + ): + out = fi_ar_comm.all_reduce(input_) + assert out is not None + return out ca_comm = self.ca_comm if ( ca_comm is not None @@ -270,6 +293,9 @@ class CudaCommunicator(DeviceCommunicatorBase): self.pynccl_comm = None if self.ca_comm is not None: self.ca_comm = None + if self.fi_ar_comm is not None: + self.fi_ar_comm.destroy() + self.fi_ar_comm = None if self.all2all_manager is not None: self.all2all_manager.destroy() self.all2all_manager = None diff --git a/vllm/distributed/device_communicators/flashinfer_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_all_reduce.py new file mode 100644 index 00000000000..ea16c93763c --- /dev/null +++ b/vllm/distributed/device_communicators/flashinfer_all_reduce.py @@ -0,0 +1,252 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + + +import torch +import torch.distributed as dist +from torch.distributed import ProcessGroup + +import vllm.envs as envs +from vllm.config.compilation import PassConfig +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +fi_ar_available = False +try: + import flashinfer.comm as flashinfer_comm # type: ignore[no-redef] + from flashinfer.comm.mnnvl import ( + TorchDistBackend, # type: ignore[import-not-found, no-redef] + ) + + fi_ar_available = hasattr(flashinfer_comm, "allreduce_fusion") +except ImportError: + pass + +# Global workspace for standalone allreduce and non-quant ar+rms fusion +_fi_ar_workspace = None +# Extra workspace for quant fusion patterns (only supported by trtllm backend) +# Only created if primary workspace is not already trtllm +_fi_ar_quant_workspace = None + + +def get_fi_ar_workspace(): + return _fi_ar_workspace + + +def get_fi_ar_quant_workspace(): + return _fi_ar_quant_workspace + + +def initialize_fi_ar_workspace( + world_size: int, + rank: int, + max_token_num: int, + hidden_dim: int, + dtype: torch.dtype, + group: ProcessGroup, +) -> None: + """ + Initialize the workspace if not already initialized. + + Currently, this function is called by either the AllReduceFusionPass + or the FlashInferAllReduce backend for standalone allreduce. + If the fusion pass is enabled via + --compilation-config.pass_config.fuse_allreduce_rms=true, + it will create the workspace first, and the standalone backend + will reuse the workspace. Otherwise, the standalone backend will + create the workspace. + """ + global _fi_ar_workspace + if _fi_ar_workspace is not None: + return + + backend = envs.VLLM_FLASHINFER_ALLREDUCE_BACKEND + comm_backend = TorchDistBackend(group=group) + _fi_ar_workspace = flashinfer_comm.create_allreduce_fusion_workspace( + backend=backend, + world_size=world_size, + rank=rank, + max_token_num=max_token_num, + hidden_dim=hidden_dim, + dtype=dtype, + comm_backend=comm_backend, + ) + assert _fi_ar_workspace is not None + logger.debug( + "Initialized FlashInfer All Reduce workspace: backend=%s, " + "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s", + backend, + world_size, + rank, + max_token_num, + hidden_dim, + dtype, + ) + + +def initialize_fi_ar_quant_workspace( + world_size: int, + rank: int, + max_token_num: int, + hidden_dim: int, + dtype: torch.dtype, + group: ProcessGroup, +) -> None: + """ + Initialize the workspace used by quantization fusion patterns. + + Currently this always creates a workspace for trtllm backend as only it + supports quantization fusion (FP8/FP4). If the primary workspace + is already trtllm, the quant workspace aliases to it. + """ + global _fi_ar_quant_workspace + if _fi_ar_quant_workspace is not None: + return + + # If primary workspace is already trtllm, reuse it + if _fi_ar_workspace is not None and _fi_ar_workspace.backend == "trtllm": + _fi_ar_quant_workspace = _fi_ar_workspace + return + + comm_backend = TorchDistBackend(group=group) + _fi_ar_quant_workspace = flashinfer_comm.create_allreduce_fusion_workspace( + backend="trtllm", + world_size=world_size, + rank=rank, + max_token_num=max_token_num, + hidden_dim=hidden_dim, + dtype=dtype, + comm_backend=comm_backend, + ) + assert _fi_ar_quant_workspace is not None + logger.debug( + "Initialized FlashInfer All Reduce workspace: backend=trtllm, " + "world_size=%d, rank=%d, max_token_num=%d, hidden_dim=%d, dtype=%s", + world_size, + rank, + max_token_num, + hidden_dim, + dtype, + ) + + +def destroy_fi_ar_workspace(): + global _fi_ar_workspace + global _fi_ar_quant_workspace + if ( + _fi_ar_quant_workspace is not None + and _fi_ar_quant_workspace is not _fi_ar_workspace + ): + _fi_ar_quant_workspace.destroy() + _fi_ar_quant_workspace = None + if _fi_ar_workspace is not None: + _fi_ar_workspace.destroy() + _fi_ar_workspace = None + + +class FlashInferAllReduce: + def __init__( + self, + group: ProcessGroup, + device: int | str | torch.device, + ): + self.disabled = True + + if not fi_ar_available: + logger.info( + "FlashInfer All Reduce is disabled because flashinfer is not available" + ) + return + + if not current_platform.is_cuda(): + logger.info( + "FlashInfer All Reduce is disabled because it requires CUDA platform" + ) + return + + self.group = group + self.world_size = dist.get_world_size(self.group) + self.rank = dist.get_rank(self.group) + self.device = device + if self.world_size == 1: + return + + # Use the same threshold as the allreduce-rms fusion pass + # TODO: tune the threshold + MiB = 1024 * 1024 + max_workspace_size = PassConfig.default_fi_allreduce_fusion_max_size_mb().get( + self.world_size, None + ) + if not max_workspace_size: + logger.warning( + "FlashInfer All Reduce is disabled because it " + "is not supported for world_size=%d.", + self.world_size, + ) + return + self.max_workspace_size = max_workspace_size * MiB + self.max_num_tokens = 0 + self.disabled = False + + def _ensure_workspace(self, hidden_dim: int, dtype: torch.dtype) -> bool: + """Ensure the all reduce workspace is initialized.""" + if get_fi_ar_workspace() is not None: + return True + if self.max_num_tokens == 0: + element_size = torch.tensor([], dtype=dtype, device="cpu").element_size() + self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size) + try: + initialize_fi_ar_workspace( + world_size=self.world_size, + rank=self.rank, + max_token_num=self.max_num_tokens, + hidden_dim=hidden_dim, + dtype=dtype, + group=self.group, + ) + return True + except Exception as e: + logger.warning( + "Failed to initialize FlashInfer All Reduce workspace: %s. " + "FlashInfer All Reduce will be disabled.", + e, + ) + self.disabled = True + return False + + def should_use_fi_ar(self, input_tensor: torch.Tensor) -> bool: + if self.disabled: + return False + + if not input_tensor.is_cuda: + return False + + if not input_tensor.is_contiguous(): + return False + + if len(input_tensor.shape) != 2: + return False + + num_tokens, hidden_dim = input_tensor.shape + if not self.max_num_tokens: + element_size = torch.tensor([], dtype=input_tensor.dtype).element_size() + self.max_num_tokens = self.max_workspace_size // (hidden_dim * element_size) + + if num_tokens > self.max_num_tokens: + return False + + return self._ensure_workspace(hidden_dim, input_tensor.dtype) + + def all_reduce(self, input_tensor: torch.Tensor) -> torch.Tensor: + workspace = get_fi_ar_workspace() + return flashinfer_comm.allreduce_fusion( + input=input_tensor, + workspace=workspace, + pattern=flashinfer_comm.AllReduceFusionPattern.kAllReduce, + ) + + def destroy(self): + if not self.disabled: + destroy_fi_ar_workspace() diff --git a/vllm/envs.py b/vllm/envs.py index d62438d5735..d560cfc7753 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -168,6 +168,7 @@ if TYPE_CHECKING: VLLM_FLASHINFER_MOE_BACKEND: Literal["throughput", "latency", "masked_gemm"] = ( "latency" ) + VLLM_FLASHINFER_ALLREDUCE_BACKEND: Literal["auto", "trtllm", "mnnvl"] = "auto" VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE: int = 394 * 1024 * 1024 VLLM_XGRAMMAR_CACHE_MB: int = 0 VLLM_MSGPACK_ZERO_COPY_THRESHOLD: int = 256 @@ -206,6 +207,7 @@ if TYPE_CHECKING: VLLM_ROCM_FP8_MFMA_PAGE_ATTN: bool = False VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True + VLLM_ALLREDUCE_USE_FLASHINFER: bool = False VLLM_TUNED_CONFIG_FOLDER: str | None = None VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set() VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False @@ -1290,6 +1292,14 @@ environment_variables: dict[str, Callable[[], Any]] = { "latency", ["throughput", "latency", "masked_gemm"], ), + # Flashinfer fused allreduce backend. + # "auto" will default to "mnnvl", which performs mostly same/better than "trtllm". + # But "mnnvl" backend does not support fuse with quantization. + "VLLM_FLASHINFER_ALLREDUCE_BACKEND": env_with_choices( + "VLLM_FLASHINFER_ALLREDUCE_BACKEND", + "auto", + ["auto", "trtllm", "mnnvl"], + ), # Control the workspace buffer size for the FlashInfer backend. "VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE": lambda: int( os.getenv("VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE", str(394 * 1024 * 1024)) @@ -1448,6 +1458,10 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ALLREDUCE_USE_SYMM_MEM": lambda: bool( int(os.getenv("VLLM_ALLREDUCE_USE_SYMM_MEM", "1")) ), + # Whether to use FlashInfer allreduce + "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool( + int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0")) + ), # Experimental: use this to enable MCP tool calling for non harmony models "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool( int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0")) From 13025e71e888330aa3277948120051ebbc2674c7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 25 Feb 2026 20:42:40 -0800 Subject: [PATCH 20/43] [Model Runner V2] Add coding style guide (#35325) Signed-off-by: Woosuk Kwon --- vllm/v1/worker/gpu/model_runner.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index ccab6cec8c7..9e0cae6feef 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1,5 +1,22 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +NOTE: Coding style guide for this file: +This model runner is shared by all models: text and multimodal, generative +and embedding, public and private. As a result, this file must only contain +code that is common to every model. Model-specific behavior belongs in the +appropriate model-specific files. + +In other words: +* Be paranoid about changing this file. It should remain stable. +* Be even more paranoid about adding new lines. It should remain minimal. + +Even for shared features (for example, different parallelism modes), keep the +complexity out of this path. The less common the feature, the more it should be +hidden. Prefer utility functions defined elsewhere and call them from here, +instead of embedding feature-specific logic directly. +""" + import functools import gc import time From 4171ff6dd9ce18f452c4e9267f5bf090c0989b04 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Thu, 26 Feb 2026 05:00:10 +0000 Subject: [PATCH 21/43] [CPU][Feat] Enable KleidiAI INT8_W4A8 for all input dtypes (#34890) Signed-off-by: Fadi Arafeh Co-authored-by: Tyler Michael Smith --- .../linear/mixed_precision/dynamic_4bit.py | 32 +++++++++++++++++-- 1 file changed, 29 insertions(+), 3 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py b/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py index 3dfe06f1b13..d0515027628 100644 --- a/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py +++ b/vllm/model_executor/kernels/linear/mixed_precision/dynamic_4bit.py @@ -42,12 +42,13 @@ class Dynamic4bitLinearKernel(MPLinearKernel): not in [ torch.float32, torch.bfloat16, + torch.float16, ] ): return ( False, "Dynamic4bitLinearKernel on Arm requires Float32 or" - " BFloat16 activations", + " BFloat16 or Float16 activations", ) if c.full_weight_shape[0] % c.group_size != 0: return ( @@ -118,8 +119,30 @@ class Dynamic4bitLinearKernel(MPLinearKernel): x: torch.Tensor, bias: torch.Tensor | None = None, ) -> torch.Tensor: + # PyTorch / KleidiAI kernels natively support the following configs: + # - channelwise with bfloat16 / float32 activations + # - groupwise with float32 activations + # To support: + # - groupwise with bfloat16/float16 activations: we need to upcast + # activations to float32 before matmul and downcast back to bfloat16/float16 + # - channelwise with float16 activations, we need to upcast activations to + # float32 before matmul and downcast back to float16 + # Note: these activations will be dynamically quantized to int8 by the kernel. + c = self.config + is_groupwise = c.group_size != c.partition_weight_shape[0] + # dtype of activations before they get dynamically quantized to int8 + original_pre_quant_act_dtype = x.dtype + pre_quant_act_dtype = original_pre_quant_act_dtype + if ( + is_groupwise and pre_quant_act_dtype == torch.bfloat16 + ) or pre_quant_act_dtype == torch.float16: + pre_quant_act_dtype = torch.float32 + x_2d = x.reshape(-1, x.shape[-1]) + if pre_quant_act_dtype != original_pre_quant_act_dtype: + x_2d = x_2d.to(pre_quant_act_dtype) + out_shape = x.shape[:-1] + (c.partition_weight_shape[1],) w_q = getattr(layer, self.w_q_name) @@ -129,5 +152,8 @@ class Dynamic4bitLinearKernel(MPLinearKernel): c.group_size, c.partition_weight_shape[0], c.partition_weight_shape[1], - ) - return output.reshape(out_shape) + ).reshape(out_shape) + + if pre_quant_act_dtype != original_pre_quant_act_dtype: + output = output.to(original_pre_quant_act_dtype) + return output From 9d379410179b649f4e7651940debc35c4ac7c0a5 Mon Sep 17 00:00:00 2001 From: Jason Li Date: Wed, 25 Feb 2026 21:00:12 -0800 Subject: [PATCH 22/43] [torch.compile] Sequence Parallelism threshold compile ranges (#28672) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: jasonlizhengjian Signed-off-by: Jason Li Co-authored-by: Claude Opus 4.6 Co-authored-by: Luka Govedič --- tests/compile/conftest.py | 34 +++++ tests/compile/fusions_e2e/conftest.py | 89 ++++++++++-- .../compile/fusions_e2e/test_tp2_async_tp.py | 133 ++++++++++++++++++ tests/compile/test_config.py | 1 + .../test_sequence_parallelism_threshold.py | 110 +++++++++++++++ .../passes/fusion/sequence_parallelism.py | 123 +++++++++++++--- vllm/config/compilation.py | 9 +- vllm/config/vllm.py | 57 +++++++- 8 files changed, 524 insertions(+), 32 deletions(-) create mode 100644 tests/compile/conftest.py create mode 100644 tests/compile/test_sequence_parallelism_threshold.py diff --git a/tests/compile/conftest.py b/tests/compile/conftest.py new file mode 100644 index 00000000000..6aafac7bcad --- /dev/null +++ b/tests/compile/conftest.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from contextlib import contextmanager +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.platforms.interface import DeviceCapability + + +@pytest.fixture +def mock_cuda_platform(): + """ + Fixture that returns a factory for creating mocked CUDA platforms. + + Usage: + def test_something(mock_cuda_platform): + with mock_cuda_platform(is_cuda=True, capability=(9, 0)): + # test code + """ + + @contextmanager + def _mock_platform(is_cuda: bool = True, capability: tuple[int, int] | None = None): + mock_platform = MagicMock() + mock_platform.is_cuda.return_value = is_cuda + if capability is not None: + mock_platform.get_device_capability.return_value = DeviceCapability( + *capability + ) + with patch("vllm.platforms.current_platform", mock_platform): + yield mock_platform + + return _mock_platform diff --git a/tests/compile/fusions_e2e/conftest.py b/tests/compile/fusions_e2e/conftest.py index 1d9f6cda9fd..40b4de57f66 100644 --- a/tests/compile/fusions_e2e/conftest.py +++ b/tests/compile/fusions_e2e/conftest.py @@ -94,7 +94,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): run_model(full_compilation_config, model_name, **model_kwargs) num_compile_ranges = len(full_compilation_config.get_compile_ranges()) - assert num_compile_ranges in [1, 2] + assert num_compile_ranges in [1, 2, 3] print(f"Compile ranges: {full_compilation_config.get_compile_ranges()}") print("Fusion results:") @@ -107,12 +107,33 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): # Now check the matches for match_name in matches_check: - num_ranges_activated = ( - 1 if match_name == "ar_rms_fusion" else num_compile_ranges - ) - n_expected = tp_size * num_ranges_activated - log_matches = list(int(ms) for ms in log_matches_dict[match_name]) + + # AR+RMS skips the largest range; SP skips the smallest. + # When both are enabled, AR+RMS activation count is + # model-dependent (hidden_size affects threshold), so derive + # from log data. + if ( + match_name == "ar_rms_fusion" + and "sequence_parallel" in matches_check + and num_compile_ranges >= 2 + ): + assert ( + len(log_matches) >= tp_size and len(log_matches) % tp_size == 0 + ), ( + f"Expected multiple of {tp_size} ar_rms log entries, " + f"found {len(log_matches)}" + ) + num_ranges_activated = len(log_matches) // tp_size + elif ( + match_name in ("ar_rms_fusion", "sequence_parallel") + and num_compile_ranges >= 2 + ): + num_ranges_activated = num_compile_ranges - 1 + else: + num_ranges_activated = num_compile_ranges + + n_expected = tp_size * num_ranges_activated assert len(log_matches) == n_expected, ( f"Could not find {n_expected} {match_name} " f"(found {len(log_matches)}) in:\n {log_holder.text}" @@ -122,8 +143,8 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): if match_name == "rms_quant_fusion" and "ar_rms_fusion" in matches_check: # AR+rms+quant takes precedence over rms+quant if activated. - # That means we get full matching where ar+rms+quant was not activated, - # and less where it was + # That means we get full matching where ar+rms+quant was not + # activated, and less where it was (only the smallest range). assert sum(m == expected_matches for m in log_matches) == tp_size * ( num_ranges_activated - 1 ), "Expecting full rms+quant fusion where ar+rms+quant not activated" @@ -135,6 +156,43 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): f"Expecting at least {expected_matches - matches.ar_rms_fusion} " f"where ar+rms+quant was activated" ) + elif ( + match_name == "async_tp" + and "sequence_parallel" in matches_check + and num_compile_ranges >= 2 + ): + # AsyncTP only finds patterns on ranges where SP ran. + n_sp_ranges = num_compile_ranges - 1 + assert ( + sum(m == expected_matches for m in log_matches) + == tp_size * n_sp_ranges + ), ( + f"Expecting {expected_matches} async_tp on " + f"{tp_size * n_sp_ranges} SP-range entries, " + f"found: {log_matches}" + ) + assert sum(m == 0 for m in log_matches) == tp_size, ( + f"Expecting 0 async_tp on {tp_size} small-range entries " + f"(no SP), found: {log_matches}" + ) + elif ( + match_name == "ar_rms_fusion" + and "sequence_parallel" in matches_check + and num_compile_ranges >= 2 + ): + # SP consumes allreduce patterns first, so AR+RMS finds + # full matches only on the smallest range (no SP). + assert sum(m == expected_matches for m in log_matches) == tp_size, ( + f"Expecting {expected_matches} ar_rms on " + f"{tp_size} small-range entries, found: {log_matches}" + ) + assert sum(m == 0 for m in log_matches) == tp_size * ( + num_ranges_activated - 1 + ), ( + f"Expecting 0 ar_rms on " + f"{tp_size * (num_ranges_activated - 1)} large-range " + f"entries (SP took precedence), found: {log_matches}" + ) else: expected_matches_list = [expected_matches] * n_expected assert sorted(log_matches) == expected_matches_list, ( @@ -142,7 +200,7 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): f"found: {sorted(log_matches)}" ) - if match_name == "ar_rms_fusion": + if match_name == "ar_rms_fusion" and num_compile_ranges >= 2: log_matches = re.findall( r"pass_manager.py:\d+] Skipping " r".*AllReduceFusionPass.* with compile range", @@ -155,4 +213,17 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn): f"(found {len(log_matches)}) in:\n {log_holder.text}" ) + if match_name == "sequence_parallel" and num_compile_ranges >= 2: + log_matches = re.findall( + r"pass_manager.py:\d+] Skipping " + r".*SequenceParallelismPass.* with compile range", + log_holder.text, + ) + + n_expected = tp_size * (num_compile_ranges - num_ranges_activated) + assert len(log_matches) == n_expected, ( + f'Could not find {n_expected} "Skipping SequenceParallelismPass" ' + f"(found {len(log_matches)}) in:\n {log_holder.text}" + ) + return run diff --git a/tests/compile/fusions_e2e/test_tp2_async_tp.py b/tests/compile/fusions_e2e/test_tp2_async_tp.py index 4769ca1e0b6..921839ea069 100644 --- a/tests/compile/fusions_e2e/test_tp2_async_tp.py +++ b/tests/compile/fusions_e2e/test_tp2_async_tp.py @@ -66,6 +66,9 @@ def test_tp2_async_tp_fp8_fusions( enable_qk_norm_rope_fusion=True, enable_sp=True, fuse_gemm_comms=True, + fuse_allreduce_rms=False, + # Override threshold for testing (models have small hidden_size) + sp_min_token_num=512, ), ) @@ -123,6 +126,9 @@ def test_tp2_async_tp_fusions( enable_qk_norm_rope_fusion=True, enable_sp=True, fuse_gemm_comms=True, + fuse_allreduce_rms=False, + # Override threshold for testing (models have small hidden_size) + sp_min_token_num=512, ), ) @@ -141,3 +147,130 @@ def test_tp2_async_tp_fusions( matches_check, tp_size=2, ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b_fp8, llama4_scout_fp8], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp2_sp_ar_rms_fp8_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, + monkeypatch, +): + matches = matches_fn(n_layers) + + if is_blackwell(): + # Disable FlashInfer scaled_mm FP8 as it's not supported in async tp patterns + monkeypatch.setenv("VLLM_DISABLED_KERNELS", "FlashInferFP8ScaledMMLinearKernel") + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + fuse_norm_quant=True, + fuse_act_quant=True, + fuse_attn_quant=True, + enable_qk_norm_rope_fusion=True, + enable_sp=True, + fuse_gemm_comms=True, + fuse_allreduce_rms=True, + # Override threshold for testing (models have small hidden_size) + sp_min_token_num=512, + ), + ) + + matches_check = [ + "rms_quant_fusion", + "act_quant_fusion", + "norm_rope_fusion", + "attn_quant_fusion", + "ar_rms_fusion", + "sequence_parallel", + "async_tp", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) + + +@multi_gpu_test(num_gpus=2) +@pytest.mark.parametrize( + "model_name, matches_fn, model_kwargs, hf_overrides", + [llama3_8b, qwen3_a3b], +) +@pytest.mark.parametrize("attn_backend", [TRITON_ATTN]) +@pytest.mark.parametrize("n_layers", [4]) +@pytest.mark.parametrize("custom_ops", custom_ops_combos("rms_norm")) +@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION) +def test_tp2_sp_ar_rms_fusions( + model_name: str, + matches_fn: Callable[[int], Matches], + model_kwargs: dict, + hf_overrides: Callable[[int], dict], + attn_backend: AttentionBackendCase, + n_layers: int, + custom_ops: str, + inductor_graph_partition: bool, + run_e2e_fusion_test, +): + matches = matches_fn(n_layers) + + # Reduce size of model and skip weight loading time + model_kwargs["hf_overrides"] = hf_overrides(n_layers) + model_kwargs["load_format"] = "dummy" + model_kwargs["max_model_len"] = 1024 + + compilation_config = dict( + use_inductor_graph_partition=inductor_graph_partition, + custom_ops=custom_ops.split(","), + pass_config=PassConfig( + enable_qk_norm_rope_fusion=True, + enable_sp=True, + fuse_gemm_comms=True, + fuse_allreduce_rms=True, + # Override threshold for testing (models have small hidden_size) + sp_min_token_num=512, + ), + ) + + matches_check = [ + "norm_rope_fusion", + "ar_rms_fusion", + "sequence_parallel", + "async_tp", + ] + + run_e2e_fusion_test( + model_name, + matches, + model_kwargs, + attn_backend, + compilation_config, + matches_check, + tp_size=2, + ) diff --git a/tests/compile/test_config.py b/tests/compile/test_config.py index eb2f0669ed5..3ba70b6aad3 100644 --- a/tests/compile/test_config.py +++ b/tests/compile/test_config.py @@ -421,6 +421,7 @@ def test_cudagraph_sizes_post_init( fuse_norm_quant=True, fuse_act_quant=True, eliminate_noops=True, + sp_min_token_num=512 if enable_sp else None, ), cudagraph_mode=cudagraph_mode, ) diff --git a/tests/compile/test_sequence_parallelism_threshold.py b/tests/compile/test_sequence_parallelism_threshold.py new file mode 100644 index 00000000000..42e374cd95d --- /dev/null +++ b/tests/compile/test_sequence_parallelism_threshold.py @@ -0,0 +1,110 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.compilation.passes.fusion.sequence_parallelism import ( + SP_MIN_HIDDEN_SIZE, + SP_MIN_PER_GPU_SIZE_MB, + get_sequence_parallelism_threshold, +) + + +class TestGetSequenceParallelismThreshold: + """Tests for get_sequence_parallelism_threshold function.""" + + def test_non_cuda_returns_none(self, mock_cuda_platform): + """Non-CUDA platforms should return None.""" + with mock_cuda_platform(is_cuda=False): + result = get_sequence_parallelism_threshold( + hidden_size=8192, tp_size=2, element_size=2 + ) + assert result is None + + def test_unsupported_device_capability_returns_none(self, mock_cuda_platform): + """Unsupported device capabilities (e.g., sm80) should return None.""" + with mock_cuda_platform(capability=(8, 0)): + result = get_sequence_parallelism_threshold( + hidden_size=8192, tp_size=2, element_size=2 + ) + assert result is None + + def test_small_hidden_size_returns_none(self, mock_cuda_platform): + """H100 with hidden_size below threshold should return None.""" + with mock_cuda_platform(capability=(9, 0)): + result = get_sequence_parallelism_threshold( + hidden_size=4096, + tp_size=2, + element_size=2, # 4096 < 8192 + ) + assert result is None + + def test_h100_large_model_returns_threshold(self, mock_cuda_platform): + """H100 with large enough hidden_size should return calculated threshold.""" + with mock_cuda_platform(capability=(9, 0)): + hidden_size = 8192 + tp_size = 2 + element_size = 2 # float16/bfloat16 + + result = get_sequence_parallelism_threshold( + hidden_size=hidden_size, + tp_size=tp_size, + element_size=element_size, + ) + + # Verify calculation: (8 * 2 * 1024 * 1024) // (8192 * 2) = 1024 + MiB = 1024 * 1024 + expected = int( + (SP_MIN_PER_GPU_SIZE_MB[90] * tp_size * MiB) + // (hidden_size * element_size) + ) + assert result == expected + assert result == 1024 + + @pytest.mark.parametrize( + "hidden_size,tp_size,element_size,expected", + [ + # Boundary: exactly at min hidden size threshold, tp_size=1 + # (8 * 1 * 1024 * 1024) // (8192 * 2) = 512 + (8192, 1, 2, 512), + # Larger hidden size reduces token threshold + # (8 * 1 * 1024 * 1024) // (16384 * 2) = 256 + (16384, 1, 2, 256), + # Larger tp_size increases token threshold + # (8 * 4 * 1024 * 1024) // (8192 * 2) = 2048 + (8192, 4, 2, 2048), + # Larger element_size (fp32) reduces token threshold + # (8 * 2 * 1024 * 1024) // (8192 * 4) = 512 + (8192, 2, 4, 512), + ], + ) + def test_threshold_calculation_variations( + self, mock_cuda_platform, hidden_size, tp_size, element_size, expected + ): + """Test threshold calculation with various parameter combinations.""" + with mock_cuda_platform(capability=(9, 0)): + result = get_sequence_parallelism_threshold( + hidden_size=hidden_size, + tp_size=tp_size, + element_size=element_size, + ) + assert result == expected + + def test_hidden_size_boundary(self, mock_cuda_platform): + """Test behavior at the exact hidden_size boundary.""" + with mock_cuda_platform(capability=(9, 0)): + # Just below threshold + result = get_sequence_parallelism_threshold( + hidden_size=SP_MIN_HIDDEN_SIZE[90] - 1, + tp_size=2, + element_size=2, + ) + assert result is None + + # Exactly at threshold + result = get_sequence_parallelism_threshold( + hidden_size=SP_MIN_HIDDEN_SIZE[90], + tp_size=2, + element_size=2, + ) + assert result is not None diff --git a/vllm/compilation/passes/fusion/sequence_parallelism.py b/vllm/compilation/passes/fusion/sequence_parallelism.py index 5fb932d7284..63de85932cb 100644 --- a/vllm/compilation/passes/fusion/sequence_parallelism.py +++ b/vllm/compilation/passes/fusion/sequence_parallelism.py @@ -27,6 +27,63 @@ from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNo logger = init_logger(__name__) +# Min hidden size per device capability for sequence parallelism +# Only apply sequence parallelism for models with hidden_size >= threshold +SP_MIN_HIDDEN_SIZE: dict[int, int] = { + 90: 8192, # H100: only for models with hidden_size >= 8192 +} + +# Min size per GPU per device capability for sequence parallelism +# Total min size = min_per_gpu_size * tp_size +# This ensures the threshold scales appropriately with tensor parallelism +SP_MIN_PER_GPU_SIZE_MB: dict[int, float] = { + 90: 8, # 8MB per GPU for H100 +} + + +def get_sequence_parallelism_threshold( + hidden_size: int, + tp_size: int, + element_size: int, +) -> int | None: + """ + Calculate the minimum token threshold for applying sequence parallelism. + + Returns None if sequence parallelism should not be applied based on model size. + + Branching logic based on device capability: + - Check if hidden_size >= SP_MIN_HIDDEN_SIZE[device_capability] + - If not, returns None (SP disabled for small models on this device) + - If yes, calculates threshold based on per-GPU size + + Formula: min_token_num = (min_per_gpu_size_mb * tp_size * MiB) // + (hidden_size * element_size) + """ + from vllm.platforms import current_platform + + if not current_platform.is_cuda(): + return None + + capability = current_platform.get_device_capability() + if capability is None: + return None + device_capability = capability.to_int() + + # Check if device has configured thresholds + min_hidden_size = SP_MIN_HIDDEN_SIZE.get(device_capability) + min_per_gpu_size_mb = SP_MIN_PER_GPU_SIZE_MB.get(device_capability) + + if min_hidden_size is None or min_per_gpu_size_mb is None: + return None + + # Only apply sequence parallelism for models meeting the size threshold + if hidden_size < min_hidden_size: + return None + + MiB = 1024 * 1024 + min_size = min_per_gpu_size_mb * MiB * tp_size + return int(min_size // (hidden_size * element_size)) + def get_first_out_wrapper( fn: Callable[..., Sequence[torch.Tensor]], @@ -309,6 +366,23 @@ class SequenceParallelismPass(VllmPatternMatcherPass): def __init__(self, config: VllmConfig) -> None: super().__init__(config) + # Get min_token_num threshold + # Read min_token_num from config (calculated during config init) + self.min_token_num = None + if config.model_config is not None: + pass_config = config.compilation_config.pass_config + self.min_token_num = pass_config.sp_min_token_num + + if self.min_token_num is not None: + # Take the min to avoid exceeding max_num_batched_tokens + max_batched = config.scheduler_config.max_num_batched_tokens + if max_batched is not None: + self.min_token_num = min(self.min_token_num, max_batched) + logger.debug_once( + f"Sequence parallelism min token threshold: {self.min_token_num}", + scope="global", + ) + # Used to clean up redundant views created temporarily # to circumvent residual shape change issues self.noop_cleanup = NoOpEliminationPass(config) @@ -339,29 +413,36 @@ class SequenceParallelismPass(VllmPatternMatcherPass): self.dump_patterns(config, self.patterns) def is_applicable_for_range(self, compile_range: Range) -> bool: - # When sequence parallelism is enabled, the residual tensor from RMSNorm - # needs to be split along the sequence dimension. However, this dimension - # is symbolic during piecewise compilation, and splitting symbolic shapes - # is not supported. - # - # This pass is therefore only applied when the sequence dimension is - # concrete: - # 1. In full-graph compilation mode (no Dynamo splitting ops are used). - # For this case we always pad num_tokens to be a multiple of - # tensor_parallel_size, so there's no need to check shape % tp_size == 0. - # 2. For specific shape provided during compilation (e.g., from - # `compile_sizes`), which must be divisible by the tensor-parallel - # size. + """ + Determines if sequence parallelism should be applied for the given + compile range. + + SP is only beneficial for larger batch sizes where the communication + overhead is amortized. For small batches, the overhead of splitting + and gathering tensors across TP ranks outweighs the benefits. + + Returns False (SP disabled) when: + - Using piecewise compilation with non-concrete or TP-indivisible sizes + - min_token_num is None (SP disabled for this device/config) + - The compile range starts below the minimum token threshold + """ + # For piecewise compilation (not using inductor graph partition), + # we need concrete sizes that are divisible by TP for correct splitting if ( - not self.compilation_config.splitting_ops - or self.compilation_config.use_inductor_graph_partition + not self.compilation_config.use_inductor_graph_partition + and self.compilation_config.splitting_ops ): - return True - tp_size = get_tensor_model_parallel_world_size() - result: bool = (compile_range.is_single_size()) and ( - compile_range.end % tp_size == 0 - ) - return result + tp_size = get_tensor_model_parallel_world_size() + if not compile_range.is_single_size() or compile_range.end % tp_size != 0: + return False + + # min_token_num is None when SP is disabled for this device/config + # (e.g., non-CUDA platform, unsupported GPU, or small hidden_size) + if self.min_token_num is None: + return False + + # Only apply SP when batch size meets the minimum threshold + return compile_range.start >= self.min_token_num @VllmInductorPass.time_and_log def __call__(self, graph: fx.Graph) -> None: diff --git a/vllm/config/compilation.py b/vllm/config/compilation.py index ab6f3da06cd..d22e9a96e0f 100644 --- a/vllm/config/compilation.py +++ b/vllm/config/compilation.py @@ -118,7 +118,9 @@ class PassConfig: eliminate_noops: bool = Field(default=True) """Eliminate no-op ops.""" enable_sp: bool = Field(default=None) - """Enable sequence parallelism.""" + """Enable sequence parallelism. Requires TP>1. Automatically disabled + if the model's hidden_size is too small for SP to be beneficial + (threshold is device-capability dependent).""" fuse_gemm_comms: bool = Field(default=None) """Enable async TP.""" fuse_allreduce_rms: bool = Field(default=None) @@ -155,6 +157,11 @@ class PassConfig: 8: 1, # 1MB }, }, where key is the device capability""" + sp_min_token_num: int | None = None + """The minimum number of tokens above which vllm should use + sequence parallelism. Specified as an integer token count. + Unspecified will fallback to default values which are compute + capability and world size dependent.""" # TODO(luka) better pass enabling system. diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index ef71a05d393..fba3c64a9af 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -853,8 +853,33 @@ class VllmConfig: logger.warning("Sequence Parallelism requires TP>1, disabling") self.compilation_config.pass_config.enable_sp = False self.compilation_config.pass_config.fuse_gemm_comms = False + else: + # Compute SP threshold early; disable if None (model too + # small) before +rms_norm gets forced into custom_ops. + pass_config = self.compilation_config.pass_config + if pass_config.sp_min_token_num is None: + from vllm.compilation.passes.fusion.sequence_parallelism import ( + get_sequence_parallelism_threshold, + ) - elif "-rms_norm" in self.compilation_config.custom_ops: + tp_size = self.parallel_config.tensor_parallel_size + hidden_size = self.model_config.get_hidden_size() + element_size = self.model_config.dtype.itemsize + pass_config.sp_min_token_num = get_sequence_parallelism_threshold( + hidden_size, tp_size, element_size + ) + + if pass_config.sp_min_token_num is None: + logger.warning( + "Model hidden_size too small for the SP " + "threshold heuristic, disabling. To force SP, " + "set pass_config.sp_min_token_num manually." + ) + self.compilation_config.pass_config.enable_sp = False + self.compilation_config.pass_config.fuse_gemm_comms = False + + if self.compilation_config.pass_config.enable_sp: + if "-rms_norm" in self.compilation_config.custom_ops: logger.warning( "RMS norm force disabled, sequence parallelism might break" ) @@ -1456,6 +1481,36 @@ class VllmConfig: "allreduce-rms fusion will be enabled for all num_tokens." ) + # Add the compile ranges for sequence parallelism + if compilation_config.pass_config.enable_sp: + pass_config = compilation_config.pass_config + + # Calculate min_token_num if not explicitly provided + # User override works regardless of hidden_size + if pass_config.sp_min_token_num is None: + from vllm.compilation.passes.fusion.sequence_parallelism import ( + get_sequence_parallelism_threshold, + ) + + tp_size = self.parallel_config.tensor_parallel_size + hidden_size = self.model_config.get_hidden_size() + element_size = self.model_config.dtype.itemsize + pass_config.sp_min_token_num = get_sequence_parallelism_threshold( + hidden_size, tp_size, element_size + ) + + min_token_num = pass_config.sp_min_token_num + max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens + if min_token_num is not None and ( + max_num_batched_tokens is not None + and min_token_num < max_num_batched_tokens + and min_token_num > 1 + ): + # Add split point at min_token_num - 1 to ensure SP applies + # starting from min_token_num + # This creates ranges: [1, min-1] (no SP), [min, max] (SP applies) + computed_compile_ranges_split_points.append(min_token_num - 1) + if compilation_config.pass_config.fuse_rope_kvcache: max_token_num = ( compilation_config.pass_config.rope_kvcache_fusion_max_token_num From 4a9c07a0a2b8308a045476b48be29e37c349274b Mon Sep 17 00:00:00 2001 From: Daniele <36171005+dtrifiro@users.noreply.github.com> Date: Thu, 26 Feb 2026 06:39:48 +0100 Subject: [PATCH 23/43] [BugFix] anthropic/serving_messages: fix tool call arguments streaming (#34887) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Daniele Trifirò Co-authored-by: Nicolò Lucchesi --- vllm/entrypoints/anthropic/serving.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 8fb347aabed..dc037313de3 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -432,6 +432,19 @@ class AnthropicServingMessages(OpenAIServingChat): data = chunk.model_dump_json(exclude_unset=True) yield wrap_data_with_event(data, "content_block_start") content_block_started = True + if tool_call.function and tool_call.function.arguments: + chunk = AnthropicStreamEvent( + index=content_block_index, + type="content_block_delta", + delta=AnthropicDelta( + type="input_json_delta", + partial_json=tool_call.function.arguments, + ), + ) + data = chunk.model_dump_json(exclude_unset=True) + yield wrap_data_with_event( + data, "content_block_delta" + ) else: chunk = AnthropicStreamEvent( From 186ea22efefd2c6f4f9b7fcb657bd00f50cb465a Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Thu, 26 Feb 2026 01:35:16 -0500 Subject: [PATCH 24/43] [Misc][Harmony] Move Responses API only harmony utils to responses/harmony.py (#35339) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../openai/parser/test_harmony_utils.py | 467 +-------------- .../openai/responses/test_harmony_utils.py | 463 +++++++++++++++ .../openai/responses/test_mcp_tools.py | 10 +- .../openai/parser/harmony_utils.py | 518 +--------------- vllm/entrypoints/openai/responses/harmony.py | 552 ++++++++++++++++++ vllm/entrypoints/openai/responses/serving.py | 20 +- 6 files changed, 1040 insertions(+), 990 deletions(-) create mode 100644 tests/entrypoints/openai/responses/test_harmony_utils.py create mode 100644 vllm/entrypoints/openai/responses/harmony.py diff --git a/tests/entrypoints/openai/parser/test_harmony_utils.py b/tests/entrypoints/openai/parser/test_harmony_utils.py index b73a0b0745c..7842a1fcd75 100644 --- a/tests/entrypoints/openai/parser/test_harmony_utils.py +++ b/tests/entrypoints/openai/parser/test_harmony_utils.py @@ -2,13 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -from openai.types.responses import ( - ResponseFunctionToolCall, - ResponseOutputMessage, - ResponseReasoningItem, -) -from openai.types.responses.response_output_item import McpCall -from openai_harmony import Author, Message, Role, TextContent +from openai_harmony import Message, Role from tests.entrypoints.openai.utils import verify_harmony_messages from vllm.entrypoints.openai.parser.harmony_utils import ( @@ -18,20 +12,21 @@ from vllm.entrypoints.openai.parser.harmony_utils import ( has_custom_tools, parse_chat_input_to_harmony_message, parse_chat_output, - parse_input_to_harmony_message, - parse_output_message, +) +from vllm.entrypoints.openai.responses.harmony import ( + response_previous_input_to_harmony, ) class TestCommonParseInputToHarmonyMessage: """ Tests for scenarios that are common to both Chat Completion - parse_chat_input_to_harmony_message and Responsees API - parse_input_to_harmony_message functions. + parse_chat_input_to_harmony_message and Responses API + response_previous_input_to_harmony functions. """ @pytest.fixture( - params=[parse_chat_input_to_harmony_message, parse_input_to_harmony_message] + params=[parse_chat_input_to_harmony_message, response_previous_input_to_harmony] ) def parse_function(self, request): return request.param @@ -216,81 +211,6 @@ class TestCommonParseInputToHarmonyMessage: assert messages[0].content[1].text == "actual text" -class TestParseInputToHarmonyMessage: - """ - Tests for scenarios that are specific to the Responses API - parse_input_to_harmony_message function. - """ - - def test_message_with_empty_content(self): - """Test parsing message with empty string content.""" - chat_msg = { - "role": "user", - "content": "", - } - - messages = parse_input_to_harmony_message(chat_msg) - - assert len(messages) == 1 - assert messages[0].content[0].text == "" - - def test_tool_message_with_string_content(self): - """Test parsing tool message with string content.""" - chat_msg = { - "role": "tool", - "name": "get_weather", - "content": "The weather in San Francisco is sunny, 72°F", - } - - messages = parse_input_to_harmony_message(chat_msg) - - assert len(messages) == 1 - assert messages[0].author.role == Role.TOOL - assert messages[0].author.name == "functions.get_weather" - assert ( - messages[0].content[0].text == "The weather in San Francisco is sunny, 72°F" - ) - assert messages[0].channel == "commentary" - - def test_tool_message_with_array_content(self): - """Test parsing tool message with array content.""" - chat_msg = { - "role": "tool", - "name": "search_results", - "content": [ - {"type": "text", "text": "Result 1: "}, - {"type": "text", "text": "Result 2: "}, - { - "type": "image", - "url": "http://example.com/img.png", - }, # Should be ignored - {"type": "text", "text": "Result 3"}, - ], - } - - messages = parse_input_to_harmony_message(chat_msg) - - assert len(messages) == 1 - assert messages[0].author.role == Role.TOOL - assert messages[0].author.name == "functions.search_results" - assert messages[0].content[0].text == "Result 1: Result 2: Result 3" - - def test_tool_message_with_empty_content(self): - """Test parsing tool message with None content.""" - chat_msg = { - "role": "tool", - "name": "empty_tool", - "content": None, - } - - messages = parse_input_to_harmony_message(chat_msg) - - assert len(messages) == 1 - assert messages[0].author.role == Role.TOOL - assert messages[0].author.name == "functions.empty_tool" - assert messages[0].content[0].text == "" - - class TestParseChatInputToHarmonyMessage: """ Tests for scenarios that are specific to the Chat Completion API @@ -888,200 +808,6 @@ class TestParseChatOutput: assert final_content == "Let me look that up.\nThe answer is 42." -class TestParseOutputMessage: - """Tests for parse_output_message function.""" - - def test_commentary_with_no_recipient_creates_message(self): - """Test that commentary with recipient=None (preambles) creates message items. - - Per Harmony format, preambles are intended to be shown to end-users, - unlike analysis channel content which is hidden reasoning. - See: https://cookbook.openai.com/articles/openai-harmony - """ - message = Message.from_role_and_content( - Role.ASSISTANT, "I will now search for the weather information." - ) - message = message.with_channel("commentary") - # recipient is None by default, representing a preamble - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseOutputMessage) - assert output_items[0].type == "message" - assert output_items[0].role == "assistant" - assert output_items[0].status == "completed" - assert len(output_items[0].content) == 1 - assert output_items[0].content[0].type == "output_text" - assert ( - output_items[0].content[0].text - == "I will now search for the weather information." - ) - - def test_commentary_with_function_recipient_creates_function_call(self): - """Test commentary with recipient='functions.X' creates function calls.""" - message = Message.from_role_and_content( - Role.ASSISTANT, '{"location": "San Francisco", "units": "celsius"}' - ) - message = message.with_channel("commentary") - message = message.with_recipient("functions.get_weather") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseFunctionToolCall) - assert output_items[0].type == "function_call" - assert output_items[0].name == "get_weather" - assert ( - output_items[0].arguments - == '{"location": "San Francisco", "units": "celsius"}' - ) - assert output_items[0].call_id.startswith("call_") - assert output_items[0].id.startswith("fc_") - - def test_commentary_with_python_recipient_creates_reasoning(self): - """Test that commentary with recipient='python' creates reasoning items.""" - message = Message.from_role_and_content( - Role.ASSISTANT, "import numpy as np\nprint(np.array([1, 2, 3]))" - ) - message = message.with_channel("commentary") - message = message.with_recipient("python") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseReasoningItem) - assert output_items[0].type == "reasoning" - assert ( - output_items[0].content[0].text - == "import numpy as np\nprint(np.array([1, 2, 3]))" - ) - - def test_commentary_with_browser_recipient_creates_reasoning(self): - """Test that commentary with recipient='browser' creates reasoning items.""" - message = Message.from_role_and_content( - Role.ASSISTANT, "Navigating to the specified URL" - ) - message = message.with_channel("commentary") - message = message.with_recipient("browser") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseReasoningItem) - assert output_items[0].type == "reasoning" - assert output_items[0].content[0].text == "Navigating to the specified URL" - - def test_commentary_with_container_recipient_creates_reasoning(self): - """Test that commentary with recipient='container' creates reasoning items.""" - message = Message.from_role_and_content( - Role.ASSISTANT, "Running command in container" - ) - message = message.with_channel("commentary") - message = message.with_recipient("container") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseReasoningItem) - assert output_items[0].type == "reasoning" - assert output_items[0].content[0].text == "Running command in container" - - def test_commentary_with_empty_content_and_no_recipient(self): - """Test edge case: empty commentary with recipient=None.""" - message = Message.from_role_and_content(Role.ASSISTANT, "") - message = message.with_channel("commentary") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseOutputMessage) - assert output_items[0].content[0].text == "" - - def test_commentary_with_multiple_contents_and_no_recipient(self): - """Test multiple content items in commentary with no recipient.""" - contents = [ - TextContent(text="Step 1: Analyze the request"), - TextContent(text="Step 2: Prepare to call functions"), - ] - message = Message.from_role_and_contents(Role.ASSISTANT, contents) - message = message.with_channel("commentary") - - output_items = parse_output_message(message) - - # _parse_final_message returns single ResponseOutputMessage with - # multiple contents - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseOutputMessage) - assert len(output_items[0].content) == 2 - assert output_items[0].content[0].text == "Step 1: Analyze the request" - assert output_items[0].content[1].text == "Step 2: Prepare to call functions" - - def test_commentary_with_multiple_function_calls(self): - """Test multiple function calls in commentary channel.""" - contents = [ - TextContent(text='{"location": "San Francisco"}'), - TextContent(text='{"location": "New York"}'), - ] - message = Message.from_role_and_contents(Role.ASSISTANT, contents) - message = message.with_channel("commentary") - message = message.with_recipient("functions.get_weather") - - output_items = parse_output_message(message) - - assert len(output_items) == 2 - assert all(isinstance(item, ResponseFunctionToolCall) for item in output_items) - assert output_items[0].name == "get_weather" - assert output_items[1].name == "get_weather" - assert output_items[0].arguments == '{"location": "San Francisco"}' - assert output_items[1].arguments == '{"location": "New York"}' - - def test_commentary_with_unknown_recipient_creates_mcp_call(self): - """Test that commentary with unknown recipient creates MCP call.""" - message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}') - message = message.with_channel("commentary") - message = message.with_recipient("custom_tool") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], McpCall) - assert output_items[0].type == "mcp_call" - assert output_items[0].name == "custom_tool" - assert output_items[0].server_label == "custom_tool" - - def test_analysis_channel_creates_reasoning(self): - """Test that analysis channel creates reasoning items.""" - message = Message.from_role_and_content( - Role.ASSISTANT, "Analyzing the problem step by step..." - ) - message = message.with_channel("analysis") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], ResponseReasoningItem) - assert output_items[0].type == "reasoning" - assert ( - output_items[0].content[0].text == "Analyzing the problem step by step..." - ) - - def test_non_assistant_message_returns_empty(self): - """Test that non-assistant messages return empty list. - - Per the implementation, tool messages to assistant (e.g., search results) - are not included in final output to align with OpenAI behavior. - """ - message = Message.from_author_and_content( - Author.new(Role.TOOL, "functions.get_weather"), - "The weather is sunny, 72°F", - ) - - output_items = parse_output_message(message) - - assert len(output_items) == 0 - - def test_has_custom_tools() -> None: assert not has_custom_tools(set()) assert not has_custom_tools({"web_search_preview", "code_interpreter", "container"}) @@ -1091,185 +817,6 @@ def test_has_custom_tools() -> None: ) -def test_parse_mcp_call_basic() -> None: - """Test that MCP calls are parsed with correct type and server_label.""" - message = Message.from_role_and_content(Role.ASSISTANT, '{"path": "/tmp"}') - message = message.with_recipient("filesystem") - message = message.with_channel("commentary") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], McpCall) - assert output_items[0].type == "mcp_call" - assert output_items[0].name == "filesystem" - assert output_items[0].server_label == "filesystem" - assert output_items[0].arguments == '{"path": "/tmp"}' - assert output_items[0].status == "completed" - - -def test_parse_mcp_call_dotted_recipient() -> None: - """Test that dotted recipients extract the tool name correctly.""" - message = Message.from_role_and_content(Role.ASSISTANT, '{"cmd": "ls"}') - message = message.with_recipient("repo_browser.list") - message = message.with_channel("commentary") - - output_items = parse_output_message(message) - - assert len(output_items) == 1 - assert isinstance(output_items[0], McpCall) - assert output_items[0].name == "list" - assert output_items[0].server_label == "repo_browser" - - -def test_mcp_vs_function_call() -> None: - """Test that function calls are not parsed as MCP calls.""" - func_message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}') - func_message = func_message.with_recipient("functions.my_tool") - func_message = func_message.with_channel("commentary") - - func_items = parse_output_message(func_message) - - assert len(func_items) == 1 - assert not isinstance(func_items[0], McpCall) - assert func_items[0].type == "function_call" - - -def test_mcp_vs_builtin_tools() -> None: - """Test that built-in tools (python, container) are not parsed as MCP calls.""" - # Test python (built-in tool) - should be reasoning, not MCP - python_message = Message.from_role_and_content(Role.ASSISTANT, "print('hello')") - python_message = python_message.with_recipient("python") - python_message = python_message.with_channel("commentary") - - python_items = parse_output_message(python_message) - - assert len(python_items) == 1 - assert not isinstance(python_items[0], McpCall) - assert python_items[0].type == "reasoning" - - -def test_parse_remaining_state_commentary_channel() -> None: - """Test parse_remaining_state with commentary channel and various recipients.""" - from unittest.mock import Mock - - from vllm.entrypoints.openai.parser.harmony_utils import parse_remaining_state - - # Test 1: functions.* recipient → should return function tool call - parser_func = Mock() - parser_func.current_content = '{"arg": "value"}' - parser_func.current_role = Role.ASSISTANT - parser_func.current_channel = "commentary" - parser_func.current_recipient = "functions.my_tool" - - func_items = parse_remaining_state(parser_func) - - assert len(func_items) == 1 - assert not isinstance(func_items[0], McpCall) - assert func_items[0].type == "function_call" - assert func_items[0].name == "my_tool" - assert func_items[0].status == "in_progress" - - # Test 2: MCP tool (not builtin) → should return MCP call - parser_mcp = Mock() - parser_mcp.current_content = '{"path": "/tmp"}' - parser_mcp.current_role = Role.ASSISTANT - parser_mcp.current_channel = "commentary" - parser_mcp.current_recipient = "filesystem" - - mcp_items = parse_remaining_state(parser_mcp) - - assert len(mcp_items) == 1 - assert isinstance(mcp_items[0], McpCall) - assert mcp_items[0].type == "mcp_call" - assert mcp_items[0].name == "filesystem" - assert mcp_items[0].server_label == "filesystem" - assert mcp_items[0].status == "in_progress" - - # Test 3: Built-in tool (python) - # should NOT return MCP call, returns reasoning (internal tool interaction) - parser_builtin = Mock() - parser_builtin.current_content = "print('hello')" - parser_builtin.current_role = Role.ASSISTANT - parser_builtin.current_channel = "commentary" - parser_builtin.current_recipient = "python" - - builtin_items = parse_remaining_state(parser_builtin) - - # Built-in tools explicitly return reasoning - assert len(builtin_items) == 1 - assert not isinstance(builtin_items[0], McpCall) - assert builtin_items[0].type == "reasoning" - - # Test 4: No recipient (preamble) → should return message, not reasoning - parser_preamble = Mock() - parser_preamble.current_content = "I'll search for that information now." - parser_preamble.current_role = Role.ASSISTANT - parser_preamble.current_channel = "commentary" - parser_preamble.current_recipient = None - - preamble_items = parse_remaining_state(parser_preamble) - - assert len(preamble_items) == 1 - assert isinstance(preamble_items[0], ResponseOutputMessage) - assert preamble_items[0].type == "message" - assert preamble_items[0].content[0].text == "I'll search for that information now." - assert preamble_items[0].status == "incomplete" # streaming - - -def test_parse_remaining_state_analysis_channel() -> None: - """Test parse_remaining_state with analysis channel and various recipients.""" - from unittest.mock import Mock - - from vllm.entrypoints.openai.parser.harmony_utils import parse_remaining_state - - # Test 1: functions.* recipient → should return function tool call - parser_func = Mock() - parser_func.current_content = '{"arg": "value"}' - parser_func.current_role = Role.ASSISTANT - parser_func.current_channel = "analysis" - parser_func.current_recipient = "functions.my_tool" - - func_items = parse_remaining_state(parser_func) - - assert len(func_items) == 1 - assert not isinstance(func_items[0], McpCall) - assert func_items[0].type == "function_call" - assert func_items[0].name == "my_tool" - assert func_items[0].status == "in_progress" - - # Test 2: MCP tool (not builtin) → should return MCP call - parser_mcp = Mock() - parser_mcp.current_content = '{"query": "test"}' - parser_mcp.current_role = Role.ASSISTANT - parser_mcp.current_channel = "analysis" - parser_mcp.current_recipient = "database" - - mcp_items = parse_remaining_state(parser_mcp) - - assert len(mcp_items) == 1 - assert isinstance(mcp_items[0], McpCall) - assert mcp_items[0].type == "mcp_call" - assert mcp_items[0].name == "database" - assert mcp_items[0].server_label == "database" - assert mcp_items[0].status == "in_progress" - - # Test 3: Built-in tool (container) - # should NOT return MCP call, falls through to reasoning - parser_builtin = Mock() - parser_builtin.current_content = "docker run" - parser_builtin.current_role = Role.ASSISTANT - parser_builtin.current_channel = "analysis" - parser_builtin.current_recipient = "container" - - builtin_items = parse_remaining_state(parser_builtin) - - # Should fall through to reasoning logic - assert len(builtin_items) == 1 - assert not isinstance(builtin_items[0], McpCall) - assert builtin_items[0].type == "reasoning" - - class TestGetSystemMessage: """Tests for get_system_message channel configuration.""" diff --git a/tests/entrypoints/openai/responses/test_harmony_utils.py b/tests/entrypoints/openai/responses/test_harmony_utils.py new file mode 100644 index 00000000000..e51538298ff --- /dev/null +++ b/tests/entrypoints/openai/responses/test_harmony_utils.py @@ -0,0 +1,463 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for vllm.entrypoints.openai.responses.harmony.""" + +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputMessage, + ResponseReasoningItem, +) +from openai.types.responses.response_output_item import McpCall +from openai_harmony import Author, Message, Role, TextContent + +from vllm.entrypoints.openai.responses.harmony import ( + harmony_to_response_output, + parser_state_to_response_output, + response_previous_input_to_harmony, +) + + +class TestResponsePreviousInputToHarmony: + """ + Tests for scenarios that are specific to the Responses API + response_previous_input_to_harmony function. + """ + + def test_message_with_empty_content(self): + """Test parsing message with empty string content.""" + chat_msg = { + "role": "user", + "content": "", + } + + messages = response_previous_input_to_harmony(chat_msg) + + assert len(messages) == 1 + assert messages[0].content[0].text == "" + + def test_tool_message_with_string_content(self): + """Test parsing tool message with string content.""" + chat_msg = { + "role": "tool", + "name": "get_weather", + "content": "The weather in San Francisco is sunny, 72°F", + } + + messages = response_previous_input_to_harmony(chat_msg) + + assert len(messages) == 1 + assert messages[0].author.role == Role.TOOL + assert messages[0].author.name == "functions.get_weather" + assert ( + messages[0].content[0].text == "The weather in San Francisco is sunny, 72°F" + ) + assert messages[0].channel == "commentary" + + def test_tool_message_with_array_content(self): + """Test parsing tool message with array content.""" + chat_msg = { + "role": "tool", + "name": "search_results", + "content": [ + {"type": "text", "text": "Result 1: "}, + {"type": "text", "text": "Result 2: "}, + { + "type": "image", + "url": "http://example.com/img.png", + }, # Should be ignored + {"type": "text", "text": "Result 3"}, + ], + } + + messages = response_previous_input_to_harmony(chat_msg) + + assert len(messages) == 1 + assert messages[0].author.role == Role.TOOL + assert messages[0].author.name == "functions.search_results" + assert messages[0].content[0].text == "Result 1: Result 2: Result 3" + + def test_tool_message_with_empty_content(self): + """Test parsing tool message with None content.""" + chat_msg = { + "role": "tool", + "name": "empty_tool", + "content": None, + } + + messages = response_previous_input_to_harmony(chat_msg) + + assert len(messages) == 1 + assert messages[0].author.role == Role.TOOL + assert messages[0].author.name == "functions.empty_tool" + assert messages[0].content[0].text == "" + + +class TestHarmonyToResponseOutput: + """Tests for harmony_to_response_output function.""" + + def test_commentary_with_no_recipient_creates_message(self): + """Test that commentary with recipient=None (preambles) creates message items. + + Per Harmony format, preambles are intended to be shown to end-users, + unlike analysis channel content which is hidden reasoning. + See: https://cookbook.openai.com/articles/openai-harmony + """ + message = Message.from_role_and_content( + Role.ASSISTANT, "I will now search for the weather information." + ) + message = message.with_channel("commentary") + # recipient is None by default, representing a preamble + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseOutputMessage) + assert output_items[0].type == "message" + assert output_items[0].role == "assistant" + assert output_items[0].status == "completed" + assert len(output_items[0].content) == 1 + assert output_items[0].content[0].type == "output_text" + assert ( + output_items[0].content[0].text + == "I will now search for the weather information." + ) + + def test_commentary_with_function_recipient_creates_function_call(self): + """Test commentary with recipient='functions.X' creates function calls.""" + message = Message.from_role_and_content( + Role.ASSISTANT, '{"location": "San Francisco", "units": "celsius"}' + ) + message = message.with_channel("commentary") + message = message.with_recipient("functions.get_weather") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseFunctionToolCall) + assert output_items[0].type == "function_call" + assert output_items[0].name == "get_weather" + assert ( + output_items[0].arguments + == '{"location": "San Francisco", "units": "celsius"}' + ) + assert output_items[0].call_id.startswith("call_") + assert output_items[0].id.startswith("fc_") + + def test_commentary_with_python_recipient_creates_reasoning(self): + """Test that commentary with recipient='python' creates reasoning items.""" + message = Message.from_role_and_content( + Role.ASSISTANT, "import numpy as np\nprint(np.array([1, 2, 3]))" + ) + message = message.with_channel("commentary") + message = message.with_recipient("python") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseReasoningItem) + assert output_items[0].type == "reasoning" + assert ( + output_items[0].content[0].text + == "import numpy as np\nprint(np.array([1, 2, 3]))" + ) + + def test_commentary_with_browser_recipient_creates_reasoning(self): + """Test that commentary with recipient='browser' creates reasoning items.""" + message = Message.from_role_and_content( + Role.ASSISTANT, "Navigating to the specified URL" + ) + message = message.with_channel("commentary") + message = message.with_recipient("browser") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseReasoningItem) + assert output_items[0].type == "reasoning" + assert output_items[0].content[0].text == "Navigating to the specified URL" + + def test_commentary_with_container_recipient_creates_reasoning(self): + """Test that commentary with recipient='container' creates reasoning items.""" + message = Message.from_role_and_content( + Role.ASSISTANT, "Running command in container" + ) + message = message.with_channel("commentary") + message = message.with_recipient("container") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseReasoningItem) + assert output_items[0].type == "reasoning" + assert output_items[0].content[0].text == "Running command in container" + + def test_commentary_with_empty_content_and_no_recipient(self): + """Test edge case: empty commentary with recipient=None.""" + message = Message.from_role_and_content(Role.ASSISTANT, "") + message = message.with_channel("commentary") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseOutputMessage) + assert output_items[0].content[0].text == "" + + def test_commentary_with_multiple_contents_and_no_recipient(self): + """Test multiple content items in commentary with no recipient.""" + contents = [ + TextContent(text="Step 1: Analyze the request"), + TextContent(text="Step 2: Prepare to call functions"), + ] + message = Message.from_role_and_contents(Role.ASSISTANT, contents) + message = message.with_channel("commentary") + + output_items = harmony_to_response_output(message) + + # _parse_final_message returns single ResponseOutputMessage with + # multiple contents + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseOutputMessage) + assert len(output_items[0].content) == 2 + assert output_items[0].content[0].text == "Step 1: Analyze the request" + assert output_items[0].content[1].text == "Step 2: Prepare to call functions" + + def test_commentary_with_multiple_function_calls(self): + """Test multiple function calls in commentary channel.""" + contents = [ + TextContent(text='{"location": "San Francisco"}'), + TextContent(text='{"location": "New York"}'), + ] + message = Message.from_role_and_contents(Role.ASSISTANT, contents) + message = message.with_channel("commentary") + message = message.with_recipient("functions.get_weather") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 2 + assert all(isinstance(item, ResponseFunctionToolCall) for item in output_items) + assert output_items[0].name == "get_weather" + assert output_items[1].name == "get_weather" + assert output_items[0].arguments == '{"location": "San Francisco"}' + assert output_items[1].arguments == '{"location": "New York"}' + + def test_commentary_with_unknown_recipient_creates_mcp_call(self): + """Test that commentary with unknown recipient creates MCP call.""" + message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}') + message = message.with_channel("commentary") + message = message.with_recipient("custom_tool") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], McpCall) + assert output_items[0].type == "mcp_call" + assert output_items[0].name == "custom_tool" + assert output_items[0].server_label == "custom_tool" + + def test_analysis_channel_creates_reasoning(self): + """Test that analysis channel creates reasoning items.""" + message = Message.from_role_and_content( + Role.ASSISTANT, "Analyzing the problem step by step..." + ) + message = message.with_channel("analysis") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], ResponseReasoningItem) + assert output_items[0].type == "reasoning" + assert ( + output_items[0].content[0].text == "Analyzing the problem step by step..." + ) + + def test_non_assistant_message_returns_empty(self): + """Test that non-assistant messages return empty list. + + Per the implementation, tool messages to assistant (e.g., search results) + are not included in final output to align with OpenAI behavior. + """ + message = Message.from_author_and_content( + Author.new(Role.TOOL, "functions.get_weather"), + "The weather is sunny, 72°F", + ) + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 0 + + +def test_parse_mcp_call_basic() -> None: + """Test that MCP calls are parsed with correct type and server_label.""" + message = Message.from_role_and_content(Role.ASSISTANT, '{"path": "/tmp"}') + message = message.with_recipient("filesystem") + message = message.with_channel("commentary") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], McpCall) + assert output_items[0].type == "mcp_call" + assert output_items[0].name == "filesystem" + assert output_items[0].server_label == "filesystem" + assert output_items[0].arguments == '{"path": "/tmp"}' + assert output_items[0].status == "completed" + + +def test_parse_mcp_call_dotted_recipient() -> None: + """Test that dotted recipients extract the tool name correctly.""" + message = Message.from_role_and_content(Role.ASSISTANT, '{"cmd": "ls"}') + message = message.with_recipient("repo_browser.list") + message = message.with_channel("commentary") + + output_items = harmony_to_response_output(message) + + assert len(output_items) == 1 + assert isinstance(output_items[0], McpCall) + assert output_items[0].name == "list" + assert output_items[0].server_label == "repo_browser" + + +def test_mcp_vs_function_call() -> None: + """Test that function calls are not parsed as MCP calls.""" + func_message = Message.from_role_and_content(Role.ASSISTANT, '{"arg": "value"}') + func_message = func_message.with_recipient("functions.my_tool") + func_message = func_message.with_channel("commentary") + + func_items = harmony_to_response_output(func_message) + + assert len(func_items) == 1 + assert not isinstance(func_items[0], McpCall) + assert func_items[0].type == "function_call" + + +def test_mcp_vs_builtin_tools() -> None: + """Test that built-in tools (python, container) are not parsed as MCP calls.""" + # Test python (built-in tool) - should be reasoning, not MCP + python_message = Message.from_role_and_content(Role.ASSISTANT, "print('hello')") + python_message = python_message.with_recipient("python") + python_message = python_message.with_channel("commentary") + + python_items = harmony_to_response_output(python_message) + + assert len(python_items) == 1 + assert not isinstance(python_items[0], McpCall) + assert python_items[0].type == "reasoning" + + +def test_parser_state_to_response_output_commentary_channel() -> None: + """Test parser_state_to_response_output with commentary + channel and various recipients.""" + from unittest.mock import Mock + + # Test 1: functions.* recipient -> should return function tool call + parser_func = Mock() + parser_func.current_content = '{"arg": "value"}' + parser_func.current_role = Role.ASSISTANT + parser_func.current_channel = "commentary" + parser_func.current_recipient = "functions.my_tool" + + func_items = parser_state_to_response_output(parser_func) + + assert len(func_items) == 1 + assert not isinstance(func_items[0], McpCall) + assert func_items[0].type == "function_call" + assert func_items[0].name == "my_tool" + assert func_items[0].status == "in_progress" + + # Test 2: MCP tool (not builtin) -> should return MCP call + parser_mcp = Mock() + parser_mcp.current_content = '{"path": "/tmp"}' + parser_mcp.current_role = Role.ASSISTANT + parser_mcp.current_channel = "commentary" + parser_mcp.current_recipient = "filesystem" + + mcp_items = parser_state_to_response_output(parser_mcp) + + assert len(mcp_items) == 1 + assert isinstance(mcp_items[0], McpCall) + assert mcp_items[0].type == "mcp_call" + assert mcp_items[0].name == "filesystem" + assert mcp_items[0].server_label == "filesystem" + assert mcp_items[0].status == "in_progress" + + # Test 3: Built-in tool (python) + # should NOT return MCP call, returns reasoning (internal tool interaction) + parser_builtin = Mock() + parser_builtin.current_content = "print('hello')" + parser_builtin.current_role = Role.ASSISTANT + parser_builtin.current_channel = "commentary" + parser_builtin.current_recipient = "python" + + builtin_items = parser_state_to_response_output(parser_builtin) + + # Built-in tools explicitly return reasoning + assert len(builtin_items) == 1 + assert not isinstance(builtin_items[0], McpCall) + assert builtin_items[0].type == "reasoning" + + # Test 4: No recipient (preamble) → should return message, not reasoning + parser_preamble = Mock() + parser_preamble.current_content = "I'll search for that information now." + parser_preamble.current_role = Role.ASSISTANT + parser_preamble.current_channel = "commentary" + parser_preamble.current_recipient = None + + preamble_items = parser_state_to_response_output(parser_preamble) + + assert len(preamble_items) == 1 + assert isinstance(preamble_items[0], ResponseOutputMessage) + assert preamble_items[0].type == "message" + assert preamble_items[0].content[0].text == "I'll search for that information now." + assert preamble_items[0].status == "incomplete" # streaming + + +def test_parser_state_to_response_output_analysis_channel() -> None: + """Test parser_state_to_response_output with analysis + channel and various recipients.""" + from unittest.mock import Mock + + # Test 1: functions.* recipient -> should return function tool call + parser_func = Mock() + parser_func.current_content = '{"arg": "value"}' + parser_func.current_role = Role.ASSISTANT + parser_func.current_channel = "analysis" + parser_func.current_recipient = "functions.my_tool" + + func_items = parser_state_to_response_output(parser_func) + + assert len(func_items) == 1 + assert not isinstance(func_items[0], McpCall) + assert func_items[0].type == "function_call" + assert func_items[0].name == "my_tool" + assert func_items[0].status == "in_progress" + + # Test 2: MCP tool (not builtin) -> should return MCP call + parser_mcp = Mock() + parser_mcp.current_content = '{"query": "test"}' + parser_mcp.current_role = Role.ASSISTANT + parser_mcp.current_channel = "analysis" + parser_mcp.current_recipient = "database" + + mcp_items = parser_state_to_response_output(parser_mcp) + + assert len(mcp_items) == 1 + assert isinstance(mcp_items[0], McpCall) + assert mcp_items[0].type == "mcp_call" + assert mcp_items[0].name == "database" + assert mcp_items[0].server_label == "database" + assert mcp_items[0].status == "in_progress" + + # Test 3: Built-in tool (container) + # should NOT return MCP call, falls through to reasoning + parser_builtin = Mock() + parser_builtin.current_content = "docker run" + parser_builtin.current_role = Role.ASSISTANT + parser_builtin.current_channel = "analysis" + parser_builtin.current_recipient = "container" + + builtin_items = parser_state_to_response_output(parser_builtin) + + # Should fall through to reasoning logic + assert len(builtin_items) == 1 + assert not isinstance(builtin_items[0], McpCall) + assert builtin_items[0].type == "reasoning" diff --git a/tests/entrypoints/openai/responses/test_mcp_tools.py b/tests/entrypoints/openai/responses/test_mcp_tools.py index 310af4308c6..55445f1889b 100644 --- a/tests/entrypoints/openai/responses/test_mcp_tools.py +++ b/tests/entrypoints/openai/responses/test_mcp_tools.py @@ -97,16 +97,16 @@ class TestMCPToolServerUnit: assert server.get_tool_description("test_server", allowed_tools=[]) is None def test_builtin_tools_consistency(self): - """MCP_BUILTIN_TOOLS must match _BUILTIN_TOOL_TO_MCP_SERVER_LABEL values.""" + """MCP_BUILTIN_TOOLS must match BUILTIN_TOOL_TO_MCP_SERVER_LABEL values.""" from vllm.entrypoints.openai.parser.harmony_utils import ( - _BUILTIN_TOOL_TO_MCP_SERVER_LABEL, + BUILTIN_TOOL_TO_MCP_SERVER_LABEL, MCP_BUILTIN_TOOLS, ) - assert set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) == MCP_BUILTIN_TOOLS, ( + assert set(BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) == MCP_BUILTIN_TOOLS, ( f"MCP_BUILTIN_TOOLS {MCP_BUILTIN_TOOLS} does not match " - f"_BUILTIN_TOOL_TO_MCP_SERVER_LABEL values " - f"{set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())}" + f"BUILTIN_TOOL_TO_MCP_SERVER_LABEL values " + f"{set(BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values())}" ) diff --git a/vllm/entrypoints/openai/parser/harmony_utils.py b/vllm/entrypoints/openai/parser/harmony_utils.py index 9dfd5f518f7..9b4264456c5 100644 --- a/vllm/entrypoints/openai/parser/harmony_utils.py +++ b/vllm/entrypoints/openai/parser/harmony_utils.py @@ -2,27 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import datetime -import json from collections.abc import Iterable, Sequence from typing import Literal -from openai.types.responses import ( - ResponseFunctionToolCall, - ResponseOutputItem, - ResponseOutputMessage, - ResponseOutputText, - ResponseReasoningItem, -) -from openai.types.responses.response_function_web_search import ( - ActionFind, - ActionOpenPage, - ActionSearch, - ResponseFunctionWebSearch, -) -from openai.types.responses.response_output_item import McpCall -from openai.types.responses.response_reasoning_item import ( - Content as ResponseReasoningTextContent, -) from openai.types.responses.tool import Tool from openai_harmony import ( Author, @@ -38,17 +20,10 @@ from openai_harmony import ( ToolDescription, load_harmony_encoding, ) -from openai_harmony import Message as OpenAIHarmonyMessage -from openai_harmony import Role as OpenAIHarmonyRole from vllm import envs from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionToolsParam -from vllm.entrypoints.openai.responses.protocol import ( - ResponseInputOutputItem, - ResponsesRequest, -) from vllm.logger import init_logger -from vllm.utils import random_uuid logger = init_logger(__name__) @@ -64,14 +39,14 @@ _harmony_encoding = None # they are available and requested by the user. # Tool args are provided by MCP tool descriptions. Output # of the tools are stringified. -_BUILTIN_TOOL_TO_MCP_SERVER_LABEL: dict[str, str] = { +BUILTIN_TOOL_TO_MCP_SERVER_LABEL: dict[str, str] = { "python": "code_interpreter", "browser": "web_search_preview", "container": "container", } # Derive MCP_BUILTIN_TOOLS from the canonical mapping -MCP_BUILTIN_TOOLS: set[str] = set(_BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) +MCP_BUILTIN_TOOLS: set[str] = set(BUILTIN_TOOL_TO_MCP_SERVER_LABEL.values()) def has_custom_tools(tool_types: set[str]) -> bool: @@ -179,55 +154,6 @@ def get_user_message(content: str) -> Message: return Message.from_role_and_content(Role.USER, content) -def parse_response_input( - response_msg: ResponseInputOutputItem, - prev_responses: list[ResponseOutputItem | ResponseReasoningItem], -) -> Message: - if not isinstance(response_msg, dict): - response_msg = response_msg.model_dump() - if "type" not in response_msg or response_msg["type"] == "message": - role = response_msg["role"] - content = response_msg["content"] - # Add prefix for developer messages. - # <|start|>developer<|message|># Instructions {instructions}<|end|> - text_prefix = "Instructions:\n" if role == "developer" else "" - if isinstance(content, str): - msg = Message.from_role_and_content(role, text_prefix + content) - else: - contents = [TextContent(text=text_prefix + c["text"]) for c in content] - msg = Message.from_role_and_contents(role, contents) - if role == "assistant": - msg = msg.with_channel("final") - elif response_msg["type"] == "function_call_output": - call_id = response_msg["call_id"] - call_response: ResponseFunctionToolCall | None = None - for prev_response in reversed(prev_responses): - if ( - isinstance(prev_response, ResponseFunctionToolCall) - and prev_response.call_id == call_id - ): - call_response = prev_response - break - if call_response is None: - raise ValueError(f"No call message found for {call_id}") - msg = Message.from_author_and_content( - Author.new(Role.TOOL, f"functions.{call_response.name}"), - response_msg["output"], - ) - elif response_msg["type"] == "reasoning": - content = response_msg["content"] - assert len(content) == 1 - msg = Message.from_role_and_content(Role.ASSISTANT, content[0]["text"]) - elif response_msg["type"] == "function_call": - msg = Message.from_role_and_content(Role.ASSISTANT, response_msg["arguments"]) - msg = msg.with_channel("commentary") - msg = msg.with_recipient(f"functions.{response_msg['name']}") - msg = msg.with_content_type("json") - else: - raise ValueError(f"Unknown input type: {response_msg['type']}") - return msg - - def parse_chat_inputs_to_harmony_messages(chat_msgs: list) -> list[Message]: """ Parse a list of messages from request.messages in the Chat Completion API to @@ -390,139 +316,6 @@ def parse_chat_input_to_harmony_message( return msgs -def parse_input_to_harmony_message(chat_msg) -> list[Message]: - """Parse a message from request.previous_input_messages - into Harmony messages. - - Supports both OpenAI chat format ({"role": "..."}) and - Harmony format ({"author": {"role": "..."}}). - """ - if not isinstance(chat_msg, dict): - chat_msg = chat_msg.model_dump(exclude_none=True) - - if "author" in chat_msg and isinstance(chat_msg.get("author"), dict): - return [_parse_harmony_format_message(chat_msg)] - - return _parse_chat_format_message(chat_msg) - - -def _parse_harmony_format_message(chat_msg: dict) -> Message: - """Reconstruct a Message from Harmony-format dict, - preserving channel, recipient, and content_type.""" - author_dict = chat_msg["author"] - role = author_dict.get("role") - name = author_dict.get("name") - - raw_content = chat_msg.get("content", "") - if isinstance(raw_content, list): - # TODO: Support refusal and non-text content types. - contents = [TextContent(text=c.get("text", "")) for c in raw_content] - elif isinstance(raw_content, str): - contents = [TextContent(text=raw_content)] - else: - contents = [TextContent(text="")] - - if name: - msg = Message.from_author_and_contents(Author.new(Role(role), name), contents) - else: - msg = Message.from_role_and_contents(Role(role), contents) - - channel = chat_msg.get("channel") - if channel: - msg = msg.with_channel(channel) - recipient = chat_msg.get("recipient") - if recipient: - msg = msg.with_recipient(recipient) - content_type = chat_msg.get("content_type") - if content_type: - msg = msg.with_content_type(content_type) - - return msg - - -def _parse_chat_format_message(chat_msg: dict) -> list[Message]: - """Parse an OpenAI chat-format dict into Harmony messages.""" - role = chat_msg.get("role") - if role is None: - raise ValueError(f"Message has no 'role' key: {chat_msg}") - - # Assistant message with tool calls - tool_calls = chat_msg.get("tool_calls") - if role == "assistant" and tool_calls: - msgs: list[Message] = [] - for call in tool_calls: - func = call.get("function", {}) - name = func.get("name", "") - arguments = func.get("arguments", "") or "" - msg = Message.from_role_and_content(Role.ASSISTANT, arguments) - msg = msg.with_channel("commentary") - msg = msg.with_recipient(f"functions.{name}") - msg = msg.with_content_type("json") - msgs.append(msg) - return msgs - - # Tool role message (tool output) - if role == "tool": - name = chat_msg.get("name", "") - if name and not name.startswith("functions."): - name = f"functions.{name}" - content = chat_msg.get("content", "") or "" - content = flatten_chat_text_content(content) - # NOTE: .with_recipient("assistant") is required on tool messages - # to match parse_chat_input_to_harmony_message behavior and ensure - # proper routing in the Harmony protocol. - msg = ( - Message.from_author_and_content(Author.new(Role.TOOL, name), content) - .with_channel("commentary") - .with_recipient("assistant") - ) - return [msg] - - # Default: user/assistant/system messages - content = chat_msg.get("content", "") - if isinstance(content, str): - contents = [TextContent(text=content)] - else: - # TODO: Support refusal. - contents = [TextContent(text=c.get("text", "")) for c in content] - msg = Message.from_role_and_contents(role, contents) - return [msg] - - -def construct_harmony_previous_input_messages( - request: ResponsesRequest, -) -> list[OpenAIHarmonyMessage]: - messages: list[OpenAIHarmonyMessage] = [] - if request.previous_input_messages: - for message in request.previous_input_messages: - # Handle both OpenAIHarmonyMessage objects and dictionary inputs - if isinstance(message, OpenAIHarmonyMessage): - message_role = message.author.role - # To match OpenAI, instructions, reasoning and tools are - # always taken from the most recent Responses API request - # not carried over from previous requests - if ( - message_role == OpenAIHarmonyRole.SYSTEM - or message_role == OpenAIHarmonyRole.DEVELOPER - ): - continue - messages.append(message) - else: - harmony_messages = parse_input_to_harmony_message(message) - for harmony_msg in harmony_messages: - message_role = harmony_msg.author.role - # To match OpenAI, instructions, reasoning and tools are - # always taken from the most recent Responses API request - # not carried over from previous requests - if ( - message_role == OpenAIHarmonyRole.SYSTEM - or message_role == OpenAIHarmonyRole.DEVELOPER - ): - continue - messages.append(harmony_msg) - return messages - - def render_for_completion(messages: list[Message]) -> list[int]: conversation = Conversation.from_messages(messages) token_ids = get_encoding().render_conversation_for_completion( @@ -531,313 +324,6 @@ def render_for_completion(messages: list[Message]) -> list[int]: return token_ids -def _parse_browser_tool_call(message: Message, recipient: str) -> ResponseOutputItem: - """Parse browser tool calls (search, open, find) into web search items.""" - if len(message.content) != 1: - raise ValueError("Invalid number of contents in browser message") - content = message.content[0] - - # Parse JSON args (with retry detection) - try: - browser_call = json.loads(content.text) - except json.JSONDecodeError: - logger.warning( - "Invalid JSON in browser tool call, using error placeholder: %s", - content.text, - ) - json_retry_output_message = ( - f"Invalid JSON args, caught and retried: {content.text}" - ) - browser_call = { - "query": json_retry_output_message, - "url": json_retry_output_message, - "pattern": json_retry_output_message, - } - - # Create appropriate action based on recipient - if recipient == "browser.search": - action = ActionSearch( - query=f"cursor:{browser_call.get('query', '')}", type="search" - ) - elif recipient == "browser.open": - action = ActionOpenPage( - url=f"cursor:{browser_call.get('url', '')}", type="open_page" - ) - elif recipient == "browser.find": - action = ActionFind( - pattern=browser_call.get("pattern", ""), - url=f"cursor:{browser_call.get('url', '')}", - type="find", - ) - else: - raise ValueError(f"Unknown browser action: {recipient}") - - return ResponseFunctionWebSearch( - id=f"ws_{random_uuid()}", - action=action, - status="completed", - type="web_search_call", - ) - - -def _parse_function_call(message: Message, recipient: str) -> list[ResponseOutputItem]: - """Parse function calls into function tool call items.""" - function_name = recipient.split(".")[-1] - output_items = [] - for content in message.content: - random_id = random_uuid() - response_item = ResponseFunctionToolCall( - arguments=content.text, - call_id=f"call_{random_id}", - type="function_call", - name=function_name, - id=f"fc_{random_id}", - ) - output_items.append(response_item) - return output_items - - -def _parse_reasoning(message: Message) -> list[ResponseOutputItem]: - """Parse reasoning/analysis content into reasoning items.""" - output_items = [] - for content in message.content: - reasoning_item = ResponseReasoningItem( - id=f"rs_{random_uuid()}", - summary=[], - type="reasoning", - content=[ - ResponseReasoningTextContent(text=content.text, type="reasoning_text") - ], - status=None, - ) - output_items.append(reasoning_item) - return output_items - - -def _parse_final_message(message: Message) -> ResponseOutputItem: - """Parse final channel messages into output message items.""" - contents = [] - for content in message.content: - output_text = ResponseOutputText( - text=content.text, - annotations=[], # TODO - type="output_text", - logprobs=None, # TODO - ) - contents.append(output_text) - return ResponseOutputMessage( - id=f"msg_{random_uuid()}", - content=contents, - role=message.author.role, - status="completed", - type="message", - ) - - -def _parse_mcp_recipient(recipient: str) -> tuple[str, str]: - """ - Parse MCP recipient into (server_label, tool_name). - - For dotted recipients like "repo_browser.list": - - server_label: "repo_browser" (namespace/server) - - tool_name: "list" (specific tool) - - For simple recipients like "filesystem": - - server_label: "filesystem" - - tool_name: "filesystem" - """ - if "." in recipient: - server_label = recipient.split(".")[0] - tool_name = recipient.split(".")[-1] - else: - server_label = recipient - tool_name = recipient - return server_label, tool_name - - -def _parse_mcp_call(message: Message, recipient: str) -> list[ResponseOutputItem]: - """Parse MCP calls into MCP call items.""" - # Handle built-in tools that need server_label mapping - if recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: - server_label = _BUILTIN_TOOL_TO_MCP_SERVER_LABEL[recipient] - tool_name = recipient - else: - server_label, tool_name = _parse_mcp_recipient(recipient) - - output_items = [] - for content in message.content: - response_item = McpCall( - arguments=content.text, - type="mcp_call", - name=tool_name, - server_label=server_label, - id=f"mcp_{random_uuid()}", - status="completed", - ) - output_items.append(response_item) - return output_items - - -def _parse_message_no_recipient( - message: Message, -) -> list[ResponseOutputItem]: - """Parse a Harmony message with no recipient based on its channel.""" - if message.channel == "analysis": - return _parse_reasoning(message) - - if message.channel in ("commentary", "final"): - # Per Harmony format, preambles (commentary with no recipient) and - # final channel content are both intended to be shown to end-users. - # See: https://cookbook.openai.com/articles/openai-harmony - return [_parse_final_message(message)] - - raise ValueError(f"Unknown channel: {message.channel}") - - -def parse_output_message(message: Message) -> list[ResponseOutputItem]: - """ - Parse a Harmony message into a list of output response items. - """ - if message.author.role != "assistant": - # This is a message from a tool to the assistant (e.g., search result). - # Don't include it in the final output for now. This aligns with - # OpenAI's behavior on models like o4-mini. - return [] - - output_items: list[ResponseOutputItem] = [] - recipient = message.recipient - - if recipient is not None: - # Browser tool calls (browser.search, browser.open, browser.find) - if recipient.startswith("browser."): - output_items.append(_parse_browser_tool_call(message, recipient)) - - # Function calls (should only happen on commentary channel) - elif message.channel == "commentary" and recipient.startswith("functions."): - output_items.extend(_parse_function_call(message, recipient)) - - # Built-in MCP tools (python, browser, container) - elif recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: - output_items.extend(_parse_reasoning(message)) - - # All other recipients are MCP calls - else: - output_items.extend(_parse_mcp_call(message, recipient)) - - # No recipient - handle based on channel for non-tool messages - else: - output_items.extend(_parse_message_no_recipient(message)) - - return output_items - - -def parse_remaining_state(parser: StreamableParser) -> list[ResponseOutputItem]: - if not parser.current_content: - return [] - if parser.current_role != Role.ASSISTANT: - return [] - current_recipient = parser.current_recipient - if current_recipient is not None and current_recipient.startswith("browser."): - return [] - - if current_recipient and parser.current_channel in ("commentary", "analysis"): - if current_recipient.startswith("functions."): - rid = random_uuid() - return [ - ResponseFunctionToolCall( - arguments=parser.current_content, - call_id=f"call_{rid}", - type="function_call", - name=current_recipient.split(".")[-1], - id=f"fc_{rid}", - status="in_progress", - ) - ] - # Built-in MCP tools (python, browser, container) - elif current_recipient in _BUILTIN_TOOL_TO_MCP_SERVER_LABEL: - return [ - ResponseReasoningItem( - id=f"rs_{random_uuid()}", - summary=[], - type="reasoning", - content=[ - ResponseReasoningTextContent( - text=parser.current_content, type="reasoning_text" - ) - ], - status=None, - ) - ] - # All other recipients are MCP calls - else: - rid = random_uuid() - server_label, tool_name = _parse_mcp_recipient(current_recipient) - return [ - McpCall( - arguments=parser.current_content, - type="mcp_call", - name=tool_name, - server_label=server_label, - id=f"mcp_{rid}", - status="in_progress", - ) - ] - - if parser.current_channel == "commentary": - # Per Harmony format, preambles (commentary with no recipient) are - # intended to be shown to end-users, unlike analysis channel content. - output_text = ResponseOutputText( - text=parser.current_content, - annotations=[], - type="output_text", - logprobs=None, - ) - return [ - ResponseOutputMessage( - id=f"msg_{random_uuid()}", - content=[output_text], - role="assistant", - status="incomplete", - type="message", - ) - ] - - if parser.current_channel == "analysis": - return [ - ResponseReasoningItem( - id=f"rs_{random_uuid()}", - summary=[], - type="reasoning", - content=[ - ResponseReasoningTextContent( - text=parser.current_content, type="reasoning_text" - ) - ], - status=None, - ) - ] - - if parser.current_channel == "final": - output_text = ResponseOutputText( - text=parser.current_content, - annotations=[], # TODO - type="output_text", - logprobs=None, # TODO - ) - text_item = ResponseOutputMessage( - id=f"msg_{random_uuid()}", - content=[output_text], - role="assistant", - # if the parser still has messages (ie if the generator got cut - # abruptly), this should be incomplete - status="incomplete", - type="message", - ) - return [text_item] - - return [] - - def get_stop_tokens_for_assistant_actions() -> list[int]: return get_encoding().stop_tokens_for_assistant_actions() diff --git a/vllm/entrypoints/openai/responses/harmony.py b/vllm/entrypoints/openai/responses/harmony.py new file mode 100644 index 00000000000..460f310926a --- /dev/null +++ b/vllm/entrypoints/openai/responses/harmony.py @@ -0,0 +1,552 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Harmony ↔ Responses API conversion utilities. + +Handles two directions: + 1. Response Input → Harmony Messages (input parsing) + 2. Harmony Messages → Response Output Items (output parsing) +""" + +import json + +from openai.types.responses import ( + ResponseFunctionToolCall, + ResponseOutputItem, + ResponseOutputMessage, + ResponseOutputText, + ResponseReasoningItem, +) +from openai.types.responses.response_function_web_search import ( + ActionFind, + ActionOpenPage, + ActionSearch, + ResponseFunctionWebSearch, +) +from openai.types.responses.response_output_item import McpCall +from openai.types.responses.response_reasoning_item import ( + Content as ResponseReasoningTextContent, +) +from openai_harmony import Author, Message, Role, StreamableParser, TextContent + +from vllm.entrypoints.openai.parser.harmony_utils import ( + BUILTIN_TOOL_TO_MCP_SERVER_LABEL, + flatten_chat_text_content, +) +from vllm.entrypoints.openai.responses.protocol import ( + ResponseInputOutputItem, + ResponsesRequest, +) +from vllm.logger import init_logger +from vllm.utils import random_uuid + +logger = init_logger(__name__) + +# --------------------------------------------------------------------------- +# 1. Private helpers for input parsing +# --------------------------------------------------------------------------- + + +def _parse_harmony_format_message(chat_msg: dict) -> Message: + """Reconstruct a Message from Harmony-format dict, + preserving channel, recipient, and content_type.""" + author_dict = chat_msg["author"] + role = author_dict.get("role") + name = author_dict.get("name") + + raw_content = chat_msg.get("content", "") + if isinstance(raw_content, list): + # TODO: Support refusal and non-text content types. + contents = [TextContent(text=c.get("text", "")) for c in raw_content] + elif isinstance(raw_content, str): + contents = [TextContent(text=raw_content)] + else: + contents = [TextContent(text="")] + + if name: + msg = Message.from_author_and_contents(Author.new(Role(role), name), contents) + else: + msg = Message.from_role_and_contents(Role(role), contents) + + channel = chat_msg.get("channel") + if channel: + msg = msg.with_channel(channel) + recipient = chat_msg.get("recipient") + if recipient: + msg = msg.with_recipient(recipient) + content_type = chat_msg.get("content_type") + if content_type: + msg = msg.with_content_type(content_type) + + return msg + + +def _parse_chat_format_message(chat_msg: dict) -> list[Message]: + """Parse an OpenAI chat-format dict into Harmony messages.""" + role = chat_msg.get("role") + if role is None: + raise ValueError(f"Message has no 'role' key: {chat_msg}") + + # Assistant message with tool calls + tool_calls = chat_msg.get("tool_calls") + if role == "assistant" and tool_calls: + msgs: list[Message] = [] + for call in tool_calls: + func = call.get("function", {}) + name = func.get("name", "") + arguments = func.get("arguments", "") or "" + msg = Message.from_role_and_content(Role.ASSISTANT, arguments) + msg = msg.with_channel("commentary") + msg = msg.with_recipient(f"functions.{name}") + msg = msg.with_content_type("json") + msgs.append(msg) + return msgs + + # Tool role message (tool output) + if role == "tool": + name = chat_msg.get("name", "") + if name and not name.startswith("functions."): + name = f"functions.{name}" + content = chat_msg.get("content", "") or "" + content = flatten_chat_text_content(content) + # NOTE: .with_recipient("assistant") is required on tool messages + # to match parse_chat_input_to_harmony_message behavior and ensure + # proper routing in the Harmony protocol. + msg = ( + Message.from_author_and_content(Author.new(Role.TOOL, name), content) + .with_channel("commentary") + .with_recipient("assistant") + ) + return [msg] + + # Default: user/assistant/system messages + content = chat_msg.get("content", "") + if isinstance(content, str): + contents = [TextContent(text=content)] + else: + # TODO: Support refusal. + contents = [TextContent(text=c.get("text", "")) for c in content] + msg = Message.from_role_and_contents(role, contents) + return [msg] + + +# --------------------------------------------------------------------------- +# 2. Public input parsing functions +# --------------------------------------------------------------------------- + + +def response_input_to_harmony( + response_msg: ResponseInputOutputItem, + prev_responses: list[ResponseOutputItem | ResponseReasoningItem], +) -> Message: + """Convert a single ResponseInputOutputItem into a Harmony Message.""" + if not isinstance(response_msg, dict): + response_msg = response_msg.model_dump() + if "type" not in response_msg or response_msg["type"] == "message": + role = response_msg["role"] + content = response_msg["content"] + # Add prefix for developer messages. + # <|start|>developer<|message|># Instructions {instructions}<|end|> + text_prefix = "Instructions:\n" if role == "developer" else "" + if isinstance(content, str): + msg = Message.from_role_and_content(role, text_prefix + content) + else: + contents = [TextContent(text=text_prefix + c["text"]) for c in content] + msg = Message.from_role_and_contents(role, contents) + if role == "assistant": + msg = msg.with_channel("final") + elif response_msg["type"] == "function_call_output": + call_id = response_msg["call_id"] + call_response: ResponseFunctionToolCall | None = None + for prev_response in reversed(prev_responses): + if ( + isinstance(prev_response, ResponseFunctionToolCall) + and prev_response.call_id == call_id + ): + call_response = prev_response + break + if call_response is None: + raise ValueError(f"No call message found for {call_id}") + msg = Message.from_author_and_content( + Author.new(Role.TOOL, f"functions.{call_response.name}"), + response_msg["output"], + ) + elif response_msg["type"] == "reasoning": + content = response_msg["content"] + assert len(content) == 1 + msg = Message.from_role_and_content(Role.ASSISTANT, content[0]["text"]) + elif response_msg["type"] == "function_call": + msg = Message.from_role_and_content(Role.ASSISTANT, response_msg["arguments"]) + msg = msg.with_channel("commentary") + msg = msg.with_recipient(f"functions.{response_msg['name']}") + msg = msg.with_content_type("json") + else: + raise ValueError(f"Unknown input type: {response_msg['type']}") + return msg + + +def response_previous_input_to_harmony(chat_msg) -> list[Message]: + """Parse a message from request.previous_input_messages + into Harmony messages. + + Supports both OpenAI chat format ({"role": "..."}) and + Harmony format ({"author": {"role": "..."}}). + """ + if not isinstance(chat_msg, dict): + chat_msg = chat_msg.model_dump(exclude_none=True) + + if "author" in chat_msg and isinstance(chat_msg.get("author"), dict): + return [_parse_harmony_format_message(chat_msg)] + + return _parse_chat_format_message(chat_msg) + + +def construct_harmony_previous_input_messages( + request: ResponsesRequest, +) -> list[Message]: + """Build a Harmony message list from request.previous_input_messages. + + Filters out system/developer messages to match OpenAI behavior where + instructions are always taken from the most recent Responses API request. + """ + messages: list[Message] = [] + if request.previous_input_messages: + for message in request.previous_input_messages: + # Handle both Message objects and dictionary inputs + if isinstance(message, Message): + message_role = message.author.role + if message_role == Role.SYSTEM or message_role == Role.DEVELOPER: + continue + messages.append(message) + else: + harmony_messages = response_previous_input_to_harmony(message) + for harmony_msg in harmony_messages: + message_role = harmony_msg.author.role + if message_role == Role.SYSTEM or message_role == Role.DEVELOPER: + continue + messages.append(harmony_msg) + return messages + + +# --------------------------------------------------------------------------- +# 3. Private helpers for output parsing +# --------------------------------------------------------------------------- + + +def _parse_browser_tool_call(message: Message, recipient: str) -> ResponseOutputItem: + """Parse browser tool calls (search, open, find) into web search items.""" + if len(message.content) != 1: + raise ValueError("Invalid number of contents in browser message") + content = message.content[0] + + # Parse JSON args (with retry detection) + try: + browser_call = json.loads(content.text) + except json.JSONDecodeError: + logger.warning( + "Invalid JSON in browser tool call, using error placeholder: %s", + content.text, + ) + json_retry_output_message = ( + f"Invalid JSON args, caught and retried: {content.text}" + ) + browser_call = { + "query": json_retry_output_message, + "url": json_retry_output_message, + "pattern": json_retry_output_message, + } + + # Create appropriate action based on recipient + if recipient == "browser.search": + action = ActionSearch( + query=f"cursor:{browser_call.get('query', '')}", type="search" + ) + elif recipient == "browser.open": + action = ActionOpenPage( + url=f"cursor:{browser_call.get('url', '')}", type="open_page" + ) + elif recipient == "browser.find": + action = ActionFind( + pattern=browser_call.get("pattern", ""), + url=f"cursor:{browser_call.get('url', '')}", + type="find", + ) + else: + raise ValueError(f"Unknown browser action: {recipient}") + + return ResponseFunctionWebSearch( + id=f"ws_{random_uuid()}", + action=action, + status="completed", + type="web_search_call", + ) + + +def _parse_function_call(message: Message, recipient: str) -> list[ResponseOutputItem]: + """Parse function calls into function tool call items.""" + function_name = recipient.split(".")[-1] + output_items = [] + for content in message.content: + random_id = random_uuid() + response_item = ResponseFunctionToolCall( + arguments=content.text, + call_id=f"call_{random_id}", + type="function_call", + name=function_name, + id=f"fc_{random_id}", + ) + output_items.append(response_item) + return output_items + + +def _parse_reasoning(message: Message) -> list[ResponseOutputItem]: + """Parse reasoning/analysis content into reasoning items.""" + output_items = [] + for content in message.content: + reasoning_item = ResponseReasoningItem( + id=f"rs_{random_uuid()}", + summary=[], + type="reasoning", + content=[ + ResponseReasoningTextContent(text=content.text, type="reasoning_text") + ], + status=None, + ) + output_items.append(reasoning_item) + return output_items + + +def _parse_final_message(message: Message) -> ResponseOutputItem: + """Parse final channel messages into output message items.""" + contents = [] + for content in message.content: + output_text = ResponseOutputText( + text=content.text, + annotations=[], # TODO + type="output_text", + logprobs=None, # TODO + ) + contents.append(output_text) + return ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=contents, + role=message.author.role, + status="completed", + type="message", + ) + + +def _parse_mcp_recipient(recipient: str) -> tuple[str, str]: + """Parse MCP recipient into (server_label, tool_name). + + For dotted recipients like "repo_browser.list": + - server_label: "repo_browser" (namespace/server) + - tool_name: "list" (specific tool) + + For simple recipients like "filesystem": + - server_label: "filesystem" + - tool_name: "filesystem" + """ + if "." in recipient: + server_label = recipient.split(".")[0] + tool_name = recipient.split(".")[-1] + else: + server_label = recipient + tool_name = recipient + return server_label, tool_name + + +def _parse_mcp_call(message: Message, recipient: str) -> list[ResponseOutputItem]: + """Parse MCP calls into MCP call items.""" + # Handle built-in tools that need server_label mapping + if recipient in BUILTIN_TOOL_TO_MCP_SERVER_LABEL: + server_label = BUILTIN_TOOL_TO_MCP_SERVER_LABEL[recipient] + tool_name = recipient + else: + server_label, tool_name = _parse_mcp_recipient(recipient) + + output_items = [] + for content in message.content: + response_item = McpCall( + arguments=content.text, + type="mcp_call", + name=tool_name, + server_label=server_label, + id=f"mcp_{random_uuid()}", + status="completed", + ) + output_items.append(response_item) + return output_items + + +def _parse_message_no_recipient( + message: Message, +) -> list[ResponseOutputItem]: + """Parse a Harmony message with no recipient based on its channel.""" + if message.channel == "analysis": + return _parse_reasoning(message) + + if message.channel in ("commentary", "final"): + # Per Harmony format, preambles (commentary with no recipient) and + # final channel content are both intended to be shown to end-users. + # See: https://cookbook.openai.com/articles/openai-harmony + return [_parse_final_message(message)] + + raise ValueError(f"Unknown channel: {message.channel}") + + +# --------------------------------------------------------------------------- +# 4. Public output parsing functions +# --------------------------------------------------------------------------- + + +def harmony_to_response_output(message: Message) -> list[ResponseOutputItem]: + """Parse a Harmony message into a list of output response items. + + This is the main dispatcher that routes based on channel and recipient. + """ + if message.author.role != "assistant": + # This is a message from a tool to the assistant (e.g., search result). + # Don't include it in the final output for now. This aligns with + # OpenAI's behavior on models like o4-mini. + return [] + + output_items: list[ResponseOutputItem] = [] + recipient = message.recipient + + if recipient is not None: + # Browser tool calls (browser.search, browser.open, browser.find) + if recipient.startswith("browser."): + output_items.append(_parse_browser_tool_call(message, recipient)) + + # Function calls (should only happen on commentary channel) + elif message.channel == "commentary" and recipient.startswith("functions."): + output_items.extend(_parse_function_call(message, recipient)) + + # Built-in MCP tools (python, browser, container) + elif recipient in BUILTIN_TOOL_TO_MCP_SERVER_LABEL: + output_items.extend(_parse_reasoning(message)) + + # All other recipients are MCP calls + else: + output_items.extend(_parse_mcp_call(message, recipient)) + + # No recipient - handle based on channel for non-tool messages + else: + output_items.extend(_parse_message_no_recipient(message)) + + return output_items + + +def parser_state_to_response_output( + parser: StreamableParser, +) -> list[ResponseOutputItem]: + """Extract in-progress response items from incomplete parser state. + + Called when the parser has buffered content that hasn't formed a + complete message yet (e.g., generation was cut short). + """ + if not parser.current_content: + return [] + if parser.current_role != Role.ASSISTANT: + return [] + current_recipient = parser.current_recipient + if current_recipient is not None and current_recipient.startswith("browser."): + return [] + + if current_recipient and parser.current_channel in ("commentary", "analysis"): + if current_recipient.startswith("functions."): + rid = random_uuid() + return [ + ResponseFunctionToolCall( + arguments=parser.current_content, + call_id=f"call_{rid}", + type="function_call", + name=current_recipient.split(".")[-1], + id=f"fc_{rid}", + status="in_progress", + ) + ] + # Built-in MCP tools (python, browser, container) + elif current_recipient in BUILTIN_TOOL_TO_MCP_SERVER_LABEL: + return [ + ResponseReasoningItem( + id=f"rs_{random_uuid()}", + summary=[], + type="reasoning", + content=[ + ResponseReasoningTextContent( + text=parser.current_content, type="reasoning_text" + ) + ], + status=None, + ) + ] + # All other recipients are MCP calls + else: + rid = random_uuid() + server_label, tool_name = _parse_mcp_recipient(current_recipient) + return [ + McpCall( + arguments=parser.current_content, + type="mcp_call", + name=tool_name, + server_label=server_label, + id=f"mcp_{rid}", + status="in_progress", + ) + ] + + if parser.current_channel == "commentary": + # Per Harmony format, preambles (commentary with no recipient) are + # intended to be shown to end-users, unlike analysis channel content. + output_text = ResponseOutputText( + text=parser.current_content, + annotations=[], + type="output_text", + logprobs=None, + ) + return [ + ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=[output_text], + role="assistant", + status="incomplete", + type="message", + ) + ] + + if parser.current_channel == "analysis": + return [ + ResponseReasoningItem( + id=f"rs_{random_uuid()}", + summary=[], + type="reasoning", + content=[ + ResponseReasoningTextContent( + text=parser.current_content, type="reasoning_text" + ) + ], + status=None, + ) + ] + + if parser.current_channel == "final": + output_text = ResponseOutputText( + text=parser.current_content, + annotations=[], # TODO + type="output_text", + logprobs=None, # TODO + ) + text_item = ResponseOutputMessage( + id=f"msg_{random_uuid()}", + content=[output_text], + role="assistant", + # if the parser still has messages (ie if the generator got cut + # abruptly), this should be incomplete + status="incomplete", + type="message", + ) + return [text_item] + + return [] diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index c0ca87a9852..b9d526e25de 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -58,15 +58,11 @@ from vllm.entrypoints.openai.engine.serving import ( ) from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.openai.parser.harmony_utils import ( - construct_harmony_previous_input_messages, get_developer_message, get_stop_tokens_for_assistant_actions, get_system_message, get_user_message, has_custom_tools, - parse_output_message, - parse_remaining_state, - parse_response_input, render_for_completion, ) from vllm.entrypoints.openai.responses.context import ( @@ -76,6 +72,12 @@ from vllm.entrypoints.openai.responses.context import ( SimpleContext, StreamingHarmonyContext, ) +from vllm.entrypoints.openai.responses.harmony import ( + construct_harmony_previous_input_messages, + harmony_to_response_output, + parser_state_to_response_output, + response_input_to_harmony, +) from vllm.entrypoints.openai.responses.protocol import ( InputTokensDetails, OutputTokensDetails, @@ -954,9 +956,9 @@ class OpenAIServingResponses(OpenAIServing): output_items: list[ResponseOutputItem] = [] num_init_messages = context.num_init_messages for msg in context.messages[num_init_messages:]: - output_items.extend(parse_output_message(msg)) + output_items.extend(harmony_to_response_output(msg)) # Handle the generation stopped in the middle (if any). - last_items = parse_remaining_state(context.parser) + last_items = parser_state_to_response_output(context.parser) if last_items: output_items.extend(last_items) return output_items @@ -1103,13 +1105,13 @@ class OpenAIServingResponses(OpenAIServing): else: prev_outputs = [] for response_msg in request.input: - new_msg = parse_response_input(response_msg, prev_outputs) + new_msg = response_input_to_harmony(response_msg, prev_outputs) if new_msg.author.role != "system": messages.append(new_msg) # User passes in a tool call request and its output. We need - # to add the tool call request to prev_outputs so that the - # parse_response_input can find the tool call request when + # to add the tool call request to prev_outputs so that + # response_input_to_harmony can find the tool call request when # parsing the tool call output. if isinstance(response_msg, ResponseFunctionToolCall): prev_outputs.append(response_msg) From d3a51da92a031f6c1758771a2b13976ace2eece2 Mon Sep 17 00:00:00 2001 From: Cyrus Leung Date: Thu, 26 Feb 2026 14:35:41 +0800 Subject: [PATCH 25/43] [Benchmark] Simplify SLA scan (#35306) Signed-off-by: DarkLight1337 --- docs/benchmarking/cli.md | 5 + docs/benchmarking/sweeps.md | 88 ++--- tests/benchmarks/sweep/test_serve_sla.py | 298 ---------------- vllm/benchmarks/sweep/plot.py | 2 +- vllm/benchmarks/sweep/serve.py | 87 +++-- vllm/benchmarks/sweep/serve_sla.py | 431 +++++++---------------- vllm/benchmarks/sweep/sla_sweep.py | 138 -------- vllm/benchmarks/sweep/startup.py | 3 +- 8 files changed, 253 insertions(+), 799 deletions(-) delete mode 100644 tests/benchmarks/sweep/test_serve_sla.py delete mode 100644 vllm/benchmarks/sweep/sla_sweep.py diff --git a/docs/benchmarking/cli.md b/docs/benchmarking/cli.md index 7bb91239c58..8bbd9b0c0e3 100644 --- a/docs/benchmarking/cli.md +++ b/docs/benchmarking/cli.md @@ -4,6 +4,11 @@ This section guides you through running benchmark tests with the extensive datas It's a living document, updated as new features and datasets become available. +!!! tip + The benchmarks described on this page are mainly for evaluating specific vLLM features as well as regression testing. + + For benchmarking production vLLM servers, we recommend [GuideLLM](https://github.com/vllm-project/guidellm), an established performance benchmarking framework with live progress updates and automatic report generation. It is also more flexible than `vllm bench serve` in terms of dataset loading, request formatting, and workload patterns. + ## Dataset Overview