From 9a16d43ef0413b0a2a222e7f5baa4498c291ae5b Mon Sep 17 00:00:00 2001 From: Hexiang Wang <56632993+whx-sjtu@users.noreply.github.com> Date: Fri, 16 Oct 2026 17:54:48 +0800 Subject: [PATCH 01/49] [Bugfix][Mooncake] Fix MultiConnector stats and proxy SSE responses (#61218) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: whx-sjtu Signed-off-by: Nicolò Lucchesi --- .buildkite/test_areas/misc.yaml | 1 + .../mooncake_connector_proxy.py | 5 +- .../kv_connector/unit/test_mooncake_proxy.py | 47 +++++++++++++++++++ .../kv_connector/unit/test_multi_connector.py | 17 +++++++ .../kv_connector/v1/multi_connector.py | 8 ++-- 5 files changed, 72 insertions(+), 6 deletions(-) create mode 100644 tests/v1/kv_connector/unit/test_mooncake_proxy.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index 29e0582ec5..c79b451749 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -214,6 +214,7 @@ steps: - vllm/v1/ - tests/v1 - tests/watermarking + - examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py device: cpu-small commands: # split the test to avoid interference diff --git a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py index 6815f42061..bd63f04627 100644 --- a/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py +++ b/examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py @@ -340,7 +340,10 @@ async def _handle_completions(api: str, request: Request): ): yield chunk - return StreamingResponse(generate_stream(), media_type="application/json") + media_type = ( + "text/event-stream" if req_data.get("stream") else "application/json" + ) + return StreamingResponse(generate_stream(), media_type=media_type) except Exception as e: import sys diff --git a/tests/v1/kv_connector/unit/test_mooncake_proxy.py b/tests/v1/kv_connector/unit/test_mooncake_proxy.py new file mode 100644 index 0000000000..2d4d93a887 --- /dev/null +++ b/tests/v1/kv_connector/unit/test_mooncake_proxy.py @@ -0,0 +1,47 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import asyncio +import itertools +import runpy +from pathlib import Path +from unittest.mock import AsyncMock + +import httpx +import pytest + +pytestmark = pytest.mark.cpu_test + + +@pytest.mark.asyncio +@pytest.mark.parametrize("endpoint", ["/v1/completions", "/v1/chat/completions"]) +@pytest.mark.parametrize("stream", [False, True]) +async def test_response_media_type_matches_payload(monkeypatch, endpoint, stream): + """Clients must select SSE parsing for streaming completion responses.""" + path = ( + Path(__file__).parents[4] + / "examples/disaggregated/mooncake_connector/mooncake_connector_proxy.py" + ) + proxy = runpy.run_path(str(path))["_handle_completions"].__globals__ + app = proxy["app"] + app.state.ready = asyncio.Event() + app.state.ready.set() + app.state.prefill_iterator = itertools.cycle([({}, 0)]) + app.state.decode_iterator = itertools.cycle([0]) + app.state.decode_clients = [{}] + payload = b'{"choices":[{"text":"hello"}]}' + body = b"data: " + payload + b"\n\ndata: [DONE]\n\n" if stream else payload + + async def decode_response(*args, **kwargs): + yield body + + monkeypatch.setitem(proxy, "send_request_to_service", AsyncMock()) + monkeypatch.setitem(proxy, "stream_service_response", decode_response) + async with httpx.AsyncClient( + transport=httpx.ASGITransport(app=app), base_url="http://proxy" + ) as client: + response = await client.post(endpoint, json={"stream": stream}) + + expected_type = "text/event-stream" if stream else "application/json" + assert response.headers["content-type"].split(";")[0] == expected_type + assert response.content == body diff --git a/tests/v1/kv_connector/unit/test_multi_connector.py b/tests/v1/kv_connector/unit/test_multi_connector.py index 4ac3eaf7c5..af111483ef 100644 --- a/tests/v1/kv_connector/unit/test_multi_connector.py +++ b/tests/v1/kv_connector/unit/test_multi_connector.py @@ -805,6 +805,23 @@ class TestMultiConnectorStats: child_metrics.observe.assert_called_once_with({"transfer_duration": [1.5]}, 2) + @pytest.mark.parametrize("with_exporter", [False, True]) + def test_prom_metrics_skips_stats_only_connectors(self, with_exporter): + """Stats-only children must not prevent supported metrics from exporting.""" + child_metrics = MagicMock() + metrics = object.__new__(MultiKVConnectorPromMetrics) + metrics._prom_metrics = ( + {"MooncakeStoreConnector": child_metrics} if with_exporter else {} + ) + payload = {"MooncakeConnector": {"bytes_transferred": [1024]}} + if with_exporter: + payload["MooncakeStoreConnector"] = {"load_get_bytes": [2048]} + + metrics.observe(payload, engine_idx=2) + + if with_exporter: + child_metrics.observe.assert_called_once_with({"load_get_bytes": [2048]}, 2) + def test_aggregate_same_connector(self): """Test aggregating stats from the same connector type.""" stats1 = MultiKVConnectorStats( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py index bc59f3cb27..70f9454a62 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/multi_connector.py @@ -125,11 +125,9 @@ class MultiKVConnectorPromMetrics(KVConnectorPromMetrics): def observe(self, transfer_stats_data: dict[str, Any], engine_idx: int = 0): for connector_id, stats_data in transfer_stats_data.items(): - assert connector_id in self._prom_metrics, ( - f"{connector_id} is not contained in the list of registered connectors " - f"with Prometheus metrics support: {self._prom_metrics.keys()}" - ) - self._prom_metrics[connector_id].observe(stats_data, engine_idx) + prom_metrics = self._prom_metrics.get(connector_id) + if prom_metrics is not None: + prom_metrics.observe(stats_data, engine_idx) class MultiConnector(KVConnectorBase_V1, SupportsHMA): From 024c5188bf5a8a464124b36519c1e7de76648625 Mon Sep 17 00:00:00 2001 From: Shimi Bandiel Date: Fri, 16 Oct 2026 03:04:15 -0700 Subject: [PATCH 02/49] [Frontend] Resolve logprobs on the parsed streaming derender path (#60437) Signed-off-by: Shimi Bandiel Co-authored-by: Claude Co-authored-by: wang.yuqi Co-authored-by: Kevin H. Luu --- docs/serving/online_serving/derenderer.md | 4 +- .../scale_out/derender/test_derender.py | 44 +++++ .../derender/test_derender_parity.py | 163 +++++++++++++++++- .../derender/test_derender_stream.py | 153 ++++++++++++++++ tests/entrypoints/scale_out/derender/utils.py | 11 ++ vllm/renderers/online_derenderer.py | 49 ++++-- 6 files changed, 404 insertions(+), 20 deletions(-) diff --git a/docs/serving/online_serving/derenderer.md b/docs/serving/online_serving/derenderer.md index b7b5e5b8f0..afc7961dd2 100644 --- a/docs/serving/online_serving/derenderer.md +++ b/docs/serving/online_serving/derenderer.md @@ -105,7 +105,7 @@ The server keeps no state between calls. Everything the next call needs is in `s - **Send `prompt_token_ids` on the first chunk, even without a parser.** It seeds detokenization so the first token's leading space matches the coupled endpoint (see [Request format](#request-format)). Later chunks already carry that context in `stream_state`. - **Don't forward `[DONE]`.** It marks the end of the generate stream and isn't a chunk. -`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Generate returns every top-k candidate the engine produced; derender keeps the ones the coupled endpoint would, with the same helpers: the first `top_logprobs` on chat (all of them for `top_logprobs=-1`) and the first `logprobs + 1` on completions (all of them for `logprobs=-1`). It reads those values from `chat_request` / `completion_request`, so without the original request, or with one that does not ask for logprobs, it returns every candidate. Requests that use `logprob_token_ids` get no output logprobs from generate (their `sampling_params.logprobs` is unset), so derender has none to return for them either. Streaming derender does the same per chunk on the plain detokenization path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). The parser path doesn't resolve them yet and drops them. +`/inference/v1/generate` returns `GenerateLogProbs` with integer `token_id`s (the generate server has no tokenizer), and derender turns those into `ChatCompletionLogProbs` / `CompletionLogProbs`, filling `token` and `bytes` from the tokenizer with the usual U+FFFD byte-fallback correction. Generate returns every top-k candidate the engine produced; derender keeps the ones the coupled endpoint would, with the same helpers: the first `top_logprobs` on chat (all of them for `top_logprobs=-1`) and the first `logprobs + 1` on completions (all of them for `logprobs=-1`). It reads those values from `chat_request` / `completion_request`, so without the original request, or with one that does not ask for logprobs, it returns every candidate. Requests that use `logprob_token_ids` get no output logprobs from generate (their `sampling_params.logprobs` is unset), so derender has none to return for them either. Streaming derender does the same per chunk on both the plain detokenization path and the parser path (see [Streaming state and logprobs](#streaming-state-and-logprobs)). As on `/v1/chat/completions`, a request with `include_reasoning: false` and a parser configured gets no logprobs at all, batch or streaming, since the decoded logprob tokens would reveal the hidden reasoning. ## Streaming cost @@ -129,7 +129,7 @@ Streaming derender is stateless on the server side: all mutable state lives in t - `logprob_context_token_ids`: the trailing sampled token IDs (at most 4) from previous chunks, used to seed byte-fallback (U+FFFD) correction so multi-byte characters whose tokens split across chunk boundaries still resolve to real strings - `logprob_text_offset`: the cumulative emitted text length, so `text_offset` in completion streaming logprobs stays absolute across chunks instead of restarting at 0 -When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the integer `token_id`s are decoded per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. +When a streamed `GenerateTokensStreamChoice` carries `logprobs`, the integer `token_id`s are decoded per chunk and the resolved logprobs are attached to the corresponding streamed choice. For chat it is as `ChatCompletionLogProbs` and for completions it is converted to the flat `CompletionLogProbs` lists. Chunks without `logprobs` produce choices with `logprobs: null`. On the parser path the chunk's logprobs are attached to whatever delta the parser emits for that chunk (reasoning, content or a tool call delta, or an empty delta while the parser is still buffering markup), which is what the coupled endpoint does; replay never touches them, since each chunk's `logprobs` cover exactly that chunk's tokens. ## Example diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index 771213309f..fc4278d3e8 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -1133,6 +1133,50 @@ async def test_e2e_parsed_reasoning(parser_client, parser_tokenizer): assert "" not in msg["content"] +@pytest.mark.asyncio +@pytest.mark.parametrize("include_reasoning", [True, False]) +async def test_e2e_parsed_logprobs_follow_include_reasoning( + parser_client, parser_tokenizer, include_reasoning +): + """With a parser configured, batch derender returns logprobs only when + reasoning is shown. Hidden reasoning drops them, as the coupled chat + endpoint does, because decoded logprob tokens would reveal it.""" + messages = [{"role": "user", "content": "What is 2+3?"}] + gen_req = await _e2e_render_chat(parser_client, PARSER_MODEL, messages) + output_ids = _require_markers_survive( + parser_tokenizer, "Two plus three.Five.", "" + ) + generate_response = _e2e_generate_response(output_ids) + generate_response["choices"][0]["logprobs"] = { + "content": [ + {"token_id": tid, "logprob": -0.5, "rank": 1, "top_logprobs": []} + for tid in output_ids + ] + } + resp = await parser_client.post( + "/v1/chat/completions/derender", + json={ + "model": PARSER_MODEL, + "generate_response": generate_response, + "prompt_tokens": len(gen_req["token_ids"]), + "chat_request": { + "model": PARSER_MODEL, + "messages": messages, + "include_reasoning": include_reasoning, + }, + }, + ) + assert resp.status_code == 200, resp.text + choice = resp.json()["choices"][0] + assert choice["message"]["content"] == "Five." + if include_reasoning: + assert choice["message"]["reasoning"] + assert len(choice["logprobs"]["content"]) == len(output_ids) + else: + assert choice["message"]["reasoning"] is None + assert choice["logprobs"] is None + + @pytest.mark.asyncio async def test_e2e_parsed_tool_call(parser_client, parser_tokenizer): """ extracted into tool_calls field.""" diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index f331873874..e2de9433dc 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -316,13 +316,18 @@ async def test_parity_logprobs(client): async def _coupled_stream( - client: httpx.AsyncClient, messages: list[dict], **extra -) -> tuple[list[dict], list[int]]: + client: httpx.AsyncClient, + messages: list[dict], + require_token_ids: bool = True, + **extra, +) -> tuple[list[dict], list[int] | None]: """Stream `/v1/chat/completions` and return the choices that carry `token_ids` along with the prompt token IDs. The leading role only chunk has no `token_ids` and no counterpart on - `/inference/v1/generate`, so it is dropped. + `/inference/v1/generate`, so it is dropped. With `require_token_ids` + off every choice is kept and `prompt_token_ids` may be `None`, for + requests where the coupled server hides token ids (hidden reasoning). """ choices: list[dict] = [] prompt_token_ids: list[int] | None = None @@ -351,9 +356,12 @@ async def _coupled_stream( if chunk.get("prompt_token_ids") is not None: prompt_token_ids = chunk["prompt_token_ids"] choices.extend( - ch for ch in chunk["choices"] if ch.get("token_ids") is not None + ch + for ch in chunk["choices"] + if not require_token_ids or ch.get("token_ids") is not None ) - assert prompt_token_ids is not None + if require_token_ids: + assert prompt_token_ids is not None return choices, prompt_token_ids @@ -473,6 +481,151 @@ async def test_stream_parity_reasoning_and_tool_call(client, chunking): pytest.skip("Model did not emit both a block and a tool call") +# --------------------------------------------------------------------------- +# Streaming logprobs on the parser path +# --------------------------------------------------------------------------- + + +def _logprob_tokens(choice: dict) -> list[tuple[str, list[int] | None]] | None: + lp = choice.get("logprobs") + if lp is None or lp.get("content") is None: + return None + return [(e["token"], e.get("bytes")) for e in lp["content"]] + + +def _sampled_generate_logprobs(choice: dict) -> dict | None: + """`GenerateLogProbs` for one coupled stream chunk from its `token_ids` + and per-token logprob values, without `top_logprobs`. + + The coupled stream's `token_ids` are the integer ids a GPU-less worker + would have sent, so one stream is both the derender input and the string + reference. A second greedy run is not needed, and greedy decoding is not + bit-reproducible across calls on every platform anyway. + """ + lp = choice.get("logprobs") + if lp is None or lp.get("content") is None: + return None + content = lp["content"] + assert len(content) == len(choice["token_ids"]), choice + return { + "content": [ + {"token_id": tid, "logprob": entry["logprob"], "top_logprobs": []} + for tid, entry in zip(choice["token_ids"], content) + ] + } + + +async def _run_stream_logprob_parity_case( + client: httpx.AsyncClient, messages: list[dict], **extra +) -> tuple[dict, list[dict], list[dict]]: + """Stream with logprobs through the coupled server, replay the same + chunks through the streaming derender endpoint and compare per chunk. + Returns the assembled coupled message, the coupled choices and the + derender choices. + """ + extra = {"logprobs": True, **extra} + coupled_choices, prompt_token_ids = await _coupled_stream(client, messages, **extra) + finish_reason = coupled_choices[-1]["finish_reason"] + derender_choices: list[dict] = [] + disagg = await stream_chat_derender( + client, + [tid for ch in coupled_choices for tid in ch["token_ids"]], + [len(ch["token_ids"]) for ch in coupled_choices], + {"model": MODEL, "messages": messages, **extra}, + len(prompt_token_ids), + prompt_token_ids, + finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, + logprobs=[_sampled_generate_logprobs(ch) for ch in coupled_choices], + raw_choices=derender_choices, + ) + coupled = assemble_stream(coupled_choices) + for message in (coupled, disagg): + for tc in message["tool_calls"]: + tc["id"] = _tool_call_id_shape(tc["id"]) + assert disagg == coupled + assert len(derender_choices) == len(coupled_choices) + for ref, der in zip(coupled_choices, derender_choices): + assert _logprob_tokens(der) == _logprob_tokens(ref), (ref, der) + return coupled, coupled_choices, derender_choices + + +@pytest.mark.asyncio +async def test_stream_parity_logprobs_reasoning(client): + """Per-chunk logprobs on the parser path match the coupled stream across + the reasoning -> content transition.""" + messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}] + coupled, _, derender_choices = await _run_stream_logprob_parity_case( + client, messages, include_reasoning=True, max_tokens=256 + ) + if not coupled["reasoning"]: + pytest.skip("Model did not emit a block") + assert any( + ch["delta"].get("reasoning") and ch.get("logprobs") for ch in derender_choices + ) + + +@pytest.mark.asyncio +async def test_stream_parity_logprobs_tool_call(client): + """Chunks the parser swallows (tool call markup) still carry their + logprobs, on both sides, so no token's logprobs are lost.""" + messages = [{"role": "user", "content": "What's the weather in Paris?"}] + coupled, _, derender_choices = await _run_stream_logprob_parity_case( + client, + messages, + tools=TOOLS, + tool_choice=FORCE_WEATHER_TOOL, + include_reasoning=True, + max_tokens=512, + ) + if not coupled["tool_calls"]: + pytest.skip("Model did not emit a tool call") + swallowed = [ + ch + for ch in derender_choices + if ch.get("logprobs") + and not ch["delta"].get("content") + and not ch["delta"].get("reasoning") + and not ch["delta"].get("tool_calls") + ] + assert swallowed, "expected at least one buffered-markup chunk with logprobs" + + +@pytest.mark.asyncio +async def test_stream_parity_logprobs_hidden_reasoning(client): + """With `include_reasoning=False` and a parser configured, neither side + returns logprobs on any chunk; derender still emits the content.""" + messages = [{"role": "user", "content": "What is 17 times 23? Think it through."}] + extra = {"logprobs": True, "max_tokens": 256} + # Token ids come from a shown-reasoning run: the coupled server hides + # them along with the logprobs when reasoning is hidden. + shown_choices, prompt_token_ids = await _coupled_stream( + client, messages, include_reasoning=True, **extra + ) + hidden_choices, _ = await _coupled_stream( + client, messages, require_token_ids=False, include_reasoning=False, **extra + ) + assert all(ch.get("logprobs") is None for ch in hidden_choices) + assert assemble_stream(hidden_choices)["reasoning"] is None + + finish_reason = shown_choices[-1]["finish_reason"] + derender_choices: list[dict] = [] + disagg = await stream_chat_derender( + client, + [tid for ch in shown_choices for tid in ch["token_ids"]], + [len(ch["token_ids"]) for ch in shown_choices], + {"model": MODEL, "messages": messages, "include_reasoning": False, **extra}, + len(prompt_token_ids), + prompt_token_ids, + finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, + logprobs=[_sampled_generate_logprobs(ch) for ch in shown_choices], + raw_choices=derender_choices, + ) + assert all(ch.get("logprobs") is None for ch in derender_choices) + shown = assemble_stream(shown_choices) + assert disagg["reasoning"] is None + assert disagg["content"] == shown["content"] + + # --------------------------------------------------------------------------- # Inline text level: generate(output_mode="text") == derender(generate(...)) # --------------------------------------------------------------------------- diff --git a/tests/entrypoints/scale_out/derender/test_derender_stream.py b/tests/entrypoints/scale_out/derender/test_derender_stream.py index 372ef31ea3..56260964a0 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_stream.py +++ b/tests/entrypoints/scale_out/derender/test_derender_stream.py @@ -1427,6 +1427,159 @@ class TestDerenderChatStreamParsed: assert delta.reasoning is None assert delta.content == "c" + @pytest.mark.asyncio + async def test_logprobs_attached_to_each_parsed_delta( + self, parsed_derenderer, tokenizer + ): + """A chunk's logprobs ride on whatever delta the parser emits for it + (reasoning, content or tool call), with token ids decoded, as on + `/v1/chat/completions`.""" + chunks = [ + [_FakeParser.REASON], + [_FakeParser.CONTENT], + [_FakeParser.TOOL_START, _FakeParser.TOOL_ARG], + ] + state = None + seen = [] + for ids in chunks: + chunk, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + ids, logprobs=_generate_logprobs(ids) + ), + state=state, + chat_request=_chat_request(), + prompt_token_ids=[1, 2], + ) + choice = chunk.choices[0] + assert choice.logprobs is not None and choice.logprobs.content is not None + assert [e.token for e in choice.logprobs.content] == ( + convert_ids_list_to_tokens(tokenizer, ids) + ) + seen.append(choice.delta) + assert seen[0].reasoning == "r" + assert seen[1].content == "c" + assert ( + seen[2].tool_calls and seen[2].tool_calls[0].function.name == "get_weather" + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("request_kwargs", "expected_candidates"), + [({"logprobs": True, "top_logprobs": 2}, 2), ({}, 4)], + ) + async def test_top_logprobs_follow_the_chat_cut( + self, parsed_derenderer, request_kwargs, expected_candidates + ): + """The parsed path trims generate's candidate list to the request's + `top_logprobs` like the plain path does, and keeps every candidate + when the request does not ask for logprobs.""" + ids = [_FakeParser.CONTENT] + logprobs = _generate_logprobs(ids) + logprobs["content"][0]["top_logprobs"] = [ + {"token_id": tid, "logprob": -float(rank)} + for rank, tid in enumerate([ids[0], 11, 12, 13]) + ] + chunk, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(ids, logprobs=logprobs), + chat_request=_chat_request(**request_kwargs), + prompt_token_ids=[1, 2], + ) + content = chunk.choices[0].logprobs.content + assert len(content) == 1 + assert len(content[0].top_logprobs) == expected_candidates + + @pytest.mark.asyncio + async def test_logprobs_kept_when_parser_swallows_chunk( + self, parsed_derenderer, tokenizer + ): + """A chunk the parser emits nothing for (buffered markup) still + carries its logprobs on an empty delta, so no token's logprobs are + lost. Mirrors the coupled endpoint, which emits a chunk whenever + logprobs are present even if the delta is empty.""" + swallowed = [_FakeParser.CONTENT + 1] # unknown to _FakeParser -> None + chunk, _ = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + swallowed, logprobs=_generate_logprobs(swallowed) + ), + chat_request=_chat_request(), + prompt_token_ids=[1, 2], + ) + choice = chunk.choices[0] + assert choice.delta.content is None and not choice.delta.tool_calls + assert [e.token for e in choice.logprobs.content] == ( + convert_ids_list_to_tokens(tokenizer, swallowed) + ) + + @pytest.mark.asyncio + async def test_include_reasoning_false_drops_logprobs_on_every_chunk( + self, parsed_derenderer + ): + """Hidden reasoning suppresses logprobs on reasoning *and* content + chunks (decoded logprob tokens would leak the reasoning), the same + rule `/v1/chat/completions` applies with a parser configured.""" + state = None + for ids in ([_FakeParser.REASON], [_FakeParser.CONTENT]): + chunk, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + ids, logprobs=_generate_logprobs(ids) + ), + state=state, + chat_request=_chat_request(include_reasoning=False), + prompt_token_ids=[1, 2], + ) + assert chunk.choices[0].logprobs is None + assert chunk.choices[0].delta.content == "c" + + @pytest.mark.asyncio + async def test_logprobs_multibyte_across_chunks(self, parsed_derenderer, tokenizer): + """Byte-fallback correction crosses chunk boundaries on the parser + path too, via the carried `logprob_context_token_ids`, which starts + empty like the engine's and never includes the prompt.""" + token_ids = tokenizer.encode("👍", add_special_tokens=False) + if len(token_ids) < 2 or "�" not in tokenizer.decode([token_ids[-1]]): + pytest.skip("Tokenizer does not byte-split this character") + state = None + last_chunk = None + for tid in token_ids: + last_chunk, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk( + [tid], logprobs=_generate_logprobs([tid]) + ), + state=state, + chat_request=_chat_request(), + prompt_token_ids=[1, 2, 3, 4, 5], + ) + assert last_chunk is not None and state is not None + final_entry = last_chunk.choices[0].logprobs.content[0] + assert not final_entry.token.endswith("�"), final_entry.token + assert state.logprob_context_token_ids == token_ids[-4:] + + @pytest.mark.asyncio + async def test_logprob_context_advances_without_logprobs(self, parsed_derenderer): + """The context tail tracks sampled ids whether or not the chunk + carries logprobs, so a later chunk that does can still be corrected.""" + ids = [_FakeParser.REASON, _FakeParser.CONTENT, _FakeParser.CONTENT] + _, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk(ids), + chat_request=_chat_request(), + prompt_token_ids=[1, 2], + ) + assert state.logprob_context_token_ids == ids + _, state = await parsed_derenderer.derender_chat_stream( + model=MODEL_NAME, + generate_chunk=_make_stream_chunk([], finish_reason="stop"), + state=state, + chat_request=_chat_request(), + prompt_token_ids=[1, 2], + ) + assert state.logprob_context_token_ids == ids + # --------------------------------------------------------------------------- # Harmony / GPT-OSS replay — unit, no server diff --git a/tests/entrypoints/scale_out/derender/utils.py b/tests/entrypoints/scale_out/derender/utils.py index c249e2053a..b6dcf2b668 100644 --- a/tests/entrypoints/scale_out/derender/utils.py +++ b/tests/entrypoints/scale_out/derender/utils.py @@ -52,6 +52,8 @@ async def stream_chat_derender( prompt_token_ids: list[int], finish_reason: str = "stop", on_chunk: Callable[[list[dict]], None] | None = None, + logprobs: list[dict | None] | None = None, + raw_choices: list[dict] | None = None, ) -> dict: """Feed `output_ids` through the streaming chat derender endpoint in the given `chunk_sizes`, threading `stream_state` across calls and @@ -64,7 +66,13 @@ async def stream_chat_derender( `tool_calls` assembled so far, letting callers assert properties of the intermediate deltas (e.g. monotonic argument growth) rather than only the final assembled result. + + `logprobs` supplies one `GenerateLogProbs` payload (or `None`) per chunk + and `raw_choices`, if given, collects every derendered choice so callers + can compare per-chunk fields such as `logprobs`. """ + if logprobs is not None: + assert len(logprobs) == len(chunk_sizes) state = None choices: list[dict] = [] pos = 0 @@ -84,6 +92,7 @@ async def stream_chat_derender( "index": 0, "token_ids": tids, "finish_reason": finish_reason if is_last else None, + "logprobs": logprobs[i] if logprobs is not None else None, } ], }, @@ -101,6 +110,8 @@ async def stream_chat_derender( assert data["chunk"]["choices"][0]["token_ids"] == tids state = data["stream_state"] choices.extend(data["chunk"]["choices"]) + if raw_choices is not None: + raw_choices.extend(data["chunk"]["choices"]) if on_chunk is not None: on_chunk(assemble_stream(choices)["tool_calls"]) diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 0c9412816a..3578c65c91 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -202,13 +202,21 @@ class OnlineDerenderer: if has_parser and chat_request is not None else {} ) - + # Like the coupled chat endpoint, hide per-token metadata (logprobs, + # token ids) when a parser is configured and reasoning is hidden: + # either would otherwise leak the reasoning text. + hide_metadata = ( + has_parser + and chat_request is not None + and not chat_request.include_reasoning + ) return_token_ids = ( chat_request is not None and chat_request.return_token_ids - and (not has_parser or chat_request.include_reasoning) + and not hide_metadata ) top_limit = _chat_top_logprobs_limit(chat_request) + for choice in generate_response.choices: if not choice.token_ids: raise ValueError(f"choice {choice.index} has empty or null token_ids") @@ -219,7 +227,7 @@ class OnlineDerenderer: tokenizer, top_limit=top_limit, ) - if choice.logprobs is not None + if choice.logprobs is not None and not hide_metadata else None ) @@ -508,13 +516,6 @@ class OnlineDerenderer: spaces_between_special_tokens=spaces_between, ) - # NOTE: parser-configured servers dispatch to - # _derender_chat_stream_parsed above and never reach this plain - # path. That parsed path does not resolve logprobs yet; when it - # does, mirror the generate chat streaming path, which suppresses - # logprobs entirely when a parser is configured and reasoning is - # hidden, because decoded logprob token text would leak hidden - # reasoning. resolved_logprobs = None if choice.logprobs is not None: resolved_logprobs = _resolve_logprobs( @@ -595,8 +596,17 @@ class OnlineDerenderer: Text comes from a fresh incremental detokenizer with special tokens preserved (``skip_special_tokens=False``), seeded from the prompt tail on every call since replay starts from scratch. + + Logprobs need no replay: a generate chunk's ``logprobs`` cover exactly + that chunk's tokens, so they are resolved once and attached to + whatever delta the parser emits for the chunk, as the coupled chat + endpoint does. Only the byte-fallback context crosses chunks, via + ``state.logprob_context_token_ids``. With ``include_reasoning=False`` + logprobs are dropped on every chunk, again like the coupled endpoint. """ tokenizer = self.renderer.get_tokenizer() + hide_metadata = not chat_request.include_reasoning + top_limit = _chat_top_logprobs_limit(chat_request) parser = parser_cls( tokenizer, @@ -651,10 +661,9 @@ class OnlineDerenderer: role_sent = state.role_sent tools_streamed = state.tools_streamed last_tool_call_ids = list(state.last_tool_call_ids) + logprob_context = list(state.logprob_context_token_ids) - return_token_ids = ( - chat_request.return_token_ids and chat_request.include_reasoning - ) + return_token_ids = chat_request.return_token_ids and not hide_metadata # At most one choice: the caller (derender_chat_stream) already # rejects >1 before dispatching here. role_sent/tools_streamed/ @@ -707,6 +716,18 @@ class OnlineDerenderer: output_token_ids.extend(delta_tids) output_chunk_lens.append(len(delta_tids)) + resolved_logprobs = None + if choice.logprobs is not None and not hide_metadata: + resolved_logprobs = _resolve_logprobs( + choice.logprobs, + tokenizer, + initial_context_token_ids=state.logprob_context_token_ids, + top_limit=top_limit, + ) + logprob_context = _logprob_context_tail( + state.logprob_context_token_ids, delta_tids + ) + if delta_message is None: delta_message = DeltaMessage() @@ -735,6 +756,7 @@ class OnlineDerenderer: stream_choice = ChatCompletionResponseStreamChoice( index=choice.index, delta=delta_message, + logprobs=resolved_logprobs, finish_reason=finish_reason, stop_reason=choice.stop_reason, token_ids=choice.token_ids if return_token_ids else None, @@ -750,6 +772,7 @@ class OnlineDerenderer: "role_sent": role_sent, "tools_streamed": tools_streamed, "last_tool_call_ids": last_tool_call_ids, + "logprob_context_token_ids": logprob_context, } ) From 4ab99ff1364735de12ac804ff4a15958001bbe06 Mon Sep 17 00:00:00 2001 From: alp Date: Fri, 16 Oct 2026 03:20:46 -0700 Subject: [PATCH 03/49] [Bugfix][Frontend] Honor the constructor chat template in offline scoring (#61990) Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> Co-authored-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> --- docs/models/pooling_models/scoring.md | 2 + .../pooling/scoring/test_io_processor_unit.py | 47 ++++++++++++++++++- vllm/entrypoints/pooling/offline.py | 4 +- .../pooling/scoring/io_processor.py | 6 ++- 4 files changed, 55 insertions(+), 4 deletions(-) diff --git a/docs/models/pooling_models/scoring.md b/docs/models/pooling_models/scoring.md index 0198810cbf..7ff186b5cd 100644 --- a/docs/models/pooling_models/scoring.md +++ b/docs/models/pooling_models/scoring.md @@ -455,6 +455,8 @@ Score templates are supported for **cross-encoder** models only. If you are usin Some scoring models require a specific prompt format to work correctly. You can specify a custom score template using the `--chat-template` parameter (see [Chat Template](../../serving/online_serving/README.md#chat-template)). +For offline scoring, pass `chat_template` to the `LLM` constructor to use it for every `LLM.score` call. A `chat_template` passed to `LLM.score` overrides the constructor's template for that call. + Like chat templates, the score template receives a `messages` list. For scoring, each message has a `role` attribute—either `"query"` or `"document"`. For the usual kind of point-wise cross-encoder, you can expect exactly two messages: one query and one document. To access the query and document content, use Jinja's `selectattr` filter: - **Query**: `{{ (messages | selectattr("role", "eq", "query") | first).content }}` diff --git a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py index cd4d58fea6..7f2c157f25 100644 --- a/tests/entrypoints/pooling/scoring/test_io_processor_unit.py +++ b/tests/entrypoints/pooling/scoring/test_io_processor_unit.py @@ -10,7 +10,7 @@ import torch from tokenizers import Tokenizer, models, pre_tokenizers, processors from transformers import TokenizersBackend -from vllm import TokensPrompt +from vllm import PoolingParams, TokensPrompt from vllm.entrypoints.pooling.scoring.io_processor import ( BiEncoderIOProcessor, CrossEncoderIOProcessor, @@ -18,7 +18,9 @@ from vllm.entrypoints.pooling.scoring.io_processor import ( LateInteractionIOProcessor, _apply_post_tokenization_to_token_type_ids, ) +from vllm.entrypoints.pooling.scoring.typing import ScoringData from vllm.entrypoints.pooling.scoring.utils import compress_token_type_ids +from vllm.entrypoints.pooling.typing import OfflineScoringInputsContext from vllm.outputs import PoolingOutput, PoolingRequestOutput, RequestError from vllm.renderers import TokenizeParams @@ -81,6 +83,49 @@ def test_llm_reranker_tokenization_is_independent_of_nonbinding_doc_limit( assert nonbinding["prompt_token_ids"] == expected +@pytest.mark.parametrize( + "default_template,request_template,expected", + [ + ("{{ messages[0].content }} {{ messages[1].content }}", None, "a bc"), + ( + "{{ messages[0].content }} {{ messages[1].content }}", + "{{ messages[1].content }} {{ messages[0].content }}", + "bc a", + ), + (None, "{{ messages[1].content }} {{ messages[0].content }}", "bc a"), + (None, None, "abc"), + ], +) +def test_offline_scoring_uses_default_template_unless_overridden( + llm_reranker_processor, default_template, request_template, expected +): + processor = llm_reranker_processor + processor.chat_template = default_template + processor.renderer = SimpleNamespace( + default_cmpl_tok_params=TokenizeParams(max_total_tokens=32) + ) + ctx = OfflineScoringInputsContext( + pooling_task="classify", + tokenization_kwargs=None, + lora_request=None, + priorities=None, + scoring_data=ScoringData(data_1=["a"], data_2=["bc"]), + chat_template=request_template, + pooling_params=PoolingParams(), + ) + factory, _ = processor.get_request_factory_offline(ctx) + (request,) = factory() + text, prompt = processor.get_score_prompt( + request["data_1"], + request["data_2"], + request["tok_params"].get_encode_kwargs(), + chat_template=request["chat_template"], + ) + + assert text == expected + assert prompt["prompt_token_ids"] == processor.tokenizer.encode(expected) + + def test_token_type_ids_stay_aligned_with_a_truncated_padded_prompt(): """The cross-encoder segment boundary must survive truncate + pad. diff --git a/vllm/entrypoints/pooling/offline.py b/vllm/entrypoints/pooling/offline.py index edf92791a8..6764730869 100644 --- a/vllm/entrypoints/pooling/offline.py +++ b/vllm/entrypoints/pooling/offline.py @@ -325,8 +325,8 @@ class PoolingOfflineMixin(OfflineInferenceMixin): it is used to create the progress bar. If `False`, no progress bar is created. lora_request: LoRA request to use for generation, if any. - chat_template: The chat template to use for the scoring. If None, we - use the model's default chat template. + chat_template: The chat template to use for scoring. If None, use + the chat template passed to the LLM constructor. tokenization_kwargs: Overrides for `tokenizer.encode`. Returns: diff --git a/vllm/entrypoints/pooling/scoring/io_processor.py b/vllm/entrypoints/pooling/scoring/io_processor.py index f6678a7bbd..7c4a1fe917 100644 --- a/vllm/entrypoints/pooling/scoring/io_processor.py +++ b/vllm/entrypoints/pooling/scoring/io_processor.py @@ -548,7 +548,11 @@ class CrossEncoderIOProcessor(ScoringIOProcessor): yield ScoringRenderParams( data_1=data_1[i], data_2=data_2[i], - chat_template=ctx.chat_template, + chat_template=( + self.chat_template + if ctx.chat_template is None + else ctx.chat_template + ), max_tokens_per_query=max_tokens_per_query, max_tokens_per_doc=max_tokens_per_doc, tok_params=tok_params, From 64175d38277f3731069cabaa9cb88983ee314f85 Mon Sep 17 00:00:00 2001 From: Martin Hickey Date: Fri, 16 Oct 2026 11:27:17 +0100 Subject: [PATCH 04/49] [Bugfix][Frontend] Keep the parser's tool call IDs in batch derender (#61852) Signed-off-by: Martin Hickey --- .../scale_out/derender/test_derender.py | 68 +++++++++++++++---- .../derender/test_derender_parity.py | 20 +++++- vllm/renderers/online_derenderer.py | 11 ++- 3 files changed, 82 insertions(+), 17 deletions(-) diff --git a/tests/entrypoints/scale_out/derender/test_derender.py b/tests/entrypoints/scale_out/derender/test_derender.py index fc4278d3e8..8bcf8c6d34 100644 --- a/tests/entrypoints/scale_out/derender/test_derender.py +++ b/tests/entrypoints/scale_out/derender/test_derender.py @@ -10,6 +10,7 @@ import pytest import pytest_asyncio from tests.utils import RemoteLaunchRenderServer +from vllm.entrypoints.generate.base.protocol import FunctionCall from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.scale_out.token_in_token_out.protocol import ( GenerateTokensResponse, @@ -1490,19 +1491,12 @@ _HERMES_TOOL_CALL_TEXT = ( ) -@pytest.mark.parametrize( - ("tool_choice", "text", "engine_finish_reason", "expected_finish_reason"), - [ - ("auto", _HERMES_TOOL_CALL_TEXT, "stop", "tool_calls"), - ("auto", _HERMES_TOOL_CALL_TEXT, "length", "length"), - ("required", "Hello there", "stop", "stop"), - ], -) -def test_derender_chat_finish_reason_matches_chat_completion( - tool_choice, text, engine_finish_reason, expected_finish_reason +def _derender_hermes_chat( + text: str, + tool_choice: str = "auto", + engine_finish_reason: str = "stop", + parser_cls: type | None = None, ): - """/derender applies the same finish_reason rule as the coupled chat - path: ``tool_calls`` only for a real call at a natural stop.""" tokenizer = get_tokenizer(MODEL_NAME) renderer = MagicMock() renderer.get_tokenizer.return_value = tokenizer @@ -1518,6 +1512,8 @@ def test_derender_chat_finish_reason_matches_chat_completion( enable_auto_tools=True, tool_parser="hermes", ) + if parser_cls is not None: + derenderer.parser = parser_cls request = ChatCompletionRequest( model=MODEL_NAME, messages=[{"role": "user", "content": "Weather in Paris?"}], @@ -1532,6 +1528,54 @@ def test_derender_chat_finish_reason_matches_chat_completion( ) (choice,) = derenderer._derender_chat(response, request) + return choice + + +@pytest.mark.parametrize( + ("tool_choice", "text", "engine_finish_reason", "expected_finish_reason"), + [ + ("auto", _HERMES_TOOL_CALL_TEXT, "stop", "tool_calls"), + ("auto", _HERMES_TOOL_CALL_TEXT, "length", "length"), + ("required", "Hello there", "stop", "stop"), + ], +) +def test_derender_chat_finish_reason_matches_chat_completion( + tool_choice, text, engine_finish_reason, expected_finish_reason +): + """/derender applies the same finish_reason rule as the coupled chat + path: ``tool_calls`` only for a real call at a natural stop.""" + choice = _derender_hermes_chat(text, tool_choice, engine_finish_reason) assert choice.finish_reason == expected_finish_reason assert bool(choice.message.tool_calls) == (text == _HERMES_TOOL_CALL_TEXT) + + +def test_derender_chat_tool_call_id_matches_chat_completion(): + """Batch /derender uses the coupled chat's default tool call ID format.""" + choice = _derender_hermes_chat(_HERMES_TOOL_CALL_TEXT) + + (tool_call,) = choice.message.tool_calls + assert tool_call.id.startswith("chatcmpl-tool-") + + +def test_derender_chat_keeps_parser_tool_call_id(): + """Batch /derender keeps the parser's tool call ID (Kimi K2 style + ``functions.:``), since chat templates reference it in + multi-turn history.""" + + class _KimiIdParser: + def __init__(self, *args, **kwargs): + pass + + def parse(self, model_output, request, **kwargs): + call = FunctionCall( + id="functions.get_weather:1", + name="get_weather", + arguments='{"city": "Paris"}', + ) + return None, None, [call] + + choice = _derender_hermes_chat(_HERMES_TOOL_CALL_TEXT, parser_cls=_KimiIdParser) + + (tool_call,) = choice.message.tool_calls + assert tool_call.id == "functions.get_weather:1" diff --git a/tests/entrypoints/scale_out/derender/test_derender_parity.py b/tests/entrypoints/scale_out/derender/test_derender_parity.py index e2de9433dc..af9bf1894a 100644 --- a/tests/entrypoints/scale_out/derender/test_derender_parity.py +++ b/tests/entrypoints/scale_out/derender/test_derender_parity.py @@ -142,12 +142,28 @@ def _tool_sig(response_choice: dict) -> list[tuple[str, dict]]: ] +def _tool_id_shapes(response_choice: dict) -> list[str | None]: + return [ + _tool_call_id_shape(tc["id"]) + for tc in (response_choice["message"].get("tool_calls") or []) + ] + + +def _engine_finish_reason(coupled_finish_reason: str) -> str: + """Undo the chat level `"tool_calls"` rewrite of `"stop"` so derender + gets what the engine reported and has to reapply the rewrite itself.""" + if coupled_finish_reason == "tool_calls": + return "stop" + return coupled_finish_reason + + def _assert_parity(coupled: dict, disagg: dict) -> None: """Both paths saw the same tokens, so they must agree unconditionally.""" c, d = coupled["choices"][0], disagg["choices"][0] assert d["message"]["content"] == c["message"]["content"] assert d["message"].get("reasoning") == c["message"].get("reasoning") assert _tool_sig(d) == _tool_sig(c) + assert _tool_id_shapes(d) == _tool_id_shapes(c) assert d["finish_reason"] == c["finish_reason"] assert d["token_ids"] == c["token_ids"] assert disagg["usage"]["prompt_tokens"] == coupled["usage"]["prompt_tokens"] @@ -166,7 +182,7 @@ async def _run_parity_case( client, ch["token_ids"], coupled["usage"]["prompt_tokens"], - ch["finish_reason"], + _engine_finish_reason(ch["finish_reason"]), chat_request, ) return coupled, disagg @@ -398,7 +414,7 @@ async def _run_stream_parity_case( {"model": MODEL, "messages": messages, "return_token_ids": True, **extra}, len(prompt_token_ids), prompt_token_ids, - finish_reason="stop" if finish_reason == "tool_calls" else finish_reason, + finish_reason=_engine_finish_reason(finish_reason), ) coupled = assemble_stream(coupled_choices) for message in (coupled, disagg): diff --git a/vllm/renderers/online_derenderer.py b/vllm/renderers/online_derenderer.py index 3578c65c91..f0169d1dba 100644 --- a/vllm/renderers/online_derenderer.py +++ b/vllm/renderers/online_derenderer.py @@ -47,14 +47,16 @@ from vllm.entrypoints.serve.utils.tool_calls_utils import ( from vllm.logger import init_logger from vllm.parser import Parser, ParserManager from vllm.renderers import BaseRenderer -from vllm.renderers.chat_utils import ChatTemplateContentFormatOption +from vllm.renderers.chat_utils import ( + ChatTemplateContentFormatOption, + make_tool_call_id, +) from vllm.tokenizers import TokenizerLike from vllm.tokenizers.detokenizer_utils import ( convert_prompt_ids_to_tokens, detokenize_incrementally, get_leading_space_marker, ) -from vllm.utils import random_uuid from vllm.utils.async_utils import make_async from vllm.v1.engine.detokenizer import uses_fast_detokenizer @@ -334,7 +336,10 @@ class OnlineDerenderer: reasoning = None tc_items = ( - [ToolCall(id=random_uuid(), function=tc) for tc in tool_calls] + [ + ToolCall(id=tc.id or make_tool_call_id(), function=tc) + for tc in tool_calls + ] if tool_calls else [] ) From de7eeaa87bc19ecf088472d5f4471cdc9a5ab995 Mon Sep 17 00:00:00 2001 From: Ayushman Singh <40520701+ayush1399@users.noreply.github.com> Date: Fri, 16 Oct 2026 08:59:42 -0400 Subject: [PATCH 05/49] [Perf] Share Triton JIT warmup compilation pool across registered kernels (#61422) --- .../test_jit_warmup_triton_launcher.py | 28 +++++---- vllm/model_executor/warmup/jit_warmup.py | 39 ++++++++++-- .../warmup/jit_warmup_triton_helper.py | 62 ++++++++++++------- 3 files changed, 90 insertions(+), 39 deletions(-) diff --git a/tests/model_executor/test_jit_warmup_triton_launcher.py b/tests/model_executor/test_jit_warmup_triton_launcher.py index c19804eae4..1c02e18b96 100644 --- a/tests/model_executor/test_jit_warmup_triton_launcher.py +++ b/tests/model_executor/test_jit_warmup_triton_launcher.py @@ -168,23 +168,28 @@ def _warmup_store_kernel(out, VALUE: tl.constexpr): not current_platform.is_cuda() or not hasattr(triton, "AsyncCompileMode"), reason="Requires CUDA Triton async compilation", ) -def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial( +def test_registered_warmup_overlaps_kernels_and_runtime_stays_serial( monkeypatch, ) -> None: from triton.runtime._async_compile import active_mode monkeypatch.setattr(envs, "VLLM_TRITON_JIT_WARMUP_NUM_THREADS", 2) - def warmup_inputs(): - return dict(out=TritonWarmupTensor(torch.int32), value=WarmupChoices(1, 2)) + def make_store(value): + def warmup_inputs(): + return dict(out=TritonWarmupTensor(torch.int32), value=value) - @triton_kernel_dispatcher_with_warmup( - kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs - ) - def store(out, value): - return (1,), dict(VALUE=value) + @triton_kernel_dispatcher_with_warmup( + kernel=_warmup_store_kernel, warmup_inputs=warmup_inputs + ) + def store(out, value): + return (1,), dict(VALUE=value) - store.get_warmup_keys() # Initialize Triton's binder before wrapping its compiler. + # Initialize Triton's binder before wrapping its compiler. + store.get_warmup_keys() + return store + + stores = [make_store(value) for value in (1, 2)] caller = threading.get_ident() barrier = threading.Barrier(2, timeout=10) compiled_on = [] @@ -200,14 +205,15 @@ def test_registered_warmup_parallelizes_compile_only_and_runtime_stays_serial( monkeypatch.setattr(_warmup_store_kernel, "compile", record_compile) registry = JitWarmupRegistry(None) with registry.activate(): - store.register_warmup() + for store in stores: + store.register_warmup() registry.warmup() assert len(compiled_on) == 2 assert all(thread != caller for thread in compiled_on) assert active_mode.get() is None out = torch.empty(1, dtype=torch.int32, device=current_platform.device_type) - for value in (1, 2): + for store, value in zip(stores, (1, 2)): store(out, value) assert out.item() == value assert len(compiled_on) == 2 diff --git a/vllm/model_executor/warmup/jit_warmup.py b/vllm/model_executor/warmup/jit_warmup.py index 84f713aa64..823c485665 100644 --- a/vllm/model_executor/warmup/jit_warmup.py +++ b/vllm/model_executor/warmup/jit_warmup.py @@ -1082,6 +1082,10 @@ class JitWarmupRegistry: from tqdm import tqdm from vllm.distributed import is_global_first_rank + from vllm.model_executor.warmup.jit_warmup_triton_helper import ( + VllmTritonJitKernel, + compile_triton_warmup, + ) kernel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = [] for kernel, registrations in self._registrations.items(): @@ -1102,17 +1106,40 @@ class JitWarmupRegistry: if not kernel_items: return + parallel_items: list[tuple[VllmJitKernel[Any], dict[Any, None]]] = [] + serial_items = [] + for kernel, compile_keys in kernel_items: + if ( + isinstance(kernel, VllmTritonJitKernel) + and kernel.supports_parallel_compile + ): + parallel_items.append((kernel, compile_keys)) + else: + serial_items.append((kernel, compile_keys)) + total_keys = sum(len(compile_keys) for _, compile_keys in kernel_items) with tqdm( - kernel_items, + total=len(kernel_items), desc=f"JIT kernel warmup ({total_keys} compile keys)", disable=not is_global_first_rank(), dynamic_ncols=True, unit="kernel", ) as progress: - for kernel, compile_keys in progress: - progress.set_postfix_str( - f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", - refresh=False, + + def compile_kernels( + items: list[tuple[VllmJitKernel[Any], dict[Any, None]]], + ) -> None: + for kernel, compile_keys in items: + progress.set_postfix_str( + f"{kernel.__class__.__name__} ({len(compile_keys)} keys)", + refresh=False, + ) + kernel.compile_many(compile_keys) + progress.update() + + if parallel_items: + compile_triton_warmup( + lambda: compile_kernels(parallel_items), + sum(len(keys) for _, keys in parallel_items), ) - kernel.compile_many(compile_keys) + compile_kernels(serial_items) diff --git a/vllm/model_executor/warmup/jit_warmup_triton_helper.py b/vllm/model_executor/warmup/jit_warmup_triton_helper.py index 70ec8cd70f..a03fe2f20b 100644 --- a/vllm/model_executor/warmup/jit_warmup_triton_helper.py +++ b/vllm/model_executor/warmup/jit_warmup_triton_helper.py @@ -6,7 +6,7 @@ from abc import abstractmethod from collections.abc import Callable, Hashable, Iterable, Mapping from concurrent.futures import ThreadPoolExecutor from contextlib import nullcontext, suppress -from contextvars import copy_context +from contextvars import ContextVar, copy_context from dataclasses import dataclass, field from functools import cache, cached_property, update_wrapper from typing import Any, Generic, ParamSpec, Protocol, TypeVar, cast, overload @@ -207,6 +207,31 @@ class TritonWarmupTensor: return result if dim is None else result[dim] +_shared_compile_pool: ContextVar[bool] = ContextVar( + "shared_triton_warmup_compile_pool", default=False +) + + +def compile_triton_warmup(compile_all: Callable[[], None], num_keys: int) -> None: + """Share one compiler pool across all eligible warmup kernels.""" + num_threads = min(envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS, num_keys) + if num_threads <= 1: + compile_all() + return + + def compile_in_context() -> None: + _shared_compile_pool.set(True) + with ( + ThreadPoolExecutor(max_workers=num_threads) as executor, + triton.AsyncCompileMode(executor), + ): + # Dispatch mutates owner state and must stay on this thread. + compile_all() + + # Older Triton versions leave async mode active on compilation failure. + copy_context().run(compile_in_context) + + class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): """Triton owner whose runtime launch specification is reused for warmup.""" @@ -230,32 +255,25 @@ class VllmTritonJitKernel(VllmJitKernel[CompileKeyT], Generic[CompileKeyT]): self._warming = False self._warming_compile_key = None + @property + def supports_parallel_compile(self) -> bool: + # AMD LLVM code generation can abort under concurrent compilation. + return ( + current_platform.is_cuda() + and not self._run_autotune + and getattr(triton, "AsyncCompileMode", None) is not None + and envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS > 1 + ) + def compile_many(self, compile_keys: Iterable[CompileKeyT]) -> None: """Compile CUDA startup warmup variants in parallel, then wait.""" keys = list(compile_keys) - num_threads = min(envs.VLLM_TRITON_JIT_WARMUP_NUM_THREADS, len(keys)) - async_compile = getattr(triton, "AsyncCompileMode", None) - # AMD LLVM code generation can abort under concurrent compilation. - if ( - not current_platform.is_cuda() - or self._run_autotune - or async_compile is None - or num_threads <= 1 - ): + if not self.supports_parallel_compile or _shared_compile_pool.get(): return super().compile_many(keys) - def compile_all() -> None: - with ( - ThreadPoolExecutor(max_workers=num_threads) as executor, - async_compile(executor), - ): - # Dispatch stays on this thread: compile() mutates owner state. - # Only Triton's underlying compiler work runs in the pool. - super(VllmTritonJitKernel, self).compile_many(keys) - - # Older Triton versions leave async mode active on compilation failure. - # Isolate that state so it cannot affect later runtime JIT compilation. - copy_context().run(compile_all) + compile_triton_warmup( + lambda: super(VllmTritonJitKernel, self).compile_many(keys), len(keys) + ) @cached_property def _kernel_arg_names(self) -> tuple[str, ...]: From 8a6c0449b36e7aba29c2a6b1720ca3c08cfafa32 Mon Sep 17 00:00:00 2001 From: Avinash Paul Date: Fri, 16 Oct 2026 19:08:55 +0530 Subject: [PATCH 06/49] [Bugfix][ROCm][GLM-5.2] Stop forcing skip_final_all_reduce on AMD DSA model path (#58589) Signed-off-by: avininjamay8 Signed-off-by: avininjamay8 Co-authored-by: avininjamay8 Co-authored-by: Kevin H. Luu Co-authored-by: vllm-agent --- vllm/models/deepseek_v32/amd/model.py | 15 +++++++++++++-- 1 file changed, 13 insertions(+), 2 deletions(-) diff --git a/vllm/models/deepseek_v32/amd/model.py b/vllm/models/deepseek_v32/amd/model.py index d2325372d9..9306180f25 100644 --- a/vllm/models/deepseek_v32/amd/model.py +++ b/vllm/models/deepseek_v32/amd/model.py @@ -8,7 +8,7 @@ from itertools import islice import torch from vllm.config import VllmConfig -from vllm.distributed import get_pp_group +from vllm.distributed import get_pp_group, get_tensor_model_parallel_rank from vllm.model_executor.layers.fused_moe import ( fused_moe_make_expert_params_mapping, ) @@ -71,6 +71,7 @@ class DeepseekV32DecoderLayer(torch.nn.Module): topk_indices_buffer=topk_indices_buffer, ) + self.zero_mlp_output = False if ( config.n_routed_experts is not None and layer_idx >= config.first_k_dense_replace @@ -82,7 +83,15 @@ class DeepseekV32DecoderLayer(torch.nn.Module): quant_config=quant_config, prefix=f"{prefix}.mlp", ) - self.mlp.experts.moe_config.skip_final_all_reduce = True + moe_cfg = self.mlp.experts.moe_config + pc = moe_cfg.moe_parallel_config + # Validated on MoRI only. DeepEP, pplx and the other combines + # keep the previous always-skip behavior. + if pc.use_mori_kernels: + # Combine is already all-reduced; the next fused norm sums TP. + self.zero_mlp_output = get_tensor_model_parallel_rank() > 0 + else: + moe_cfg.skip_final_all_reduce = True else: self.mlp = DeepseekV2MLP( hidden_size=config.hidden_size, @@ -116,6 +125,8 @@ class DeepseekV32DecoderLayer(torch.nn.Module): hidden_states, residual, self.post_attention_layernorm ) hidden_states = self.mlp(hidden_states) + if self.zero_mlp_output: + hidden_states = torch.zeros_like(hidden_states) return hidden_states, residual From 8e25bbca19f0605b347e103dde4bf5355d049e72 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:01:07 -0400 Subject: [PATCH 07/49] [Refactor] Remove dead code multiple places (#61843) Signed-off-by: yewentao256 --- csrc/activation_kernels.cu | 27 ------------- csrc/ops.h | 4 -- csrc/torch_bindings.cpp | 4 -- tests/kernels/core/test_activation.py | 40 ------------------- tests/kernels/mhc/test_mhc_kernels.py | 14 ++++++- vllm/_aiter_ops.py | 34 +--------------- vllm/model_executor/kernels/mhc/tilelang.py | 13 ------ .../layers/fused_moe/expert_map_manager.py | 7 ---- .../layers/quantization/utils/int8_utils.py | 29 -------------- vllm/v1/sample/ops/topk_topp_sampler.py | 10 ----- 10 files changed, 14 insertions(+), 168 deletions(-) diff --git a/csrc/activation_kernels.cu b/csrc/activation_kernels.cu index 9e28c3343b..0341236f77 100644 --- a/csrc/activation_kernels.cu +++ b/csrc/activation_kernels.cu @@ -1059,33 +1059,6 @@ void situ_and_mul(torch::stable::Tensor& out, // [..., d] }); } -void masked_situ_and_mul(torch::stable::Tensor& out, // [E, T, d] - torch::stable::Tensor& input, // [E, T, 2 * d] - const torch::stable::Tensor& expert_num_tokens, - double beta, double linear_beta) { - int num_experts = input.size(0); - int max_num_tokens = input.size(1); - int d = input.size(2) / 2; - if (num_experts == 0 || max_num_tokens == 0 || d == 0) { - return; - } - constexpr int block_size = 256; - dim3 grid((d + block_size - 1) / block_size, num_experts); - dim3 block(block_size); - const torch::stable::accelerator::DeviceGuard device_guard( - input.get_device_index()); - const cudaStream_t stream = get_current_cuda_stream(); - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - input.scalar_type(), "masked_situ_and_mul_kernel", [&] { - vllm::masked_situ_and_mul_kernel - <<>>( - out.mutable_data_ptr(), - input.const_data_ptr(), - expert_num_tokens.const_data_ptr(), max_num_tokens, d, - (float)beta, (float)linear_beta); - }); -} - #define LAUNCH_MASKED_ACT_AND_MUL(KERNEL, ACT_FIRST, HAS_CLAMP, LIMIT, ALPHA, \ BETA) \ vllm::masked_act_and_mul_kernel, ACT_FIRST, \ diff --git a/csrc/ops.h b/csrc/ops.h index 549dc84a1c..6e40dc7e54 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -579,10 +579,6 @@ void situ_and_mul_quant( int64_t group_size = 0, std::optional num_valid_tokens = std::nullopt, int64_t topk = 1); -void masked_situ_and_mul(torch::stable::Tensor& out, - torch::stable::Tensor& input, - const torch::stable::Tensor& expert_num_tokens, - double beta = 1.0, double linear_beta = -1.0); void masked_moe_activation(torch::stable::Tensor& out, torch::stable::Tensor& input, const torch::stable::Tensor& valid_token_counts, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index a92d9cc114..223cf5e101 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -640,9 +640,6 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "situ_and_mul_quant(Tensor! out, Tensor! scale, Tensor input, " "float beta=1.0, float linear_beta=-1.0, int group_size=0, " "Tensor? num_valid_tokens=None, int topk=1) -> ()"); - ops.def( - "masked_situ_and_mul(Tensor! out, Tensor input, Tensor " - "expert_num_tokens, float beta=1.0, float linear_beta=-1.0) -> ()"); ops.def( "masked_moe_activation(Tensor! out, Tensor input, Tensor " "valid_token_counts, str activation, float clamp_limit=0.0, float " @@ -875,7 +872,6 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { ops.impl("swigluoai_and_mul", TORCH_BOX(&swigluoai_and_mul)); ops.impl("situ_and_mul", TORCH_BOX(&situ_and_mul)); ops.impl("situ_and_mul_quant", TORCH_BOX(&situ_and_mul_quant)); - ops.impl("masked_situ_and_mul", TORCH_BOX(&masked_situ_and_mul)); ops.impl("masked_moe_activation", TORCH_BOX(&masked_moe_activation)); ops.impl("gelu_new", TORCH_BOX(&gelu_new)); ops.impl("gelu_fast", TORCH_BOX(&gelu_fast)); diff --git a/tests/kernels/core/test_activation.py b/tests/kernels/core/test_activation.py index 85ec4903ee..9303e23553 100644 --- a/tests/kernels/core/test_activation.py +++ b/tests/kernels/core/test_activation.py @@ -353,46 +353,6 @@ def test_silu_and_mul_with_clamp( ) -@pytest.mark.parametrize("linear_beta", [-1.0, 2.0]) -@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16]) -@torch.inference_mode() -def test_masked_situ_and_mul( - default_vllm_config, - linear_beta: float, - dtype: torch.dtype, -) -> None: - """Masked SITU computes valid expert rows and preserves padded zeros.""" - device = CUDA_DEVICES[0] - num_experts, max_num_tokens, d = 4, 7, 512 - beta = 1.5 - input = torch.randn(num_experts, max_num_tokens, 2 * d, dtype=dtype, device=device) - expert_num_tokens = torch.tensor([0, 1, 4, 7], dtype=torch.int32, device=device) - output = torch.zeros(num_experts, max_num_tokens, d, dtype=dtype, device=device) - - torch.ops._C.masked_situ_and_mul( - output, input, expert_num_tokens, beta, linear_beta - ) - - gate, up = input.float().chunk(2, dim=-1) - expected = beta * torch.tanh(gate / beta) * torch.sigmoid(gate) - if linear_beta > 0: - up = linear_beta * torch.tanh(up / linear_beta) - expected = (expected * up).to(dtype) - for expert, num_tokens in enumerate(expert_num_tokens.cpu().tolist()): - torch.testing.assert_close( - output[expert, :num_tokens], - expected[expert, :num_tokens], - atol=get_default_atol(output), - rtol=get_default_rtol(output), - ) - assert torch.count_nonzero(output[expert, num_tokens:]) == 0 - - opcheck( - torch.ops._C.masked_situ_and_mul, - (output, input, expert_num_tokens, beta, linear_beta), - ) - - MOE_ACTIVATION_CASES = [ pytest.param(MoEActivation.SILU, ApplyMoEActivationConfig(), id="silu"), pytest.param( diff --git a/tests/kernels/mhc/test_mhc_kernels.py b/tests/kernels/mhc/test_mhc_kernels.py index 4afb45f503..7969c82dda 100644 --- a/tests/kernels/mhc/test_mhc_kernels.py +++ b/tests/kernels/mhc/test_mhc_kernels.py @@ -10,7 +10,6 @@ import torch.nn.functional as F import vllm.model_executor.kernels.mhc # noqa: F401 import vllm.model_executor.layers.mhc as mhc_layers from vllm.model_executor.kernels.mhc.tilelang import ( - _torch_hc_prenorm_gemm, mhc_post_tilelang, mhc_pre_delayed_tilelang, ) @@ -1002,6 +1001,19 @@ def test_mhc_pre_tilelang(num_tokens, hidden_size, hc_mult, fused_norm): torch.testing.assert_close(actual, expected, atol=5e-2, rtol=1e-2) +def _torch_hc_prenorm_gemm( + x: torch.Tensor, + fn: torch.Tensor, + out: torch.Tensor, + sqrsum: torch.Tensor, +) -> None: + assert out.shape[0] == 1 + assert sqrsum.shape[0] == 1 + x_float = x.float() + out[0].copy_(x_float @ fn.t()) + sqrsum[0].copy_(x_float.square().sum(dim=-1)) + + @pytest.mark.skipif( not HAS_TILELANG_MHC, reason="TileLang MHC support required", diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 0f57c9ecb4..fadb053915 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -21,7 +21,7 @@ from vllm.v1.attention.ops.rocm_aiter_mla_sparse import ( ) if TYPE_CHECKING: - from aiter import ActivationType, QuantType + from aiter import ActivationType logger = init_logger(__name__) @@ -2202,38 +2202,6 @@ class rocm_aiter_ops: } return mapping.get(name) - @staticmethod - def get_aiter_quant_type(quant_type_str: str) -> "QuantType | None": - """Given a quantization type as a string, returns the corresponding aiter QuantType enum. - Supported quantization types: "no", "per_tensor", "per_token", "per_1x32", "per_1x128", "per_128x128". - Returns None if the mapping fails. - - Args: - quant_type_str (str): Quantization type as string. - - Returns: - Aiter QuantType enum value, or None if not found. - - """ - try: - from aiter import QuantType - except ImportError: - return None - - if not isinstance(quant_type_str, str): - return None - - name = quant_type_str.strip().lower() - mapping = { - "no": QuantType.No, - "per_tensor": QuantType.per_Tensor, - "per_token": QuantType.per_Token, - "per_1x32": QuantType.per_1x32, - "per_1x128": QuantType.per_1x128, - "per_128x128": QuantType.per_128x128, - } - return mapping.get(name) - @classmethod @if_aiter_supported def is_enabled(cls) -> bool: diff --git a/vllm/model_executor/kernels/mhc/tilelang.py b/vllm/model_executor/kernels/mhc/tilelang.py index 62741bbda1..2c0f4c12a0 100644 --- a/vllm/model_executor/kernels/mhc/tilelang.py +++ b/vllm/model_executor/kernels/mhc/tilelang.py @@ -6,19 +6,6 @@ from vllm.utils.math_utils import cdiv from vllm.utils.torch_utils import direct_register_custom_op -def _torch_hc_prenorm_gemm( - x: torch.Tensor, - fn: torch.Tensor, - out: torch.Tensor, - sqrsum: torch.Tensor, -) -> None: - assert out.shape[0] == 1 - assert sqrsum.shape[0] == 1 - x_float = x.float() - out[0].copy_(x_float @ fn.t()) - sqrsum[0].copy_(x_float.square().sum(dim=-1)) - - def _hc_prenorm_gemm_outputs( x: torch.Tensor, fn: torch.Tensor, diff --git a/vllm/model_executor/layers/fused_moe/expert_map_manager.py b/vllm/model_executor/layers/fused_moe/expert_map_manager.py index daa2b3ff8e..1035654378 100644 --- a/vllm/model_executor/layers/fused_moe/expert_map_manager.py +++ b/vllm/model_executor/layers/fused_moe/expert_map_manager.py @@ -359,13 +359,6 @@ class ExpertMapManager: return True return self._expert_map[global_id] != -1 - def get_local_expert_ids(self) -> list[int]: - """Get list of global IDs for experts on this rank.""" - if self._expert_map is None: - return list(range(self.global_num_experts)) - - return torch.where(self._expert_map != -1)[0].tolist() - def update( self, moe_parallel_config: FusedMoEParallelConfig, diff --git a/vllm/model_executor/layers/quantization/utils/int8_utils.py b/vllm/model_executor/layers/quantization/utils/int8_utils.py index 1180e2df36..eb3ff4fe27 100644 --- a/vllm/model_executor/layers/quantization/utils/int8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/int8_utils.py @@ -12,35 +12,6 @@ from vllm.triton_utils import tl, triton logger = logging.getLogger(__name__) -def block_dequant( - x_q_block: torch.Tensor, - x_s: torch.Tensor, - block_size: list[int], -) -> torch.Tensor: - """This function conducts block-wise dequantization. - The inputs are block-wise quantization tensor `x_q_block`, - block-wise quantization scale and the block size. - The outputs are dequantized tensor. - """ - block_n, block_k = block_size[0], block_size[1] - n, k = x_q_block.shape - n_tiles = (n + block_n - 1) // block_n - k_tiles = (k + block_k - 1) // block_k - assert n_tiles == x_s.shape[0] - assert k_tiles == x_s.shape[1] - - x_dq_block = x_q_block.to(torch.float32) - - for i in range(k_tiles): - for j in range(n_tiles): - x_dq_block[ - j * block_n : min((j + 1) * block_n, n), - i * block_k : min((i + 1) * block_k, k), - ] *= x_s[j][i] - - return x_dq_block - - if current_platform.is_rocm(): @triton.jit diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index d38497b2ff..085f092341 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -483,16 +483,6 @@ class TopKTopPSampler(nn.Module): return xpu_sample(logits, k, p, self.logprobs_mode) -# Note: this is a workaround for -# https://github.com/pytorch/pytorch/pull/151218 -@torch.compile(dynamic=True, backend=current_platform.simple_compile_backend) -def compiled_random_sample(logits: torch.Tensor) -> torch.Tensor: - probs = logits.softmax(dim=-1, dtype=torch.float32) - q = torch.empty_like(probs) - q.exponential_() - return probs.div(q).argmax(dim=-1).view(-1) - - def apply_top_k_top_p( logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None ) -> torch.Tensor: From 1a5b187066c918b2e23184560d703174dd542828 Mon Sep 17 00:00:00 2001 From: bnellnm <49004751+bnellnm@users.noreply.github.com> Date: Fri, 16 Oct 2026 11:16:03 -0400 Subject: [PATCH 08/49] Add @bnellnm to CODEOWNERS (#62025) Signed-off-by: Bill Nell --- .github/CODEOWNERS | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS index b14edb0d5b..c621233907 100644 --- a/.github/CODEOWNERS +++ b/.github/CODEOWNERS @@ -8,7 +8,7 @@ /vllm/distributed/kv_transfer/kv_connector/v1/hisparse @MatthewBonanni @LucasWilkinson /vllm/lora @jeejeelee /vllm/model_executor/layers/attention @LucasWilkinson @MatthewBonanni -/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18 +/vllm/model_executor/layers/fused_moe @mgoin @pavanimajety @zyongye @wzhao18 @bnellnm /vllm/model_executor/layers/quantization @mgoin @robertgshaw2-redhat @tlrmchlsmth @yewentao256 @pavanimajety @zyongye /vllm/model_executor/layers/mamba @tdoublep @tomeras91 @ZJY0516 /vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @tdoublep @ZJY0516 @vadiklyutiy @@ -138,6 +138,7 @@ /tests/evals @mgoin @vadiklyutiy @AndreasKaratzas @ZJY0516 /tests/kernels @mgoin @tlrmchlsmth @WoosukKwon @yewentao256 @zyongye @AndreasKaratzas /tests/kernels/ir @ProExpertProg @tjtanaa +/tests/kernels/moe @bnellnm /tests/models @DarkLight1337 @ywang96 @AndreasKaratzas /tests/multimodal @DarkLight1337 @ywang96 @NickLucche @shen-shanshan /tests/multimodal/media @Isotr0py From 76081b9ee4bfc73fc307505d633ada5bce079880 Mon Sep 17 00:00:00 2001 From: Sawyer Bowerman Date: Fri, 16 Oct 2026 11:49:26 -0400 Subject: [PATCH 09/49] [Bugfix][Model] Honor quantization config in Inkling MTP (#61921) Signed-off-by: Sawyer Bowerman Co-authored-by: Codex --- vllm/models/inkling/amd/mtp.py | 38 +++++++++++++++++++++++-------- vllm/models/inkling/nvidia/mtp.py | 38 +++++++++++++++++++++++-------- 2 files changed, 58 insertions(+), 18 deletions(-) diff --git a/vllm/models/inkling/amd/mtp.py b/vllm/models/inkling/amd/mtp.py index f2ed51c075..b11ea511fb 100644 --- a/vllm/models/inkling/amd/mtp.py +++ b/vllm/models/inkling/amd/mtp.py @@ -4,7 +4,7 @@ Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj`` -(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the +(``2H -> H``) and a full Inkling transformer block (dense MLP, with the same short convolutions as the backbone; its attention is full or sliding-window per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared ``chain_norm`` is applied after every depth; its output is both the logits input @@ -28,18 +28,23 @@ from torch import nn from vllm.config import VllmConfig from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.model_loader.mtp_validation import ( is_mtp_completeness_check_enabled, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsMultiModalEmbeddings -from vllm.model_executor.models.utils import maybe_prefix +from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from ..configs import InklingModelConfig from .layernorm import InklingRMSNorm -from .model import InklingDecoderLayer, InklingReplicatedEmbedding +from .model import ( + InklingDecoderLayer, + InklingReplicatedEmbedding, + _TmlForCausalLMBase, +) from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm # Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. @@ -71,7 +76,13 @@ def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int: class InklingMTPDepthLayer(nn.Module): """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" - def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: + def __init__( + self, + config: InklingModelConfig, + prefix: str, + is_local: bool, + quant_config: QuantizationConfig | None = None, + ) -> None: super().__init__() self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -80,16 +91,17 @@ class InklingMTPDepthLayer(nn.Module): config.hidden_size, bias=False, return_bias=False, + quant_config=quant_config, prefix=f"{prefix}.input_proj", ) - # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs + # A force-dense-MLP block; ``is_local`` selects sliding-window vs # full attention (the swa_* head config and sliding_window window) to # match this depth's checkpoint transformer_block weights. self.transformer_block = InklingDecoderLayer( config, layer_id=0, is_local=is_local, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.transformer_block", force_dense_mlp=True, ) @@ -131,7 +143,10 @@ class InklingMultiTokenPredictor(nn.Module): self.layers = nn.ModuleDict( { str(idx): InklingMTPDepthLayer( - config, f"{prefix}.layers.{idx}", idx in local_ids + config, + f"{prefix}.layers.{idx}", + idx in local_ids, + quant_config=vllm_config.quant_config, ) for idx in range(self.num_mtp_layers) } @@ -244,6 +259,11 @@ class InklingMultiTokenPredictor(nn.Module): class InklingMTP(nn.Module, SupportsMultiModalEmbeddings): + hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"model.mtp.": "model."}, + ) + packed_modules_mapping = _TmlForCausalLMBase.packed_modules_mapping + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() assert vllm_config.speculative_config is not None @@ -418,9 +438,9 @@ def _load_inkling_mtp_weights( continue # Dense MLP fused gate/up + down. - if ".mlp.w13_dn.weight" in name: + if ".mlp.w13_dn." in name: loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) - elif ".mlp.w2_md.weight" in name: + elif ".mlp.w2_md." in name: loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) else: if name.endswith(".bias") and name not in params: diff --git a/vllm/models/inkling/nvidia/mtp.py b/vllm/models/inkling/nvidia/mtp.py index f2ed51c075..b11ea511fb 100644 --- a/vllm/models/inkling/nvidia/mtp.py +++ b/vllm/models/inkling/nvidia/mtp.py @@ -4,7 +4,7 @@ Mirrors the reference ``mtp_model.py`` shipped with the checkpoint: each MTP depth ``i`` owns ``hidden_norm`` / ``embed_norm`` RMSNorms, an ``input_proj`` -(``2H -> H``) and a full Inkling transformer block (dense bf16 MLP, with the +(``2H -> H``) and a full Inkling transformer block (dense MLP, with the same short convolutions as the backbone; its attention is full or sliding-window per depth, selected by ``mtp_config.local_layer_ids``). When enabled, a shared ``chain_norm`` is applied after every depth; its output is both the logits input @@ -28,18 +28,23 @@ from torch import nn from vllm.config import VllmConfig from vllm.model_executor.layers.linear import ReplicatedLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.model_loader.mtp_validation import ( is_mtp_completeness_check_enabled, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsMultiModalEmbeddings -from vllm.model_executor.models.utils import maybe_prefix +from vllm.model_executor.models.utils import WeightsMapper, maybe_prefix from vllm.sequence import IntermediateTensors from ..configs import InklingModelConfig from .layernorm import InklingRMSNorm -from .model import InklingDecoderLayer, InklingReplicatedEmbedding +from .model import ( + InklingDecoderLayer, + InklingReplicatedEmbedding, + _TmlForCausalLMBase, +) from .ops.norm import embed_dual_rmsnorm_cat, embed_rmsnorm # Checkpoint attention projections (wq_du/wk_dv/wv_dv/wr_du) -> fused qkvr. @@ -71,7 +76,13 @@ def _select_mtp_depth_count(n_predict: int, num_spec: int | None) -> int: class InklingMTPDepthLayer(nn.Module): """One MTP depth: norm both inputs, fuse (2H->H), run a Inkling block.""" - def __init__(self, config: InklingModelConfig, prefix: str, is_local: bool) -> None: + def __init__( + self, + config: InklingModelConfig, + prefix: str, + is_local: bool, + quant_config: QuantizationConfig | None = None, + ) -> None: super().__init__() self.hidden_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) self.embed_norm = InklingRMSNorm(config.hidden_size, eps=config.rms_norm_eps) @@ -80,16 +91,17 @@ class InklingMTPDepthLayer(nn.Module): config.hidden_size, bias=False, return_bias=False, + quant_config=quant_config, prefix=f"{prefix}.input_proj", ) - # A force-dense-MLP bf16 block; ``is_local`` selects sliding-window vs + # A force-dense-MLP block; ``is_local`` selects sliding-window vs # full attention (the swa_* head config and sliding_window window) to # match this depth's checkpoint transformer_block weights. self.transformer_block = InklingDecoderLayer( config, layer_id=0, is_local=is_local, - quant_config=None, + quant_config=quant_config, prefix=f"{prefix}.transformer_block", force_dense_mlp=True, ) @@ -131,7 +143,10 @@ class InklingMultiTokenPredictor(nn.Module): self.layers = nn.ModuleDict( { str(idx): InklingMTPDepthLayer( - config, f"{prefix}.layers.{idx}", idx in local_ids + config, + f"{prefix}.layers.{idx}", + idx in local_ids, + quant_config=vllm_config.quant_config, ) for idx in range(self.num_mtp_layers) } @@ -244,6 +259,11 @@ class InklingMultiTokenPredictor(nn.Module): class InklingMTP(nn.Module, SupportsMultiModalEmbeddings): + hf_to_vllm_mapper = _TmlForCausalLMBase.hf_to_vllm_mapper | WeightsMapper( + orig_to_new_prefix={"model.mtp.": "model."}, + ) + packed_modules_mapping = _TmlForCausalLMBase.packed_modules_mapping + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: super().__init__() assert vllm_config.speculative_config is not None @@ -418,9 +438,9 @@ def _load_inkling_mtp_weights( continue # Dense MLP fused gate/up + down. - if ".mlp.w13_dn.weight" in name: + if ".mlp.w13_dn." in name: loaded_weight = _load(name.replace(".w13_dn.", ".gate_up_proj."), weight) - elif ".mlp.w2_md.weight" in name: + elif ".mlp.w2_md." in name: loaded_weight = _load(name.replace(".w2_md.", ".down_proj."), weight) else: if name.endswith(".bias") and name not in params: From 706791d5162355dfe2a924dd0257770fc529ec10 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Fri, 16 Oct 2026 10:58:59 -0500 Subject: [PATCH 10/49] [CI][AMD] Move 17 test groups from MI300 to MI355 and test large DeepSeek/Nemotron evals on AITER (#61955) Signed-off-by: Andreas Karatzas Co-authored-by: Claude --- .buildkite/test-amd.yaml | 449 ++++++++---------- .buildkite/test_areas/engine.yaml | 4 +- .buildkite/test_areas/lm_eval.yaml | 23 +- .buildkite/test_areas/models_language.yaml | 4 +- .buildkite/test_areas/models_multimodal.yaml | 12 +- .buildkite/test_areas/spec_decode.yaml | 28 +- .../gsm8k/configs/DeepSeek-R1-DP_MI325.yaml | 12 - .../gsm8k/configs/DeepSeek-R1-TP_MI325.yaml | 12 - .../gsm8k/configs/DeepSeek-V3.2-DP_MI325.yaml | 12 - .../gsm8k/configs/DeepSeek-V3.2-TP_MI325.yaml | 12 - ...Nemotron-3-Super-120B-A12B-BF16_MI355.yaml | 1 + .../Nemotron-3-Super-120B-A12B-FP8_MI355.yaml | 1 + tests/evals/gsm8k/configs/models-mi3xx.txt | 8 - 13 files changed, 228 insertions(+), 350 deletions(-) delete mode 100644 tests/evals/gsm8k/configs/DeepSeek-R1-DP_MI325.yaml delete mode 100644 tests/evals/gsm8k/configs/DeepSeek-R1-TP_MI325.yaml delete mode 100644 tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI325.yaml delete mode 100644 tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI325.yaml delete mode 100644 tests/evals/gsm8k/configs/models-mi3xx.txt diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 7a9876032d..0b04b36878 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -1094,25 +1094,6 @@ steps: #----------------------------------------------------------- mi300 · evals -----------------------------------------------------------# -- label: ":amd: (MI300) LM Eval Small Models" # TBD - timeout_in_minutes: 45 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - csrc/ - - vllm/model_executor/layers/quantization - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - optional: true - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt - - label: ":amd: (MI300) DeepSeek V2-Lite Prefetch Offload Accuracy" # TBD timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1153,23 +1134,6 @@ steps: commands: - pytest -s -v evals/mrcr/test_mrcr_correctness.py --config-list-file=evals/mrcr/configs/models-small.txt -- label: ":amd: (MI300) LM Eval KV-Offload" # TBD - timeout_in_minutes: 55 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py - - vllm/v1/kv_offload/ - - vllm/v1/simple_kv_offload/ - - tests/evals/gsm8k/test_gsm8k_offloading.py - - vllm/platforms/rocm.py - commands: - - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" - - label: ":amd: (MI300) LM Eval KV-Offload Medium" # TBD timeout_in_minutes: 70 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1188,24 +1152,6 @@ steps: commands: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "qwen3.5-35b or deepseek-v2-lite" -- label: ":amd: (MI300) LM Eval KV-Offload Large" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ - - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py - - vllm/v1/kv_offload/ - - vllm/v1/simple_kv_offload/ - - tests/evals/gsm8k/test_gsm8k_offloading.py - - vllm/platforms/rocm.py - commands: - - VLLM_ROCM_USE_AITER=1 pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" - - label: ":amd: (MI300) LM Eval Small Models Harness" # TBD timeout_in_minutes: 35 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1383,52 +1329,6 @@ steps: commands: - VLLM_ENGINE_READY_TIMEOUT_S=1800 bash .buildkite/scripts/scheduled_integration_test/qwen3_next_mtp_async_eplb.sh 0.8 1319 8040 -- label: ":amd: (MI300) LM Eval Large Models" # TBD - timeout_in_minutes: 65 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_8 - optional: true - num_gpus: 8 - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/layernorm.py - - csrc/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - - tests/evals/ - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export PYTORCH_ROCM_ARCH=gfx942 - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt - -- label: ":amd: (MI300) LM Eval Large Models ROCm Harness" # TBD - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_8 - optional: true - num_gpus: 8 - working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" - source_file_dependencies: - - vllm/model_executor/models/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/quantization/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/layers/layernorm.py - - csrc/ - - vllm/_aiter_ops.py - - vllm/platforms/rocm.py - commands: - - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400 - - label: ":amd: (MI300) LM Eval Large Models FP8" # TBD timeout_in_minutes: 40 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1939,20 +1839,6 @@ steps: commands: - pytest -v -s models/language/pooling_mteb_test -- label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N" - timeout_in_minutes: 95 - parallelism: 4 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/language/pooling - commands: - - pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB - - label: ":amd: (MI300) Language Models (Extended Generation)" # TBD timeout_in_minutes: 105 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -1970,21 +1856,6 @@ steps: #---------------------------------------------------- mi300 · models / multimodal ----------------------------------------------------# -- label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" # TBD - timeout_in_minutes: 100 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation - - tests/models/multimodal/test_mapping.py - commands: - - pytest -v -s models/multimodal/generation -m 'not core_model' --ignore models/multimodal/generation/test_common.py - - pytest -v -s models/multimodal/test_mapping.py - - label: ":amd: (MI300) Multimodal Models (Extended Generation 2) Shard %N" timeout_in_minutes: 100 parallelism: 4 @@ -2013,19 +1884,6 @@ steps: commands: - pytest -v -s models/multimodal/generation/test_common.py -m 'split(group=1) and not core_model' -- label: ":amd: (MI300) Multimodal Models (PPL)" - timeout_in_minutes: 90 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/generation_ppl_test/ - commands: - - pytest -v -s models/multimodal/generation_ppl_test/ - - label: ":amd: (MI300) Multimodal Models (Standard) 3: llava + qwen2_vl" # TBD timeout_in_minutes: 55 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2057,19 +1915,6 @@ steps: commands: - pytest -v -s models/multimodal/processing/test_tensor_schema.py --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB -- label: ":amd: (MI300) Multimodal Models (Extended Pooling)" # TBD - timeout_in_minutes: 75 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/ - - tests/models/multimodal/pooling - commands: - - pytest -v -s models/multimodal/pooling -m 'not core_model' - #----------------------------------------------------- mi300 · models / quantized -----------------------------------------------------# - label: ":amd: (MI300) Quantized Models" # TBD @@ -2581,60 +2426,6 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py -- label: ":amd: (MI300) Spec Decode AL DFlash Nightly" - timeout_in_minutes: 120 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/models/qwen3_dflash.py - - vllm/model_executor/models/qwen3_dflash2.py - - vllm/model_executor/models/qwen3_5.py - - vllm/model_executor/models/laguna_dflash.py - - tests/v1/e2e/spec_decode/ - - vllm/v1/attention/ - - vllm/v1/sample/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/qwen2.py - - vllm/model_executor/models/qwen3.py - - vllm/model_executor/models/laguna.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" - -- label: ":amd: (MI300) Spec Decode AL DSpark Nightly" - timeout_in_minutes: 90 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/models/qwen3_dspark.py - - vllm/model_executor/models/gemma4_dspark.py - - tests/v1/e2e/spec_decode/ - - vllm/v1/attention/ - - vllm/v1/sample/ - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/qwen2.py - - vllm/model_executor/models/qwen3.py - - vllm/model_executor/models/qwen3_dflash.py - - vllm/model_executor/models/gemma4.py - - vllm/model_executor/models/gemma4_mm.py - - vllm/model_executor/models/gemma4_unified.py - - vllm/model_executor/models/gemma4_mtp.py - - vllm/platforms/rocm.py - commands: - - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ - - label: ":amd: (MI300) V1 Spec Decode" # TBD timeout_in_minutes: 50 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -2719,39 +2510,6 @@ steps: - pytest -v -s v1/engine --ignore v1/engine/test_preprocess_error_handling.py - pytest -v -s v1/test_tensor_ipc_queue.py -- label: ":amd: (MI300) Speculators Correctness Nightly" - timeout_in_minutes: 40 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/config/speculative.py - - vllm/engine/arg_utils.py - - vllm/transformers_utils/config.py - - vllm/transformers_utils/configs/speculators/ - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/v1/worker/gpu_model_runner.py - - vllm/v1/sample/ - - vllm/v1/attention/backends/ - - vllm/v1/attention/selector.py - - vllm/model_executor/model_loader/ - - vllm/model_executor/layers/ - - vllm/model_executor/models/llama_eagle3.py - - vllm/model_executor/models/qwen3.py - - vllm/model_executor/models/qwen3_dflash.py - - vllm/model_executor/models/registry.py - - vllm/_aiter_ops.py - - tests/evals/gsm8k/ - - tests/v1/spec_decode/test_speculators_correctness.py - - vllm/platforms/rocm.py - - vllm/v1/worker/gpu/spec_decode/dflash/ - commands: - - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test - - label: ":amd: (MI300) Extract Hidden States Integration" # TBD timeout_in_minutes: 30 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -3106,22 +2864,6 @@ steps: - uv pip install --system -r /vllm-workspace/requirements/kv_connectors_rocm.txt - ATTENTION_BACKEND=TRITON_ATTN bash v1/kv_connector/nixl_integration/run_multi_connector_edge_case_test.sh -- label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" # TBD - timeout_in_minutes: 35 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_4 - num_gpus: 4 - optional: true - source_file_dependencies: - - vllm/v1/attention/backends/utils.py - - vllm/v1/worker/gpu_model_runner.py - - tests/v1/e2e/test_hybrid_chunked_prefill.py - - vllm/v1/core/ - working_dir: /vllm-workspace/tests - commands: - - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py - ######################################################################################################################################### # # # MI355 (gfx950) tests # @@ -4604,6 +4346,82 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-nemotron-3-super.txt +- label: ":amd: (MI355) LM Eval Small Models" # TBD + timeout_in_minutes: 45 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - csrc/ + - vllm/model_executor/layers/quantization + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + optional: true + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt + +- label: ":amd: (MI355 DPX) LM Eval KV-Offload" # TBD + timeout_in_minutes: 55 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/kv_offload/ + - vllm/v1/simple_kv_offload/ + - tests/evals/gsm8k/test_gsm8k_offloading.py + - vllm/platforms/rocm.py + commands: + - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" + +- label: ":amd: (MI355) LM Eval KV-Offload Large" # TBD + timeout_in_minutes: 65 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_4 + num_gpus: 4 + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/offloading/ + - vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py + - vllm/v1/kv_offload/ + - vllm/v1/simple_kv_offload/ + - tests/evals/gsm8k/test_gsm8k_offloading.py + - vllm/platforms/rocm.py + commands: + - VLLM_ROCM_USE_AITER=1 pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" + +- label: ":amd: (MI355) LM Eval Large Models ROCm Harness" # TBD + timeout_in_minutes: 40 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_8 + optional: true + num_gpus: 8 + working_dir: "/vllm-workspace/.buildkite/lm-eval-harness" + source_file_dependencies: + - vllm/model_executor/models/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/quantization/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/layers/layernorm.py + - csrc/ + - vllm/_aiter_ops.py + - vllm/platforms/rocm.py + commands: + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large-rocm.txt --tp-size=8 --timeout=5400 + #--------------------------------------------------------- mi355 · examples ----------------------------------------------------------# - label: ":amd: (MI355 DPX) Examples" # TBD @@ -5421,6 +5239,19 @@ steps: - pytest models/multimodal/generation/test_memory_leak.py -m core_model - cd .. && VLLM_WORKER_MULTIPROC_METHOD=spawn pytest -v -s tests/models/multimodal/generation/test_whisper.py -m core_model +- label: ":amd: (MI355 DPX) Multimodal Models (PPL)" + timeout_in_minutes: 90 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/ + - tests/models/multimodal/generation_ppl_test/ + commands: + - pytest -v -s models/multimodal/generation_ppl_test/ + #----------------------------------------------------- mi355 · models / quantized -----------------------------------------------------# - label: ":amd: (MI355 DPX) Quantized Models" # TBD @@ -6041,6 +5872,108 @@ steps: - if [ "$$BUILDKITE_PARALLEL_JOB" = "1" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[False]"; fi - if [ "$$BUILDKITE_PARALLEL_JOB" = "2" ]; then pytest -v -s "v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py::test_gemma4_mtp_acceptance_lengths[True]"; fi +- label: ":amd: (MI355 DPX) V1 E2E Hybrid Chunked Prefill" # TBD + timeout_in_minutes: 35 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + source_file_dependencies: + - vllm/v1/attention/backends/utils.py + - vllm/v1/worker/gpu_model_runner.py + - tests/v1/e2e/test_hybrid_chunked_prefill.py + - vllm/v1/core/ + working_dir: /vllm-workspace/tests + commands: + - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py + +- label: ":amd: (MI355 DPX) Speculators Correctness Nightly" + timeout_in_minutes: 40 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/config/speculative.py + - vllm/engine/arg_utils.py + - vllm/transformers_utils/config.py + - vllm/transformers_utils/configs/speculators/ + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/worker/gpu_model_runner.py + - vllm/v1/sample/ + - vllm/v1/attention/backends/ + - vllm/v1/attention/selector.py + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/model_executor/models/llama_eagle3.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/registry.py + - vllm/_aiter_ops.py + - tests/evals/gsm8k/ + - tests/v1/spec_decode/test_speculators_correctness.py + - vllm/platforms/rocm.py + - vllm/v1/worker/gpu/spec_decode/dflash/ + commands: + - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 + - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test + +- label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly" + timeout_in_minutes: 120 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/qwen3_dflash2.py + - vllm/model_executor/models/qwen3_5.py + - vllm/model_executor/models/laguna_dflash.py + - tests/v1/e2e/spec_decode/ + - vllm/v1/attention/ + - vllm/v1/sample/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/laguna.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" + +- label: ":amd: (MI355 DPX) Spec Decode AL DSpark Nightly" + timeout_in_minutes: 90 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/models/qwen3_dspark.py + - vllm/model_executor/models/gemma4_dspark.py + - tests/v1/e2e/spec_decode/ + - vllm/v1/attention/ + - vllm/v1/sample/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/gemma4.py + - vllm/model_executor/models/gemma4_mm.py + - vllm/model_executor/models/gemma4_unified.py + - vllm/model_executor/models/gemma4_mtp.py + - vllm/platforms/rocm.py + commands: + - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ + #----------------------------------------------------------- mi355 · misc ------------------------------------------------------------# - label: ":amd: (MI355 DPX) Regression" # TBD diff --git a/.buildkite/test_areas/engine.yaml b/.buildkite/test_areas/engine.yaml index bddee97a3a..ef4b7951ef 100644 --- a/.buildkite/test_areas/engine.yaml +++ b/.buildkite/test_areas/engine.yaml @@ -238,9 +238,9 @@ steps: - pytest -v -s v1/e2e/test_hybrid_chunked_prefill.py mirror: amd: - label: ":amd: (MI300) V1 E2E Hybrid Chunked Prefill" + label: ":amd: (MI355 DPX) V1 E2E Hybrid Chunked Prefill" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 35 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/lm_eval.yaml b/.buildkite/test_areas/lm_eval.yaml index af0b419de9..cb4e7eaedc 100644 --- a/.buildkite/test_areas/lm_eval.yaml +++ b/.buildkite/test_areas/lm_eval.yaml @@ -15,9 +15,9 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-small.txt mirror: amd: - label: ":amd: (MI300) LM Eval Small Models" + label: ":amd: (MI355) LM Eval Small Models" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 45 depends_on: - image-build-amd @@ -707,18 +707,17 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi355-nemotron-3-super.txt -- label: ":amd: (MI300) LM Eval Large Models" +- label: ":amd: (MI355) LM Eval Large Models" key: amd-lm-eval-large-models dind: false - device: mi300_8 + device: mi355_8 optional: true timeout_in_minutes: 65 depends_on: - image-build-amd commands: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - - export PYTORCH_ROCM_ARCH=gfx942 # Limit Quark compilation to save time - - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-mi3xx.txt + - pytest -s -v evals/gsm8k/test_gsm8k_correctness.py --config-list-file=configs/models-gfx950-large.txt source_file_dependencies: - vllm/_aiter_ops.py - vllm/platforms/rocm.py @@ -1298,9 +1297,9 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "nemotron-h-8b or gemma-4-e4b-it" mirror: amd: - label: ":amd: (MI300) LM Eval KV-Offload" + label: ":amd: (MI355 DPX) LM Eval KV-Offload" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 55 depends_on: - image-build-amd @@ -1360,9 +1359,9 @@ steps: - pytest -s -v evals/gsm8k/test_gsm8k_offloading.py -k "deepseek-v4-flash" mirror: amd: - label: ":amd: (MI300) LM Eval KV-Offload Large" + label: ":amd: (MI355) LM Eval KV-Offload Large" dind: false - device: mi300_4 + device: mi355_4 timeout_in_minutes: 65 depends_on: - image-build-amd @@ -1427,10 +1426,10 @@ steps: # Qwen2-57B has 28 attention heads and cannot run at TP8. - pytest -s -v test_lm_eval_correctness.py --config-list-file=configs/models-large.txt --tp-size=4 -- label: ":amd: (MI300) LM Eval Large Models ROCm Harness" +- label: ":amd: (MI355) LM Eval Large Models ROCm Harness" key: amd-lm-eval-large-models-rocm-harness dind: false - device: mi300_8 + device: mi355_8 optional: true timeout_in_minutes: 40 depends_on: diff --git a/.buildkite/test_areas/models_language.yaml b/.buildkite/test_areas/models_language.yaml index c30a25993c..be9f6927b3 100644 --- a/.buildkite/test_areas/models_language.yaml +++ b/.buildkite/test_areas/models_language.yaml @@ -216,9 +216,9 @@ steps: parallelism: 4 commands: - "pytest -v -s models/language/pooling -m 'not core_model' --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT --shard-id=$$BUILDKITE_PARALLEL_JOB" - label: ":amd: (MI300) Language Models (Extended Pooling) Shard %N" + label: ":amd: (MI355 DPX) Language Models (Extended Pooling) Shard %N" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 95 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/models_multimodal.yaml b/.buildkite/test_areas/models_multimodal.yaml index d281d16f7f..2cb2079d9a 100644 --- a/.buildkite/test_areas/models_multimodal.yaml +++ b/.buildkite/test_areas/models_multimodal.yaml @@ -163,9 +163,9 @@ steps: - pytest -v -s models/multimodal/test_mapping.py mirror: amd: - label: ":amd: (MI300) Multimodal Models (Extended Generation 1)" + label: ":amd: (MI355 DPX) Multimodal Models (Extended Generation 1)" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 100 depends_on: - image-build-amd @@ -182,9 +182,9 @@ steps: - pytest -v -s models/multimodal/generation_ppl_test/ mirror: amd: - label: ":amd: (MI300) Multimodal Models (PPL)" + label: ":amd: (MI355 DPX) Multimodal Models (PPL)" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 90 depends_on: - image-build-amd @@ -240,9 +240,9 @@ steps: - pytest -v -s models/multimodal/pooling -m 'not core_model' mirror: amd: - label: ":amd: (MI300) Multimodal Models (Extended Pooling)" + label: ":amd: (MI355 DPX) Multimodal Models (Extended Pooling)" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 75 depends_on: - image-build-amd diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index c2742b4c51..bfc81b7d7f 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -137,9 +137,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/speculators/ mirror: amd: - label: ":amd: (MI300) Spec Decode Speculators" + label: ":amd: (MI355 DPX) Spec Decode Speculators" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 40 depends_on: - image-build-amd @@ -167,9 +167,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/mtp/gemma4/ mirror: amd: - label: ":amd: (MI300) Spec Decode MTP Gemma4" + label: ":amd: (MI355 DPX) Spec Decode MTP Gemma4" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 40 depends_on: - image-build-amd @@ -198,9 +198,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/mtp/mimo/ mirror: amd: - label: ":amd: (MI300) Spec Decode MTP DeepSeek + MiMo" + label: ":amd: (MI355 DPX) Spec Decode MTP DeepSeek + MiMo" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 40 depends_on: - image-build-amd @@ -228,9 +228,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/mtp/qwen3_5/ mirror: amd: - label: ":amd: (MI300) Spec Decode MTP Qwen3.5" + label: ":amd: (MI355 DPX) Spec Decode MTP Qwen3.5" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 40 depends_on: - image-build-amd @@ -387,9 +387,9 @@ steps: - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test mirror: amd: - label: ":amd: (MI300) Speculators Correctness Nightly" + label: ":amd: (MI355 DPX) Speculators Correctness Nightly" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 40 depends_on: - image-build-amd @@ -452,9 +452,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" mirror: amd: - label: ":amd: (MI300) Spec Decode AL DFlash Nightly" + label: ":amd: (MI355 DPX) Spec Decode AL DFlash Nightly" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 120 commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ -k "not dflash2" @@ -531,9 +531,9 @@ steps: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ mirror: amd: - label: ":amd: (MI300) Spec Decode AL DSpark Nightly" + label: ":amd: (MI355 DPX) Spec Decode AL DSpark Nightly" dind: false - device: mi300_1 + device: mi355_dpx timeout_in_minutes: 90 depends_on: - image-build-amd diff --git a/tests/evals/gsm8k/configs/DeepSeek-R1-DP_MI325.yaml b/tests/evals/gsm8k/configs/DeepSeek-R1-DP_MI325.yaml deleted file mode 100644 index 0171cb4b19..0000000000 --- a/tests/evals/gsm8k/configs/DeepSeek-R1-DP_MI325.yaml +++ /dev/null @@ -1,12 +0,0 @@ -model_name: "deepseek-ai/DeepSeek-R1" -accuracy_threshold: 0.95 -num_questions: 1319 -num_fewshot: 5 -startup_max_wait_seconds: 1200 -server_args: >- - --enforce-eager - --max-model-len 4096 - --data-parallel-size 8 - --enable-expert-parallel - --attention-backend=TRITON_ATTN - --speculative-config '{"method":"mtp","num_speculative_tokens":3}' diff --git a/tests/evals/gsm8k/configs/DeepSeek-R1-TP_MI325.yaml b/tests/evals/gsm8k/configs/DeepSeek-R1-TP_MI325.yaml deleted file mode 100644 index ef92f574c7..0000000000 --- a/tests/evals/gsm8k/configs/DeepSeek-R1-TP_MI325.yaml +++ /dev/null @@ -1,12 +0,0 @@ -model_name: "deepseek-ai/DeepSeek-R1" -accuracy_threshold: 0.95 -num_questions: 1319 -num_fewshot: 5 -startup_max_wait_seconds: 1200 -server_args: >- - --enforce-eager - --max-model-len 4096 - --tensor-parallel-size 8 - --enable-expert-parallel - --attention-backend=TRITON_ATTN - --speculative-config '{"method":"mtp","num_speculative_tokens":3}' diff --git a/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI325.yaml b/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI325.yaml deleted file mode 100644 index 8d207878d4..0000000000 --- a/tests/evals/gsm8k/configs/DeepSeek-V3.2-DP_MI325.yaml +++ /dev/null @@ -1,12 +0,0 @@ -model_name: "deepseek-ai/DeepSeek-V3.2" -accuracy_threshold: 0.95 -num_questions: 1319 -num_fewshot: 5 -startup_max_wait_seconds: 1200 -server_args: >- - --enforce-eager - --max-model-len 4096 - --data-parallel-size 8 - --enable-expert-parallel - --attention-backend=TRITON_ATTN - --speculative-config '{"method":"mtp","num_speculative_tokens":3}' diff --git a/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI325.yaml b/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI325.yaml deleted file mode 100644 index 46853d3f5e..0000000000 --- a/tests/evals/gsm8k/configs/DeepSeek-V3.2-TP_MI325.yaml +++ /dev/null @@ -1,12 +0,0 @@ -model_name: "deepseek-ai/DeepSeek-V3.2" -accuracy_threshold: 0.95 -num_questions: 1319 -num_fewshot: 5 -startup_max_wait_seconds: 1200 -server_args: >- - --enforce-eager - --max-model-len 4096 - --tensor-parallel-size 8 - --enable-expert-parallel - --attention-backend=TRITON_ATTN - --speculative-config '{"method":"mtp","num_speculative_tokens":3}' diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16_MI355.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16_MI355.yaml index bf67ca1f7d..44fc3e3bb0 100644 --- a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16_MI355.yaml +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-BF16_MI355.yaml @@ -12,4 +12,5 @@ server_args: >- --enable-expert-parallel --mamba-backend triton --moe-backend triton + --attention-backend ROCM_AITER_FA --speculative-config '{"method":"mtp","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-FP8_MI355.yaml b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-FP8_MI355.yaml index 9d76e8f3ab..877b4e7e2f 100644 --- a/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-FP8_MI355.yaml +++ b/tests/evals/gsm8k/configs/Nemotron-3-Super-120B-A12B-FP8_MI355.yaml @@ -12,4 +12,5 @@ server_args: >- --enable-expert-parallel --mamba-backend triton --moe-backend triton + --attention-backend ROCM_AITER_FA --speculative-config '{"method":"mtp","num_speculative_tokens":5}' diff --git a/tests/evals/gsm8k/configs/models-mi3xx.txt b/tests/evals/gsm8k/configs/models-mi3xx.txt deleted file mode 100644 index 70d9be32e3..0000000000 --- a/tests/evals/gsm8k/configs/models-mi3xx.txt +++ /dev/null @@ -1,8 +0,0 @@ -DeepSeek-R1-TP_MI325.yaml -DeepSeek-R1-DP_MI325.yaml -DeepSeek-V3.2-TP_MI325.yaml -DeepSeek-V3.2-DP_MI325.yaml -Qwen3-30B-A3B-NVFP4.yaml -Qwen3-30B-A3B-NVFP4-quark.yaml -Qwen3.5-35B-A3B-MXFP4-AITER-TP2.yaml -Qwen3.5-35B-A3B-MXFP4-EMU-TP2.yaml From cffc0cfa311f4997be58cfe8f5f6473ab4898269 Mon Sep 17 00:00:00 2001 From: Misha Goin Date: Fri, 16 Oct 2026 12:00:12 -0400 Subject: [PATCH 11/49] [CI] Fix B200 FusedMoE Layer Kernels hang on NVSHMEM IB setup (#61695) Signed-off-by: mgoin Co-authored-by: Claude Opus 5.5 Co-authored-by: Kevin H. Luu --- .buildkite/test_areas/kernels.yaml | 2 +- .../moe/modular_kernel_tools/parallel_utils.py | 18 ++++++++++++++++-- tests/kernels/moe/test_moe_layer.py | 12 +++++++++++- 3 files changed, 28 insertions(+), 4 deletions(-) diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index 12b8516926..a1612bc3e8 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -939,7 +939,7 @@ steps: - label: ":nvidia: (B200) FusedMoE Layer Kernels" key: kernels-fusedmoe-layer-test-2-b200s - timeout_in_minutes: 90 + timeout_in_minutes: 45 device: b200-k8s num_devices: 2 source_file_dependencies: diff --git a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py index d7394d72cb..235d1b81a4 100644 --- a/tests/kernels/moe/modular_kernel_tools/parallel_utils.py +++ b/tests/kernels/moe/modular_kernel_tools/parallel_utils.py @@ -2,6 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import dataclasses import os +import time import traceback from collections.abc import Callable from typing import Any, Concatenate @@ -23,6 +24,8 @@ from vllm.utils.network_utils import get_open_port P = ParamSpec("P") +LAUNCH_TIMEOUT_S = 900 + @dataclasses.dataclass class ProcessGroupInfo: @@ -148,7 +151,7 @@ def parallel_launch_with_config( *args: P.args, **kwargs: P.kwargs, ) -> None: - spawn( + context = spawn( _worker_parallel_launch, args=( world_size, @@ -162,5 +165,16 @@ def parallel_launch_with_config( ) + args, nprocs=world_size, - join=True, + join=False, ) + # A worker stuck in communicator setup never exits, so bound the wait. + deadline = time.monotonic() + LAUNCH_TIMEOUT_S + while not context.join(timeout=max(0.0, deadline - time.monotonic())): + if time.monotonic() >= deadline: + for process in context.processes: + process.kill() + for process in context.processes: + process.join() + raise TimeoutError( + f"Parallel workers did not finish within {LAUNCH_TIMEOUT_S}s" + ) diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index ba79cc5c34..91a6bc3f70 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -127,6 +127,14 @@ if has_nixl_ep(): BACKENDS += ["nixl_ep"] DEEPEP_BACKENDS = {"deepep_high_throughput", "deepep_low_latency"} + +# These tests run on one node, where NVSHMEM reaches peers over NVLink, so +# keep it off the NICs to avoid depending on IB setup on shared CI hosts. +DEEPEP_SINGLE_NODE_NVSHMEM_ENV = { + "NVSHMEM_REMOTE_TRANSPORT": "none", + # DeepEP forces IBGDA on, so also hide the NICs from it. + "NVSHMEM_HCA_LIST": "none", +} MORI_BACKENDS = {"mori_high_throughput", "mori_low_latency"} QUANT_METHODS = [ @@ -2187,6 +2195,8 @@ def test_moe_layer( rocm_deepep = current_platform.is_rocm() and backend in DEEPEP_BACKENDS parallel_worker = _parallel_worker_rocm_deepep if rocm_deepep else _parallel_worker + cuda_deepep = current_platform.is_cuda() and backend in DEEPEP_BACKENDS + env_dict = DEEPEP_SINGLE_NODE_NVSHMEM_ENV if cuda_deepep else None launch_failures: list[str] = [] try: @@ -2197,7 +2207,7 @@ def test_moe_layer( world_size, parallel_worker, vllm_config, - None, + env_dict, test_config_batch, verbosity, failure_report_path=failure_report_path, From ae8f41238cc7316781ef4359890940749735b209 Mon Sep 17 00:00:00 2001 From: alberto Date: Fri, 16 Oct 2026 17:54:42 +0100 Subject: [PATCH 12/49] [Bugfix][Parser] Honor a prompt-closed reasoning block in non-streaming parse and reasoning token counts (#61857) Signed-off-by: Alberto Perdomo --- tests/parser/engine/test_qwen3_reasoning.py | 34 ++++ tests/parser/test_parse.py | 179 ++++++++++++++++++++ vllm/parser/abstract_parser.py | 5 + vllm/parser/engine/adapters.py | 8 + vllm/parser/engine/parser_engine.py | 20 ++- vllm/reasoning/abs_reasoning_parsers.py | 9 + 6 files changed, 253 insertions(+), 2 deletions(-) diff --git a/tests/parser/engine/test_qwen3_reasoning.py b/tests/parser/engine/test_qwen3_reasoning.py index 3f285e4d03..e9c751c3cb 100644 --- a/tests/parser/engine/test_qwen3_reasoning.py +++ b/tests/parser/engine/test_qwen3_reasoning.py @@ -380,6 +380,40 @@ class TestDelegatingPromptDetection: assert delta.reasoning is None assert delta.content == "answer" + def test_prior_turn_think_end_keeps_non_streaming_reasoning( + self, mock_tokenizer, mock_request + ): + parser = _Qwen3DelegatingParser(mock_tokenizer) + parser.set_prompt_token_ids([_THINK_END_ID, _IM_END_ID, _IM_START_ID, _TEXT_ID]) + output_ids = [ord("a"), _THINK_END_ID, ord("b")] + + reasoning, content, _ = parser.parse( + mock_tokenizer.decode(output_ids), + mock_request, + model_output_token_ids=output_ids, + ) + + assert reasoning == "a" + assert content == "b" + assert parser.count_reasoning_tokens(output_ids) > 0 + + def test_think_end_in_current_turn_parses_output_as_content( + self, mock_tokenizer, mock_request + ): + parser = _Qwen3DelegatingParser(mock_tokenizer) + parser.set_prompt_token_ids([_IM_START_ID, _THINK_START_ID, _THINK_END_ID]) + output_ids = [ord("a"), ord("b")] + + reasoning, content, _ = parser.parse( + mock_tokenizer.decode(output_ids), + mock_request, + model_output_token_ids=output_ids, + ) + + assert reasoning is None + assert content == "ab" + assert parser.count_reasoning_tokens(output_ids) == 0 + class TestStreaming: def test_basic_streaming(self, parser): diff --git a/tests/parser/test_parse.py b/tests/parser/test_parse.py index 2379b8a4c2..4eea8db73e 100644 --- a/tests/parser/test_parse.py +++ b/tests/parser/test_parse.py @@ -9,6 +9,9 @@ import pytest from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest from vllm.entrypoints.openai.responses.protocol import ResponsesRequest from vllm.parser.abstract_parser import DelegatingParser +from vllm.parser.engine.parser_engine import ParserEngine +from vllm.parser.parser_manager import ParserManager +from vllm.parser.qwen3 import qwen3_config from vllm.parser.utils import count_history_tool_calls from vllm.reasoning.basic_parsers import BaseThinkingReasoningParser from vllm.tool_parsers.hermes_tool_parser import Hermes2ProToolParser @@ -417,3 +420,179 @@ def test_parse_auto_tools_no_calls_returns_none(tokenizer): assert reasoning is None assert content == PLAIN_TEXT assert tool_calls is None + + +# -------------------------------------------------------------------------- +# Prompt that already closes the reasoning block +# -------------------------------------------------------------------------- + +CONTINUE_MESSAGES = [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "Sure,"}, +] +REASONED_OUTPUT = "\nlet me think\n\n\n" + PLAIN_TEXT +TOOL_TURN_MESSAGES = [ + {"role": "user", "content": "Weather in Paris?"}, + { + "role": "assistant", + "content": "\nneed a tool\n\n\n", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, +] + + +def reasoning_parser(tokenizer, name="qwen3"): + parser_cls = ParserManager.get_parser( + reasoning_parser_name=name, tokenizer=tokenizer + ) + return parser_cls(tokenizer, None) + + +def encode(tokenizer, text): + return tokenizer.encode(text, add_special_tokens=False) + + +def prompt_case(tokenizer, case): + """(rendered prompt ids, request, model output) for each prompt shape.""" + if case == "continue_final_message": + # The template renders an empty, closed think block before "Sure,". + prompt = tokenizer.apply_chat_template( + CONTINUE_MESSAGES, + tokenize=False, + add_generation_prompt=False, + continue_final_message=True, + ) + request = make_request( + messages=CONTINUE_MESSAGES, + add_generation_prompt=False, + continue_final_message=True, + ) + return encode(tokenizer, prompt), request, " here is the answer" + if case == "template_closed_thinking": + # The template switches thinking off, but the parser is not told + # through chat_template_kwargs (#53284). + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": "hi"}], + tokenize=False, + add_generation_prompt=True, + enable_thinking=False, + ) + return encode(tokenizer, prompt), make_request(), PLAIN_TEXT + if case == "earlier_turn_closed_reasoning": + # Only a previous assistant turn closes reasoning; this one is open. + prompt = tokenizer.apply_chat_template( + TOOL_TURN_MESSAGES, + tools=[{"type": "function", "function": {"name": "get_weather"}}], + tokenize=False, + add_generation_prompt=True, + ) + return encode(tokenizer, prompt), make_request(), REASONED_OUTPUT + assert case == "reasoning_open" + prompt = tokenizer.apply_chat_template( + [{"role": "user", "content": "hi"}], + tokenize=False, + add_generation_prompt=True, + ) + return encode(tokenizer, prompt), make_request(), REASONED_OUTPUT + + +@pytest.mark.parametrize("case", ["continue_final_message", "template_closed_thinking"]) +def test_parse_prompt_closed_reasoning_is_content(tokenizer, case): + prompt_ids, request, output = prompt_case(tokenizer, case) + parser = reasoning_parser(tokenizer) + parser.set_prompt_token_ids(prompt_ids) + + reasoning, content, _ = parser.parse( + output, request, model_output_token_ids=encode(tokenizer, output) + ) + + assert reasoning is None + assert content == output + assert parser.count_reasoning_tokens(encode(tokenizer, output)) == 0 + + +@pytest.mark.parametrize( + ("parser_name", "case"), + [ + ("qwen3", "reasoning_open"), + ("qwen3", "earlier_turn_closed_reasoning"), + # BaseThinkingReasoningParser scans the prompt without turn + # boundaries, so it must not decide non-streaming classification. + ("deepseek_r1", "earlier_turn_closed_reasoning"), + ], +) +def test_parse_prompt_open_reasoning_still_extracted(tokenizer, parser_name, case): + prompt_ids, request, output = prompt_case(tokenizer, case) + parser = reasoning_parser(tokenizer, parser_name) + parser.set_prompt_token_ids(prompt_ids) + + reasoning, content, _ = parser.parse( + output, request, model_output_token_ids=encode(tokenizer, output) + ) + + assert reasoning.strip() == "let me think" + assert content.strip() == PLAIN_TEXT + assert parser.count_reasoning_tokens(encode(tokenizer, output)) > 0 + + +def test_parse_prompt_ignored_without_turn_boundaries(tokenizer): + """Without turn boundaries, an earlier turn's looks like this + turn's, so the engine must not start the output in content.""" + prompt_ids, request, output = prompt_case( + tokenizer, "earlier_turn_closed_reasoning" + ) + parser = ParserEngine(tokenizer, parser_engine_config=qwen3_config(thinking=True)) + parser.set_prompt_token_ids(prompt_ids) + + reasoning, content, _ = parser.parse( + output, request, model_output_token_ids=encode(tokenizer, output) + ) + + assert reasoning.strip() == "let me think" + assert content.strip() == PLAIN_TEXT + + +@pytest.mark.parametrize( + "case", + [ + "continue_final_message", + "template_closed_thinking", + "reasoning_open", + "earlier_turn_closed_reasoning", + ], +) +def test_prompt_reasoning_streaming_matches_non_streaming(tokenizer, case): + prompt_ids, request, output = prompt_case(tokenizer, case) + output_ids = encode(tokenizer, output) + + parser = reasoning_parser(tokenizer) + parser.set_prompt_token_ids(prompt_ids) + reasoning, content, _ = parser.parse( + output, request, model_output_token_ids=output_ids + ) + non_streaming = (reasoning or "", content or "") + non_streaming_count = parser.count_reasoning_tokens(output_ids) + + parser = reasoning_parser(tokenizer) + streamed_reasoning = streamed_content = "" + for i, token_id in enumerate(output_ids): + delta = parser.parse_delta( + tokenizer.decode([token_id]), + [token_id], + request, + prompt_token_ids=prompt_ids, + finished=i == len(output_ids) - 1, + ) + if delta is not None: + streamed_reasoning += delta.reasoning or "" + streamed_content += delta.content or "" + + assert (streamed_reasoning, streamed_content) == non_streaming + assert parser.count_reasoning_tokens(output_ids) == non_streaming_count diff --git a/vllm/parser/abstract_parser.py b/vllm/parser/abstract_parser.py index fbddc45783..287cad5c24 100644 --- a/vllm/parser/abstract_parser.py +++ b/vllm/parser/abstract_parser.py @@ -769,6 +769,10 @@ class DelegatingParser(Parser): return False return self._reasoning_parser.is_reasoning_end(input_ids) + def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: + if self._reasoning_parser is not None: + self._reasoning_parser.set_prompt_token_ids(prompt_token_ids) + def _is_reasoning_end_streaming( self, input_ids: list[int], delta_ids: list[int] ) -> bool: @@ -858,6 +862,7 @@ class DelegatingParser(Parser): if not state.prompt_reasoning_checked and prompt_token_ids is not None: state.prompt_reasoning_checked = True + self.set_prompt_token_ids(prompt_token_ids) if self._reasoning_parser is None or self.is_reasoning_end( prompt_token_ids ): diff --git a/vllm/parser/engine/adapters.py b/vllm/parser/engine/adapters.py index a4a307a43a..52764e9fa8 100644 --- a/vllm/parser/engine/adapters.py +++ b/vllm/parser/engine/adapters.py @@ -68,6 +68,11 @@ class ParserEngineReasoningAdapter(ReasoningParser): def find_reasoning_end(self, token_ids: Sequence[int]) -> ReasoningEnd | None: return self._parser_engine.find_reasoning_end(token_ids) + def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: + self._parser_engine.set_prompt_token_ids(prompt_token_ids) + if self._counting_parser_engine is not None: + self._counting_parser_engine.set_prompt_token_ids(prompt_token_ids) + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: self._parser_engine.adjust_initial_state_from_prompt(prompt_token_ids) @@ -148,6 +153,9 @@ class ParserEngineReasoningAdapter(ReasoningParser): self._counting_parser_engine = self._parser_engine_cls( self.model_tokenizer, **self._parser_engine_kwargs ) # type: ignore[call-arg] + self._counting_parser_engine._prompt_closes_reasoning = ( + self._parser_engine._prompt_closes_reasoning + ) self._counting_parser_engine._single_pass_parse( self.model_tokenizer.decode(token_ids), token_ids ) diff --git a/vllm/parser/engine/parser_engine.py b/vllm/parser/engine/parser_engine.py index bbb3605ed7..de695f4bdf 100644 --- a/vllm/parser/engine/parser_engine.py +++ b/vllm/parser/engine/parser_engine.py @@ -147,6 +147,7 @@ class ParserEngine(Parser): or parser_engine_config.initial_state == ParserState.REASONING ) self._reasoning_ended: bool = not self._has_reasoning + self._prompt_closes_reasoning: bool = False self._streaming_initialized: bool = False self._prompt_streaming_prepared: bool = False @@ -231,6 +232,18 @@ class ParserEngine(Parser): self._streaming_initialized = True self._reset(initial_state=initial_state) + def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: + """Start parsing in content when the prompt already closes reasoning, + as with continue_final_message or a template that switches thinking + off. Only configs with turn boundaries can tell a closed block in + this turn from one in an earlier turn, so others are left as is. + """ + self._prompt_closes_reasoning = ( + self.parser_engine_config.initial_state is ParserState.REASONING + and bool(self._turn_boundary_token_ids) + and self.is_reasoning_end(list(prompt_token_ids)) + ) + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: """See :meth:`ReasoningParser.adjust_initial_state_from_prompt`.""" return @@ -242,8 +255,11 @@ class ParserEngine(Parser): return None def _reset(self, initial_state: ParserState | None = None) -> None: - self._engine.reset(initial_state=initial_state) - self._reasoning_ended = not self._has_reasoning + prompt_closed = initial_state is None and self._prompt_closes_reasoning + self._engine.reset( + initial_state=ParserState.CONTENT if prompt_closed else initial_state + ) + self._reasoning_ended = not self._has_reasoning or prompt_closed self._tool_slots.clear() self._deferred_content = "" self._deferred_reasoning = "" diff --git a/vllm/reasoning/abs_reasoning_parsers.py b/vllm/reasoning/abs_reasoning_parsers.py index accc8c9c27..f55cbfd686 100644 --- a/vllm/reasoning/abs_reasoning_parsers.py +++ b/vllm/reasoning/abs_reasoning_parsers.py @@ -192,6 +192,15 @@ class ReasoningParser: """Adjust request parameters; override in subclasses as needed.""" return request + def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: + """Provide the rendered prompt before the output is parsed. + + Parsers that can tell whether the prompt already closes reasoning + (e.g. continue_final_message) use it to classify the output. + Default is a no-op; override in subclasses as needed. + """ + return + def adjust_initial_state_from_prompt(self, prompt_token_ids: Sequence[int]) -> None: """Hook called once at the start of streaming with the prompt tokens. From e7a5bdb4542ca13e70a1ef0a1d9fb628b632c691 Mon Sep 17 00:00:00 2001 From: djramic Date: Fri, 16 Oct 2026 20:22:38 +0200 Subject: [PATCH 13/49] [ROCm][CI] Move Spec Decode Speculators + MTP to MI355 (#62053) Signed-off-by: Djordje Ramic --- .buildkite/test-amd.yaml | 63 ++++++++++++++-------------------------- 1 file changed, 21 insertions(+), 42 deletions(-) diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index 0b04b36878..a28227689d 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -399,27 +399,6 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ -- label: ":amd: (MI250) Spec Decode Speculators + MTP" # TBD - timeout_in_minutes: 75 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx90anightly, amdmi250] - dind: false - agent_pool: mi250_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - vllm/transformers_utils/configs/speculators/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - - vllm/v1/attention/backends/ - commands: - - pytest -v -s v1/e2e/spec_decode/speculators/ - - pytest -v -s v1/e2e/spec_decode/mtp/ - #------------------------------------------------------------- mi250 · misc ------------------------------------------------------------# - label: ":amd: (MI300) Python-only Installation" # TBD @@ -2378,27 +2357,6 @@ steps: - VLLM_TP_SIZE=4 VLLM_NEEDLE_TRIALS=2 pytest -v -s v1/determinism/test_batch_invariance.py::test_v1_generation_is_deterministic_across_batch_sizes_with_needle -k TRITON_ATTN - VLLM_TEST_MODEL=allenai/OLMoE-1B-7B-0924 VLLM_TEST_SERVER_ARGS="--data-parallel-size 4 --enable-expert-parallel --max-model-len 4096" pytest -v -s v1/determinism/test_online_batch_invariance.py -- label: ":amd: (MI300) Spec Decode Speculators + MTP" # TBD - timeout_in_minutes: 75 - mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] - dind: false - agent_pool: mi300_1 - optional: true - working_dir: "/vllm-workspace/tests" - source_file_dependencies: - - vllm/v1/spec_decode/ - - vllm/v1/worker/gpu/spec_decode/ - - vllm/model_executor/model_loader/ - - vllm/v1/sample/ - - vllm/model_executor/layers/ - - vllm/transformers_utils/configs/speculators/ - - tests/v1/e2e/spec_decode/ - - vllm/platforms/rocm.py - - vllm/v1/attention/backends/ - commands: - - pytest -v -s v1/e2e/spec_decode/speculators/ - - pytest -v -s v1/e2e/spec_decode/mtp/ - - label: ":amd: (MI300) Spec Decode DeepSeek MTP Parallel Load" timeout_in_minutes: 45 mirror_hardwares: [amdexperimental, amdproduction, amdgfx942nightly, amdmi300] @@ -5396,6 +5354,27 @@ steps: commands: - pytest -v -s v1/e2e/spec_decode/draft_model/ +- label: ":amd: (MI355 DPX) Spec Decode Speculators + MTP" # TBD + timeout_in_minutes: 75 + mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] + dind: false + agent_pool: mi355_dpx + optional: true + working_dir: "/vllm-workspace/tests" + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - vllm/transformers_utils/configs/speculators/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + - vllm/v1/attention/backends/ + commands: + - pytest -v -s v1/e2e/spec_decode/speculators/ + - pytest -v -s v1/e2e/spec_decode/mtp/ + - label: ":amd: (MI355 DPX) Spec Decode Eagle 1: DeepSeek + Qwen" timeout_in_minutes: 60 mirror_hardwares: [amdexperimental, amdproduction, amdgfx950nightly, amdmi355] From 3dc6def36088a4139182e1a79cfe9b4e5953aa8b Mon Sep 17 00:00:00 2001 From: Guan-Ming Chiu <105915352+guan404ming@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:23:59 +0100 Subject: [PATCH 14/49] [Attention][Spec Decode] Enable varlen Mamba2 decode for adaptive verification (#53036) Signed-off-by: Guan-Ming (Wesley) Chiu <105915352+guan404ming@users.noreply.github.com> Co-authored-by: vllm-agent --- tests/kernels/mamba/test_causal_conv1d.py | 112 ++++++++++++++++++++++ vllm/v1/attention/backends/mamba2_attn.py | 21 +++- 2 files changed, 131 insertions(+), 2 deletions(-) diff --git a/tests/kernels/mamba/test_causal_conv1d.py b/tests/kernels/mamba/test_causal_conv1d.py index ddd532ad13..06031842bd 100644 --- a/tests/kernels/mamba/test_causal_conv1d.py +++ b/tests/kernels/mamba/test_causal_conv1d.py @@ -2,6 +2,8 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from collections.abc import Iterator + import pytest import torch import torch.nn.functional as F @@ -358,3 +360,113 @@ def test_causal_conv1d_varlen( ) unpadded_out = out[:, : out_ref_tensor.shape[-1]] assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol) + + +@pytest.mark.parametrize("itype", [torch.bfloat16]) +@pytest.mark.parametrize("silu_activation", [True]) +@pytest.mark.parametrize("has_bias", [True]) +@pytest.mark.parametrize("width", [4]) +@pytest.mark.parametrize("dim", [64, 2048]) +@pytest.mark.parametrize("q_lens_cfg", [[3, 1, 0, 4, 2], [1, 5, 5, 1], [2, 0, 0, 3]]) +@pytest.mark.skipif( + current_platform.is_cpu(), + reason="CPU kernel for causal_conv1d_update ignores num_accepted_tokens.", +) +def test_causal_conv1d_update_ragged_varlen( + q_lens_cfg, dim, width, has_bias, silu_activation, itype +): + mk_inputs = lambda: _create_ragged_varlen_inputs( + q_lens_cfg, dim, width, has_bias, silu_activation, itype + ) + inputs, state = mk_inputs() + actual = causal_conv1d_update(conv_state=state, **inputs) + + ref_inputs, ref_state = mk_inputs() + ref_outs = _create_ragged_varlen_ref_outs(conv_state=ref_state, **ref_inputs) + + for i, ref_out in enumerate(ref_outs): + start = int(inputs["query_start_loc"][i]) + end = int(inputs["query_start_loc"][i + 1]) + entry = inputs["conv_state_indices"][i : i + 1] + assert torch.allclose(actual[start:end], ref_out, rtol=1e-2, atol=5e-2) + assert torch.equal(state[entry], ref_state[entry]) + + +def _create_ragged_varlen_inputs( + q_lens_cfg, dim, width, has_bias, silu_activation, itype +): + device = DEVICE + set_random_seed(0) + + batch_size = len(q_lens_cfg) + max_query_len = 5 + assert max(q_lens_cfg) <= max_query_len + state_len = width - 1 + max_query_len + total_entries = 10 * batch_size + + q_lens = torch.tensor(q_lens_cfg, dtype=torch.int32, device=device) + query_start_loc = torch.cat( + [ + torch.zeros(1, dtype=torch.int32, device=device), + torch.cumsum(q_lens, dim=0, dtype=torch.int32), + ] + ) + total_tokens = int(q_lens.sum()) + + x = torch.randn(total_tokens, dim, device=device, dtype=itype) + conv_state = torch.randn( + total_entries, state_len, dim, device=device, dtype=itype + ).transpose(1, 2) + weight = torch.randn(dim, width, device=device, dtype=itype) + bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None + activation = None if not silu_activation else "silu" + + conv_state_indices = ( + torch.randperm(total_entries - 1, device=device)[:batch_size] + 1 + ).to(torch.int32) + num_accepted_tokens = torch.randint( + 1, max_query_len + 1, (batch_size,), dtype=torch.int32, device=device + ) + inputs = dict( + x=x, + weight=weight, + bias=bias, + activation=activation, + conv_state_indices=conv_state_indices, + num_accepted_tokens=num_accepted_tokens, + query_start_loc=query_start_loc, + max_query_len=max_query_len, + ) + return inputs, conv_state + + +def _create_ragged_varlen_ref_outs( + x, + conv_state, + weight, + bias, + activation, + conv_state_indices, + num_accepted_tokens, + query_start_loc, + max_query_len, +) -> Iterator[torch.Tensor]: + for seq_idx in range(len(conv_state_indices)): + start = int(query_start_loc[seq_idx]) + end = int(query_start_loc[seq_idx + 1]) + entry = conv_state_indices[seq_idx : seq_idx + 1] + if start == end: + yield x[start:end] + continue + out_ref = causal_conv1d_update( + x[start:end], + conv_state, + weight, + bias, + activation=activation, + conv_state_indices=entry, + num_accepted_tokens=num_accepted_tokens[seq_idx : seq_idx + 1], + query_start_loc=query_start_loc.new_tensor([0, end - start]), + max_query_len=max_query_len, + ) + yield out_ref diff --git a/vllm/v1/attention/backends/mamba2_attn.py b/vllm/v1/attention/backends/mamba2_attn.py index fc81cf680e..1b9e573182 100644 --- a/vllm/v1/attention/backends/mamba2_attn.py +++ b/vllm/v1/attention/backends/mamba2_attn.py @@ -6,7 +6,7 @@ from typing import Any import torch -from vllm.config import VllmConfig +from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.model_executor.layers.mamba.checkpoint import ( MambaPrefillCheckpointBuilder, MambaPrefillCheckpointMetadata, @@ -15,12 +15,13 @@ from vllm.utils.torch_utils import async_tensor_h2d from vllm.v1.attention.backend import ( AttentionBackend, CommonAttentionMetadata, + max_decode_query_len, ) from vllm.v1.attention.backends.mamba_attn import ( BaseMambaAttentionMetadata, BaseMambaAttentionMetadataBuilder, ) -from vllm.v1.kv_cache_interface import MambaSpec +from vllm.v1.kv_cache_interface import KVCacheSpec, MambaSpec def compute_varlen_chunk_metadata( @@ -104,6 +105,14 @@ class Mamba2AttentionBackend(AttentionBackend): def is_ssm(cls) -> bool: return True + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + # Decode kernels take device offsets; ReplaySSM plans off CPU counts. + vllm_config = get_current_vllm_config_or_none() + if vllm_config is None: + return True + return not vllm_config.cache_config.use_replayssm + @dataclass class Mamba2AttentionMetadata(BaseMambaAttentionMetadata): @@ -124,6 +133,14 @@ class Mamba2AttentionMetadataBuilder( ): metadata_cls = Mamba2AttentionMetadata + @classmethod + def get_varlen_cudagraph_max_query_len( + cls, + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + return max_decode_query_len(vllm_config) + def __init__( self, kv_cache_spec: MambaSpec, From 307c5844ad2d3003154f3a768540e686e37dddec Mon Sep 17 00:00:00 2001 From: djramic Date: Fri, 16 Oct 2026 20:36:14 +0200 Subject: [PATCH 15/49] [ROCm][CI] Enable batch invariance for the ngram scheduling accuracy test (#62048) Signed-off-by: Djordje Ramic --- tests/v1/e2e/general/test_async_scheduling_accuracy.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/tests/v1/e2e/general/test_async_scheduling_accuracy.py b/tests/v1/e2e/general/test_async_scheduling_accuracy.py index 3ffcaefdc3..249f2c5ef7 100644 --- a/tests/v1/e2e/general/test_async_scheduling_accuracy.py +++ b/tests/v1/e2e/general/test_async_scheduling_accuracy.py @@ -12,6 +12,7 @@ from typing import Any import pytest import torch +import vllm.envs as envs from tests.conftest import HfRunner, VllmRunner from tests.utils import ( large_gpu_mark, @@ -168,9 +169,16 @@ def test_with_eagle3_spec_decoding( reason=("XPU matmul/attention kernels are not batch-invariant"), ) def test_with_ngram_gpu_spec_decoding( - vllm_runner: type[VllmRunner], hf_runner: type[HfRunner] + vllm_runner: type[VllmRunner], + hf_runner: type[HfRunner], + monkeypatch: pytest.MonkeyPatch, ): """Check ngram accuracy and acceptance across scheduling configurations.""" + if current_platform.is_rocm(): + # hipBLASLt picks GEMM solutions by batch shape, so sync and async runs + # can flip near-tie greedy tokens and shift ngram acceptance. + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + monkeypatch.setattr(envs, "VLLM_BATCH_INVARIANT", True) ngram_gpu_config = { "method": "ngram_gpu", "num_speculative_tokens": 3, From 5bf3d9f6653e53485f0ca74149f958dc6cbd3b82 Mon Sep 17 00:00:00 2001 From: Oxana Korzh Date: Fri, 16 Oct 2026 12:53:49 -0600 Subject: [PATCH 16/49] [ROCm][Kimi-K3] Support DSpark adaptive verification on MI355X (#59824) Signed-off-by: Oxana Korzh Co-authored-by: Claude Opus 5.5 (1M context) --- .../test_amd_kda_adaptive_verification.py | 93 +++++++++++ .../v1/attention/test_expand_page_indices.py | 3 + tests/v1/attention/test_mla_backends.py | 84 +++++++++- .../test_rocm_aiter_mla_mtp_split.py | 32 ++++ vllm/models/kimi_k3/amd/kda_metadata.py | 29 ++++ .../attention/backends/mla/rocm_aiter_mla.py | 155 ++++++++++++++++-- 6 files changed, 383 insertions(+), 13 deletions(-) create mode 100644 tests/models/kimi_k3/test_amd_kda_adaptive_verification.py diff --git a/tests/models/kimi_k3/test_amd_kda_adaptive_verification.py b/tests/models/kimi_k3/test_amd_kda_adaptive_verification.py new file mode 100644 index 0000000000..9cb4307e56 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_kda_adaptive_verification.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""ROCm Kimi-K3 KDA metadata under adaptive verification. + +Adaptive verification splits the draft budget evenly on the host, then trims +each verify request on device. The KDA spec path must therefore plan from the +device offsets, never from host lengths that merely look uniform. +""" + +import torch + +from tests.v1.attention.utils import ( + BatchSpec, + create_common_attn_metadata, + create_vllm_config, +) +from vllm.config import SpeculativeConfig +from vllm.config.compilation import CUDAGraphMode +from vllm.models.kimi_k3.amd.kda_metadata import ( + KimiK3ROCmKDABackend, + KimiK3ROCmKDAMetadataBuilder, +) +from vllm.v1.kv_cache_interface import MambaSpec + +BLOCK_SIZE = 16 +DEVICE = torch.device("cpu") +NUM_SPEC = 3 + + +def _config(adaptive: bool): + vllm_config = create_vllm_config( + model_name="Qwen/Qwen3.5-0.8B", + block_size=BLOCK_SIZE, + ) + vllm_config.compilation_config.cudagraph_mode = CUDAGraphMode.FULL_AND_PIECEWISE + vllm_config.speculative_config = SpeculativeConfig( + method="ngram", + num_speculative_tokens=NUM_SPEC, + enable_adaptive_verification=adaptive, + ) + vllm_config.cache_config.mamba_cache_mode = "none" + return vllm_config + + +def _spec() -> MambaSpec: + return MambaSpec( + block_size=BLOCK_SIZE, + shapes=((16, 64),), + dtypes=(torch.float16,), + num_speculative_blocks=NUM_SPEC, + ) + + +def _builder(adaptive: bool) -> KimiK3ROCmKDAMetadataBuilder: + return KimiK3ROCmKDAMetadataBuilder( + kv_cache_spec=_spec(), + layer_names=["layer.0"], + vllm_config=_config(adaptive), + device=DEVICE, + ) + + +def _trimmed_verify_batch(builder: KimiK3ROCmKDAMetadataBuilder): + """Host sees [3, 3]; the device was trimmed to [4, 2] (same total).""" + batch = BatchSpec(seq_lens=[64, 64], query_lens=[4, 2]) + common = create_common_attn_metadata( + batch, BLOCK_SIZE, DEVICE, arange_block_indices=True + ) + common = common.replace( + query_start_loc_cpu=torch.tensor([0, 3, 6], dtype=torch.int32) + ) + accepted = torch.ones(2, dtype=torch.int32, device=DEVICE) + num_decode_draft_tokens_cpu = torch.tensor([NUM_SPEC, NUM_SPEC], dtype=torch.int32) + return builder.build(0, common, accepted, num_decode_draft_tokens_cpu) + + +def test_adaptive_verification_plans_from_device_offsets(): + """Equal host lengths would select the fixed-length recurrent kernel, + which places sequence i at i * L and ignores the device cu_seqlens.""" + assert KimiK3ROCmKDABackend.supports_device_cpu_query_lens_mismatch() + bound = KimiK3ROCmKDAMetadataBuilder.get_varlen_cudagraph_max_query_len + assert bound(_config(adaptive=True), _spec()) == NUM_SPEC + 1 + assert bound(_config(adaptive=False), _spec()) is None + + meta = _trimmed_verify_batch(_builder(adaptive=True)) + assert meta.num_spec_decodes == 2 + assert meta.uniform_spec_sequence_length is None + assert torch.equal( + meta.spec_query_start_loc.cpu(), torch.tensor([0, 4, 6], dtype=torch.int32) + ) + # Without adaptive verification the host lengths are exact and kept. + meta = _trimmed_verify_batch(_builder(adaptive=False)) + assert meta.uniform_spec_sequence_length == 3 diff --git a/tests/v1/attention/test_expand_page_indices.py b/tests/v1/attention/test_expand_page_indices.py index 73163fe13d..1fc078cf56 100644 --- a/tests/v1/attention/test_expand_page_indices.py +++ b/tests/v1/attention/test_expand_page_indices.py @@ -51,6 +51,7 @@ def _run(block_table, cu_num_tokens, kernel_block_size, out_len, device): block_table.stride(0), block_table.stride(1), cu_num_tokens, + cu_num_tokens[1:] - cu_num_tokens[:-1], KERNEL_BLOCK_SIZE=kernel_block_size, BLOCK_SIZE=1024, ) @@ -109,6 +110,7 @@ def test_ragged_batch_leaves_other_rows_untouched(): bt.stride(0), bt.stride(1), cu, + cu[1:] - cu[:-1], KERNEL_BLOCK_SIZE=kbs, BLOCK_SIZE=1024, ) @@ -153,6 +155,7 @@ def test_over_provisioned_chunks_do_not_change_result(): bt.stride(0), bt.stride(1), cu, + cu[1:] - cu[:-1], KERNEL_BLOCK_SIZE=kbs, BLOCK_SIZE=1024, ) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index b696e35b77..891aa8880d 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -418,6 +418,11 @@ for backend in BACKENDS_TO_TEST: if query_len_support != QueryLenSupport.SINGLE_ONLY: SPEC_DECODE_BACKENDS.append(backend) +# Backends whose decode pipeline replays device-trimmed verify batches. +ADAPTIVE_VARLEN_BACKENDS = [ + b for b in BACKENDS_TO_TEST if b == AttentionBackendEnum.ROCM_AITER_MLA +] + BACKEND_BLOCK_SIZES = {} for backend in BACKENDS_TO_TEST: supported_sizes = backend.get_class().get_supported_kernel_block_sizes() @@ -470,6 +475,10 @@ BATCH_SPECS = { "spec_decode_small": BatchSpec( seq_lens=[128, 256, 512, 1024], query_lens=[4, 4, 4, 4] ), + # Adaptive verification: verify requests trimmed to 1..k+1 tokens on device. + "spec_decode_ragged_small": BatchSpec( + seq_lens=[128, 256, 512, 1024], query_lens=[8, 3, 1, 5] + ), "spec_decode_medium": BatchSpec( seq_lens=[512, 1024, 2048, 512, 1024, 2048], query_lens=[8, 8, 8, 8, 8, 8] ), @@ -1581,6 +1590,34 @@ def run_attention_backend( return output +def _with_even_cpu_query_split( + common_attn_metadata: CommonAttentionMetadata, +) -> CommonAttentionMetadata: + """Host query lengths as adaptive verification leaves them. + + The device keeps the trimmed per-request lengths; the host sees the same + total spread evenly, the first requests taking the remainder + (AdaptiveVerificationManager.compact_batch). + """ + num_reqs = common_attn_metadata.num_reqs + total = int(common_attn_metadata.query_start_loc_cpu[num_reqs]) + cpu_lens = torch.full((num_reqs,), total // num_reqs, dtype=torch.int32) + cpu_lens[: total % num_reqs] += 1 + query_start_loc_cpu = torch.zeros_like(common_attn_metadata.query_start_loc_cpu) + torch.cumsum(cpu_lens, dim=0, out=query_start_loc_cpu[1 : num_reqs + 1]) + return common_attn_metadata.replace( + query_start_loc_cpu=query_start_loc_cpu, + max_query_len=int( + ( + common_attn_metadata.query_start_loc[1:] + - common_attn_metadata.query_start_loc[:-1] + ) + .max() + .item() + ), + ) + + def _run_backend_correctness( default_vllm_config, dist_init, @@ -1595,6 +1632,7 @@ def _run_backend_correctness( qk_nope_head_dim: int, v_head_dim: int, chunked_prefill_workspace_size: int | None = None, + adaptive_verification: bool = False, ): """Test that all backends produce similar outputs to a reference implementation using torch.nn.functional.scaled_dot_product_attention. @@ -1626,6 +1664,10 @@ def _run_backend_correctness( ) and AttentionBackendEnum.CUTLASS_MLA in backends_to_test: # CUTLASS_MLA does not support non-1 Q/K scales backends_to_test.remove(AttentionBackendEnum.CUTLASS_MLA) + if adaptive_verification: + backends_to_test = [ + b for b in backends_to_test if b in ADAPTIVE_VARLEN_BACKENDS + ] if not backends_to_test: pytest.skip(f"No backends support kv_cache_dtype={kv_cache_dtype}") @@ -1675,13 +1717,16 @@ def _run_backend_correctness( if is_spec_decode_test: from vllm.config import SpeculativeConfig - # Get the query length from the batch spec (they should all be uniform) - query_len = batch_spec.query_lens[0] + # Get the query length from the batch spec (they should all be uniform, + # except under adaptive verification, where the longest is k+1) + query_len = max(batch_spec.query_lens) # Set num_speculative_tokens to query_len - 1 # (since threshold is 1 + num_spec_tokens) # Use ngram method which doesn't require a draft model vllm_config.speculative_config = SpeculativeConfig( - method="ngram", num_speculative_tokens=query_len - 1 + method="ngram", + num_speculative_tokens=query_len - 1, + enable_adaptive_verification=adaptive_verification, ) device = torch.device(f"{DEVICE_TYPE}:0") @@ -1951,6 +1996,13 @@ def _run_backend_correctness( ) kv_cache_per_block_size[block_size] = kv_cache + if adaptive_verification: + # Only after the cache is filled: prepopulation reads the context + # lengths off the host query lengths. + metadata_per_block_size[block_size] = _with_even_cpu_query_split( + common_attn_metadata + ) + # 4. Run vLLM backends and compare rtol = 1e-2 atol = { @@ -2136,3 +2188,29 @@ def test_chunked_context_backend_correctness( v_head_dim=v_head_dim, chunked_prefill_workspace_size=1024, ) + + +@pytest.mark.skipif( + not ADAPTIVE_VARLEN_BACKENDS, reason="No adaptive varlen decode backend" +) +def test_adaptive_varlen_decode_correctness( + default_vllm_config, dist_init, workspace_init +): + """Device query lengths 1..k+1 against an even host split, as adaptive + verification hands them to the target model's attention. TP 8 gives 16 + heads, the padded Kimi-K3 count.""" + _run_backend_correctness( + default_vllm_config, + dist_init, + workspace_init, + "spec_decode_ragged_small", + "deepseek-ai/DeepSeek-R1", + 8, + "fp8", + 1.0, + 1.0, + None, + 128, + 128, + adaptive_verification=True, + ) diff --git a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py index 187bbea4bf..fdc89f4bbe 100644 --- a/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py +++ b/tests/v1/attention/test_rocm_aiter_mla_mtp_split.py @@ -63,9 +63,11 @@ class _ExpandPageIndicesKernel: stride_0, stride_1, paged_kv_indptr, + num_real_tokens, *, KERNEL_BLOCK_SIZE, BLOCK_SIZE, + CLAMP_TAIL=False, ): # Both block-table strides are passed explicitly: the kernel no longer # assumes the column stride is 1. This stand-in indexes the tensor @@ -145,6 +147,7 @@ def _builder( ) ), _mtp_decode_qlen=mtp_decode_qlen, + _varlen_decode_qlen=0, _uniform_padded_mtp_qo_len=(AiterMLAMetadataBuilder._uniform_padded_mtp_qo_len), _use_persistent_metadata=False, kernel_block_size=kernel_block_size, @@ -500,6 +503,7 @@ def test_non_causal_dcp_block_bypasses_cprr_in_forward(monkeypatch): impl.qk_rope_head_dim = 64 impl.scale = 576**-0.5 decode = SimpleNamespace( + varlen_q_rows=None, max_qo_len=4, qo_indptr=torch.tensor([0, 4], dtype=torch.int32), paged_kv_indptr=torch.tensor([0, 1], dtype=torch.int32), @@ -563,6 +567,7 @@ def test_single_token_dcp_decode_returns_unpadded_lse(monkeypatch): impl.qk_rope_head_dim = 64 impl.scale = head_dim**-0.5 decode = SimpleNamespace( + varlen_q_rows=None, max_qo_len=1, qo_indptr=torch.arange(num_tokens + 1, dtype=torch.int32), paged_kv_indptr=torch.zeros(num_tokens + 1, dtype=torch.int32), @@ -787,6 +792,7 @@ def test_segmented_dcp_verify_matches_causal_attention(monkeypatch): torch.tensor([global_seq_len], dtype=torch.int32, device=device), ) decode = SimpleNamespace( + varlen_q_rows=None, max_qo_len=qlen, paged_kv_indptr=torch.tensor( [0, local_kv.shape[0]], dtype=torch.int32, device=device @@ -1452,3 +1458,29 @@ def test_a_two_token_bf16_block_is_allowed(monkeypatch): monkeypatch, num_heads=12, kv_cache_dtype="auto", qlen=2, mtp_qlen=8 ) assert metadata.has_persistent_metadata + + +def _adaptive_config( + *, adaptive: bool = True, dcp: int = 1, num_speculative_tokens: int = 7 +) -> SimpleNamespace: + return SimpleNamespace( + speculative_config=SimpleNamespace( + num_speculative_tokens=num_speculative_tokens, + enable_adaptive_verification=adaptive, + parallel_drafting=False, + ), + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp), + use_v2_model_runner=True, + ) + + +def test_varlen_cudagraph_bound_follows_adaptive_verification(): + bound = AiterMLAMetadataBuilder.get_varlen_cudagraph_max_query_len + assert bound(_adaptive_config(), SimpleNamespace()) == 8 + assert bound(_adaptive_config(adaptive=False), SimpleNamespace()) is None + # DCP routes plan per-row windows from host lengths. + assert bound(_adaptive_config(dcp=2), SimpleNamespace()) is None + # Kimi-K3 DSpark draft layers share the target's group; causality is + # decided per batch, so the merged non-causal spec keeps the bound. + shared_draft_group = SimpleNamespace(non_causal_multi_token_decode=True) + assert bound(_adaptive_config(), shared_draft_group) == 8 diff --git a/vllm/models/kimi_k3/amd/kda_metadata.py b/vllm/models/kimi_k3/amd/kda_metadata.py index dcc03d2814..90887c0491 100644 --- a/vllm/models/kimi_k3/amd/kda_metadata.py +++ b/vllm/models/kimi_k3/amd/kda_metadata.py @@ -11,6 +11,7 @@ from dataclasses import replace import torch +from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.layers.mamba.checkpoint import ( MambaPrefillCheckpointMetadata, @@ -18,12 +19,14 @@ from vllm.model_executor.layers.mamba.checkpoint import ( from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE from vllm.triton_utils import tl, triton from vllm.utils.math_utils import next_power_of_2 +from vllm.v1.attention.backend import max_decode_query_len from vllm.v1.attention.backends.gdn_attn import ( GDNAttentionBackend, GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.v1.kv_cache_interface import KVCacheSpec logger = init_logger(__name__) @@ -92,7 +95,24 @@ def prepare_chunk_metadata_device( return chunk_indices, chunk_offsets +def _adaptive_verification_enabled(vllm_config: VllmConfig) -> bool: + speculative_config = vllm_config.speculative_config + return speculative_config is not None and bool( + getattr(speculative_config, "enable_adaptive_verification", False) + ) + + class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder): + @classmethod + def get_varlen_cudagraph_max_query_len( + cls: type["KimiK3ROCmKDAMetadataBuilder"], + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + if not _adaptive_verification_enabled(vllm_config): + return None + return max_decode_query_len(vllm_config) + def build( # type: ignore[override] self, common_prefix_len: int, @@ -108,6 +128,11 @@ class KimiK3ROCmKDAMetadataBuilder(GDNAttentionMetadataBuilder): num_decode_draft_tokens_cpu, fast_build, ) + if ( + attn_metadata.uniform_spec_sequence_length is not None + and _adaptive_verification_enabled(self.vllm_config) + ): + attn_metadata = replace(attn_metadata, uniform_spec_sequence_length=None) if attn_metadata.checkpoint is None: return attn_metadata # The shared builder keys checkpoint rows on every request; the ROCm @@ -163,3 +188,7 @@ class KimiK3ROCmKDABackend(GDNAttentionBackend): @staticmethod def get_builder_cls() -> type[KimiK3ROCmKDAMetadataBuilder]: return KimiK3ROCmKDAMetadataBuilder + + @classmethod + def supports_device_cpu_query_lens_mismatch(cls) -> bool: + return True diff --git a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py index e485f5462a..bb29e55826 100644 --- a/vllm/v1/attention/backends/mla/rocm_aiter_mla.py +++ b/vllm/v1/attention/backends/mla/rocm_aiter_mla.py @@ -35,6 +35,7 @@ from vllm.v1.attention.backend import ( AttentionLayer, CommonAttentionMetadata, MultipleOf, + max_decode_query_len, ) from vllm.v1.attention.backends.mla.triton_mla import _compute_num_kv_splits from vllm.v1.attention.backends.utils import ( @@ -49,7 +50,11 @@ from vllm.v1.attention.ops.rocm_aiter_mla_prefill import ( gather_compressed_context, ) from vllm.v1.attention.ops.triton_decode_attention import decode_attention_fwd -from vllm.v1.kv_cache_interface import AttentionSpec, is_quantized_kv_cache +from vllm.v1.kv_cache_interface import ( + AttentionSpec, + KVCacheSpec, + is_quantized_kv_cache, +) from vllm.v1.worker.workspace import ( current_workspace_manager, is_workspace_manager_initialized, @@ -457,6 +462,16 @@ def _aiter_mla_small_head_mode() -> str: return mode +def _adaptive_varlen_verify_enabled(vllm_config: VllmConfig) -> bool: + speculative_config = vllm_config.speculative_config + return ( + speculative_config is not None + and bool(getattr(speculative_config, "enable_adaptive_verification", False)) + and vllm_config.parallel_config.decode_context_parallel_size == 1 + and max_decode_query_len(vllm_config) > 1 + ) + + class AiterMLABackend(MLACommonBackend): supported_dtypes: ClassVar[list[torch.dtype]] = [torch.float16, torch.bfloat16] supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [ @@ -620,6 +635,8 @@ class AiterMLADecodeMetadata(MLACommonDecodeMetadata): use_gluon_verify: bool = False # Whether persistent MLA metadata was computed has_persistent_metadata: bool = False + varlen_q_rows: torch.Tensor | None = None + varlen_o_rows: torch.Tensor | None = None @dataclass @@ -720,6 +737,21 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): return uniform_qo_len + @classmethod + def get_varlen_cudagraph_max_query_len( + cls: type["AiterMLAMetadataBuilder"], + vllm_config: VllmConfig, + kv_cache_spec: KVCacheSpec, + ) -> int | None: + if not _adaptive_varlen_verify_enabled(vllm_config): + return None + if ( + cls.get_cudagraph_support(vllm_config, kv_cache_spec) + != AttentionCGSupport.UNIFORM_BATCH + ): + return None + return max_decode_query_len(vllm_config) + def __init__( self, kv_cache_spec: AttentionSpec, @@ -823,6 +855,26 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # persistent gate below and makes aiter raise a KeyError mid-run. self._mtp_decode_qlen = self.reorder_batch_threshold or 1 + self._varlen_decode_qlen = 0 + if _adaptive_varlen_verify_enabled(vllm_config): + assert self._mtp_decode_qlen == max_decode_query_len(vllm_config), ( + self._mtp_decode_qlen, + max_decode_query_len(vllm_config), + ) + self._varlen_decode_qlen = self._mtp_decode_qlen + self.query_len_support = QueryLenSupport.VARLEN # type: ignore[misc] + max_rows = vllm_config.scheduler_config.max_num_seqs * ( + self._varlen_decode_qlen + ) + row_ids = torch.arange(max_rows, dtype=torch.int64, device=device) + self._varlen_row_req = row_ids // self._varlen_decode_qlen + self._varlen_row_pos = row_ids % self._varlen_decode_qlen + self._varlen_row_ids = row_ids + self._varlen_q_rows = torch.zeros_like(row_ids) + self._varlen_o_rows = torch.zeros( + max_rows + 1, dtype=torch.int64, device=device + ) + # Store the kernel block size from the spec. When kernel_block_size=1 # (no spec-dec), behavior is identical to the original. When > 1 # (e.g. 16 with Eagle3), we expand block-level indices into per-token @@ -850,6 +902,12 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): max_num_pages_per_req = vllm_config.model_config.max_model_len max_num_reqs = vllm_config.scheduler_config.max_num_seqs max_num_pages = max_num_reqs * max_num_pages_per_req + speculative_config = vllm_config.speculative_config + if _adaptive_varlen_verify_enabled(vllm_config): + assert speculative_config is not None + max_num_pages += max_num_reqs * ( + speculative_config.num_speculative_tokens or 0 + ) # Preparing persistent buffers # TODO: we can disambiguate between decode and mixed-prefill decode here @@ -1437,13 +1495,18 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): device = self.device num_reqs = seq_lens_device.size(0) qo_len = query_start_loc_cpu[1:] - query_start_loc_cpu[:-1] - max_qo_len = qo_len.max().item() - padded_mtp_qo_len = self._uniform_padded_mtp_qo_len( - qo_len, max_qo_len, num_decode_tokens - ) - if padded_mtp_qo_len > 0: - max_qo_len = padded_mtp_qo_len - pad_uniform_mtp = padded_mtp_qo_len > 0 + varlen_decode = self._varlen_decode_qlen > 0 and causal + if varlen_decode: + max_qo_len = self._varlen_decode_qlen + pad_uniform_mtp = True + else: + max_qo_len = qo_len.max().item() + padded_mtp_qo_len = self._uniform_padded_mtp_qo_len( + qo_len, max_qo_len, num_decode_tokens + ) + if padded_mtp_qo_len > 0: + max_qo_len = padded_mtp_qo_len + pad_uniform_mtp = padded_mtp_qo_len > 0 seq_lens_for_kernel = seq_lens_device # the global lengths need the same dummy-row @@ -1477,6 +1540,13 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): g_tot_seq_lens.new_full((), max_qo_len * self.dcp_world_size), ) + varlen_q_rows = varlen_o_rows = None + if varlen_decode: + seq_lens_for_kernel = seq_lens_device + (max_qo_len - qo_lens_device) + varlen_q_rows, varlen_o_rows = self._fill_varlen_rows( + query_start_loc_device, qo_lens_device, num_reqs + ) + if self._graph_seq_lens is not None: self._graph_seq_lens[:num_kernel_reqs].copy_( seq_lens_for_kernel, non_blocking=True @@ -1537,6 +1607,8 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): # device synchronisation. Programs whose chunk lies past # num_tokens mask out entirely. max_tokens_per_req = block_table_tensor.shape[1] * self.kernel_block_size + if varlen_decode: + max_tokens_per_req += max_qo_len - 1 num_chunks = max(1, cdiv(max_tokens_per_req, 1024)) _expand_page_indices_kernel[(num_reqs, num_chunks)]( self.paged_kv_indices, @@ -1544,8 +1616,10 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): block_table_tensor.stride(0), block_table_tensor.stride(1), paged_kv_indptr, + seq_lens_device, KERNEL_BLOCK_SIZE=self.kernel_block_size, BLOCK_SIZE=1024, + CLAMP_TAIL=varlen_decode, ) paged_kv_indices = self.paged_kv_indices @@ -1740,10 +1814,38 @@ class AiterMLAMetadataBuilder(MLACommonMetadataBuilder[AiterMLAMetadata]): use_gluon_verify=use_gluon_verify, attn_out_dtype=self.decode_attn_out_dtype, has_persistent_metadata=has_persistent_metadata, + varlen_q_rows=varlen_q_rows, + varlen_o_rows=varlen_o_rows, ) return attn_metadata + def _fill_varlen_rows( + self, + query_start_loc_device: torch.Tensor, + qo_lens_device: torch.Tensor, + num_reqs: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + width = self._varlen_decode_qlen + num_rows = num_reqs * width + row_req = self._varlen_row_req[:num_rows] + row_pos = self._varlen_row_pos[:num_rows] + row_qlen = qo_lens_device.to(torch.int64).index_select(0, row_req) + row_start = query_start_loc_device.to(torch.int64).index_select(0, row_req) + is_real = row_pos < row_qlen + packed_row = row_start + row_pos + q_rows = self._varlen_q_rows[:num_rows] + q_rows.copy_(torch.where(is_real, packed_row, 0)) + o_rows = self._varlen_o_rows + o_rows.zero_() + dump_slot = o_rows.size(0) - 1 + o_rows.scatter_( + 0, + torch.where(is_real, packed_row, dump_slot), + self._varlen_row_ids[:num_rows], + ) + return q_rows, o_rows + def build( self, common_prefix_len: int, @@ -1795,8 +1897,10 @@ def _expand_page_indices_kernel( block_table_stride_0, block_table_stride_1, cu_num_tokens, + num_real_tokens, KERNEL_BLOCK_SIZE: tl.constexpr, BLOCK_SIZE: tl.constexpr, + CLAMP_TAIL: tl.constexpr = False, ): """Expand block table entries into per-token flat page indices. @@ -1810,6 +1914,10 @@ def _expand_page_indices_kernel( When KERNEL_BLOCK_SIZE=K: block table entry b (covering K tokens) is expanded to flat indices b*K, b*K+1, ..., b*K+(K-1). + + With CLAMP_TAIL, entries past num_real_tokens[req] (padding a request's + page list out to its kernel length) repeat the request's first page, so + they never index past its block-table row. """ # One program per (request, token-chunk). Parallelising over requests # alone degenerates at low concurrency: with num_reqs == 1 a single @@ -1819,6 +1927,8 @@ def _expand_page_indices_kernel( row_ptr = block_table + req_idx * block_table_stride_0 start_idx = tl.load(cu_num_tokens + req_idx) num_tokens = tl.load(cu_num_tokens + req_idx + 1) - start_idx + if CLAMP_TAIL: + real_tokens = tl.load(num_real_tokens + req_idx) # The grid is sized from the block table width, an upper bound over all # requests, so a ragged batch launches chunks past a short request's end. @@ -1830,11 +1940,14 @@ def _expand_page_indices_kernel( token_offsets = chunk_start + tl.arange(0, BLOCK_SIZE) mask = token_offsets < num_tokens + source_offsets = token_offsets + if CLAMP_TAIL: + source_offsets = tl.where(token_offsets < real_tokens, token_offsets, 0) # Which block in the block table does this token belong to? - block_idx = token_offsets // KERNEL_BLOCK_SIZE + block_idx = source_offsets // KERNEL_BLOCK_SIZE # Offset within that block - offset_in_block = token_offsets % KERNEL_BLOCK_SIZE + offset_in_block = source_offsets % KERNEL_BLOCK_SIZE # Load the block ID from the block table # Both strides are taken from the caller: the block table is a view owned @@ -2600,6 +2713,28 @@ class AiterMLAImpl(MLACommonImpl[AiterMLAMetadata]): kv_c_and_k_pe_cache: torch.Tensor, attn_metadata: AiterMLAMetadata, layer: AttentionLayer, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + decode = attn_metadata.decode + if decode is None or decode.varlen_q_rows is None: + return self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) + assert decode.varlen_o_rows is not None + q_rows = decode.varlen_q_rows + if isinstance(q, tuple): + num_tokens = q[0].shape[0] + q = (q[0].index_select(0, q_rows), q[1].index_select(0, q_rows)) + else: + num_tokens = q.shape[0] + q = q.index_select(0, q_rows) + o, lse = self._forward_mqa(q, kv_c_and_k_pe_cache, attn_metadata, layer) + assert lse is None + return o.index_select(0, decode.varlen_o_rows[:num_tokens]), None + + def _forward_mqa( + self, + q: torch.Tensor | tuple[torch.Tensor, torch.Tensor], + kv_c_and_k_pe_cache: torch.Tensor, + attn_metadata: AiterMLAMetadata, + layer: AttentionLayer, ) -> tuple[torch.Tensor, torch.Tensor | None]: assert kv_c_and_k_pe_cache.numel() > 0 assert attn_metadata.decode is not None From 663e5356404332d6c764180aef41cc07015cd34a Mon Sep 17 00:00:00 2001 From: cjackal <44624812+cjackal@users.noreply.github.com> Date: Sat, 17 Oct 2026 04:42:48 +0900 Subject: [PATCH 17/49] [Perf][Spec Decode] Skip bonus-logit reduction for short drafts in rejection sampling (#61451) Signed-off-by: cjackal <44624812+cjackal@users.noreply.github.com> --- .../test_rejection_sampler_utils.py | 85 +++++++++++++++++++ .../spec_decode/rejection_sampler_utils.py | 30 ++++--- 2 files changed, 105 insertions(+), 10 deletions(-) diff --git a/tests/v1/spec_decode/test_rejection_sampler_utils.py b/tests/v1/spec_decode/test_rejection_sampler_utils.py index ab94ec54d2..b93cd2b306 100644 --- a/tests/v1/spec_decode/test_rejection_sampler_utils.py +++ b/tests/v1/spec_decode/test_rejection_sampler_utils.py @@ -1157,6 +1157,91 @@ def test_chunked_requests_match_full_batch(has_draft_logits: bool): assert torch.equal(chunked_sampled[valid], sampled[valid]) +@pytest.mark.parametrize("use_block_verification", [False, True]) +@pytest.mark.parametrize("has_draft_logits", [False, True]) +def test_adaptive_draft_lengths_match_and_bonus_oracle( + use_block_verification: bool, has_draft_logits: bool +): + torch.manual_seed(7) + device = "cuda" + K = 3 + num_drafts = [3, 1, 0, 2, 3, 2] + temperatures = [1.0, 0.0, 0.0, 0.6, 0.0, 1.0] + greedy_num_matches = {1: 1, 2: 0, 4: 2} + num_reqs = len(num_drafts) + + num_logits_per_req = torch.tensor(num_drafts, device=device) + 1 + cu_num_logits = torch.zeros(num_reqs + 1, dtype=torch.int32, device=device) + cu_num_logits[1:] = num_logits_per_req.cumsum(0) + num_logits = int(cu_num_logits[-1]) + idx_mapping = torch.arange(num_reqs, dtype=torch.int32, device=device) + expanded_idx_mapping = idx_mapping.repeat_interleave(num_logits_per_req) + expanded_local_pos = ( + torch.arange(num_logits, dtype=torch.int32, device=device) + - cu_num_logits[expanded_idx_mapping] + ) + + draft_logits = torch.randn(num_reqs, K, VOCAB_SIZE, device=device) + target_logits = torch.randn(num_logits, VOCAB_SIZE, device=device) + # Allocate at exactly num_logits so a lookahead past the batch is OOB. + draft_sampled = torch.randint(VOCAB_SIZE, (num_logits,), device=device) + target_argmax = target_logits.argmax(dim=-1) + for req_idx, num_matches in greedy_num_matches.items(): + lo = int(cu_num_logits[req_idx]) + for step in range(num_drafts[req_idx]): + token = target_argmax[lo + step] + if step >= num_matches: + token = (token + 1) % VOCAB_SIZE + draft_sampled[lo + step + 1] = token + + inputs = dict( + target_logits=target_logits, + draft_logits=draft_logits if has_draft_logits else None, + draft_sampled=draft_sampled, + cu_num_logits=cu_num_logits, + pos=torch.arange(num_logits, dtype=torch.int32, device=device), + idx_mapping=idx_mapping, + expanded_idx_mapping=expanded_idx_mapping, + expanded_local_pos=expanded_local_pos, + temperature=torch.tensor(temperatures, device=device), + seed=torch.arange(num_reqs, dtype=torch.int64, device=device), + num_speculative_steps=K, + use_block_verification=use_block_verification, + ) + sampled, num_sampled = rejection_sample(**inputs) + + for req_idx, num_matches in greedy_num_matches.items(): + lo = int(cu_num_logits[req_idx]) + n = num_matches + 1 + assert int(num_sampled[req_idx]) == n, f"greedy request {req_idx}" + assert torch.equal(sampled[req_idx, :n], target_argmax[lo : lo + n]), ( + f"greedy request {req_idx}" + ) + + for req_idx in range(num_reqs): + lo, hi = int(cu_num_logits[req_idx]), int(cu_num_logits[req_idx + 1]) + req_inputs = dict(inputs) + for name in ( + "target_logits", + "draft_sampled", + "pos", + "expanded_idx_mapping", + "expanded_local_pos", + ): + req_inputs[name] = inputs[name][lo:hi] + req_inputs["cu_num_logits"] = cu_num_logits[req_idx : req_idx + 2] - lo + req_inputs["idx_mapping"] = idx_mapping[req_idx : req_idx + 1] + + req_sampled, req_num_sampled = rejection_sample(**req_inputs) + + n = int(req_num_sampled[0]) + assert n <= num_drafts[req_idx] + 1 + assert int(num_sampled[req_idx]) == n, f"request {req_idx} mismatch" + assert torch.equal(sampled[req_idx, :n], req_sampled[0, :n]), ( + f"request {req_idx} mismatch" + ) + + def _residual_recovery_logits( inputs: dict, num_sampled: torch.Tensor, diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py index f35468d62b..060c64caa4 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler_utils.py @@ -194,6 +194,16 @@ def _compute_global_logprobs_and_logsumexp( return target_log_prob, draft_log_prob, target_lse, draft_lse +@triton.jit +def _is_bonus_logit(expanded_local_pos_ptr, logit_idx, num_logits): + next_local_pos = tl.load( + expanded_local_pos_ptr + logit_idx + 1, + mask=logit_idx + 1 < num_logits, + other=0, + ) + return next_local_pos == 0 + + @triton.jit def _compute_local_logits_stats_kernel( # [num_logits, num_blocks] @@ -225,17 +235,17 @@ def _compute_local_logits_stats_kernel( # [max_num_reqs] temp_ptr, vocab_size, - num_speculative_steps, + num_logits, BLOCK_SIZE: tl.constexpr, HAS_DRAFT_LOGITS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) - draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) - if draft_step_idx >= num_speculative_steps: - # Bonus token. Max/argmax and summed exponentials are not needed. + if _is_bonus_logit(expanded_local_pos_ptr, logit_idx, num_logits): + # Max/argmax and summed exponentials are not needed. return + draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) req_state_idx = tl.load(expanded_idx_mapping_ptr + logit_idx).to(tl.int64) temp = tl.load(temp_ptr + req_state_idx).to(tl.float32) @@ -423,17 +433,17 @@ def _compute_local_residual_mass_kernel( # [max_num_reqs] temp_ptr, vocab_size, - num_speculative_steps, + num_logits, vocab_num_blocks, BLOCK_SIZE: tl.constexpr, PADDED_VOCAB_NUM_BLOCKS: tl.constexpr, ): logit_idx = tl.program_id(0).to(tl.int64) # Adaptive verification can end before num_speculative_steps. - if logit_idx + 1 >= tl.num_programs(0): - return draft_step_idx = tl.load(expanded_local_pos_ptr + logit_idx) - if draft_step_idx == 0 or draft_step_idx >= num_speculative_steps: + if draft_step_idx == 0 or _is_bonus_logit( + expanded_local_pos_ptr, logit_idx, num_logits + ): # The acceptance threshold, h, looks one position ahead and sums # over: max(p_i * M_b(x|x_{ Date: Fri, 16 Oct 2026 15:57:49 -0400 Subject: [PATCH 18/49] [Build] Remove unused `tools/install_deepgemm.sh` (#61265) Signed-off-by: Tyler Michael Smith Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: Kevin H. Luu --- .buildkite/test_areas/kernels.yaml | 1 - cmake/external_projects/deepgemm.cmake | 1 - tools/install_deepgemm.sh | 131 ------------------------- 3 files changed, 133 deletions(-) delete mode 100755 tools/install_deepgemm.sh diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index a1612bc3e8..aa2124e33f 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -545,7 +545,6 @@ steps: num_devices: 1 source_file_dependencies: - cmake/external_projects/deepgemm.cmake - - tools/install_deepgemm.sh - tools/build_deepgemm_C.py - tools/setup_deepgemm_pythons.sh - tools/check_wheel_deepgemm.py diff --git a/cmake/external_projects/deepgemm.cmake b/cmake/external_projects/deepgemm.cmake index 26eb5ffe6b..f92321db42 100644 --- a/cmake/external_projects/deepgemm.cmake +++ b/cmake/external_projects/deepgemm.cmake @@ -27,7 +27,6 @@ if(DEEPGEMM_SRC_DIR) set(deepgemm_SOURCE_DIR "${DEEPGEMM_SRC_DIR}") message(STATUS "DeepGEMM using local DEEPGEMM_SRC_DIR: ${deepgemm_SOURCE_DIR}") else() - # Keep in sync with tools/install_deepgemm.sh set(_DEEPGEMM_UPSTREAM_REPO "https://github.com/vllm-project/DeepGEMM.git") # Pinned to the fork's dev branch: upstream's 26/09/30 release (2.8.1) # plus the SM120 port, the SM90 paged-MQA kv_block=32/next_n=4 port, diff --git a/tools/install_deepgemm.sh b/tools/install_deepgemm.sh deleted file mode 100755 index b3f7ceb43f..0000000000 --- a/tools/install_deepgemm.sh +++ /dev/null @@ -1,131 +0,0 @@ -#!/bin/bash -# Script to build and/or install DeepGEMM from source -# Default: build and install immediately -# Optional: build wheels to a directory for later installation (useful in multi-stage builds) -set -e - -# Default values -# Keep DEEPGEMM_GIT_REF in sync with cmake/external_projects/deepgemm.cmake -DEEPGEMM_GIT_REPO="https://github.com/vllm-project/DeepGEMM.git" -# NOTE: This targets the vLLM fork's dev branch: upstream's 26/09/30 release -# plus the sm120 and sm90 paged-MQA ports, the SwiGLU alpha/beta and SiTU -# Mega MoE activations, NVFP4 Mega MoE, and BF16/FP8/NVFP4 swiglustep. -# Uses the PyTorch 2.11 stable ABI and CPython 3.10 limited API. -DEEPGEMM_GIT_REF="a7a927731b109fdfec770df7c2a5be5d7e2216b5" -WHEEL_DIR="" - -# Parse command line arguments -while [[ $# -gt 0 ]]; do - case $1 in - --ref) - if [[ -z "$2" || "$2" =~ ^- ]]; then - echo "Error: --ref requires an argument." >&2 - exit 1 - fi - DEEPGEMM_GIT_REF="$2" - shift 2 - ;; - --cuda-version) - if [[ -z "$2" || "$2" =~ ^- ]]; then - echo "Error: --cuda-version requires an argument." >&2 - exit 1 - fi - CUDA_VERSION="$2" - shift 2 - ;; - --wheel-dir) - if [[ -z "$2" || "$2" =~ ^- ]]; then - echo "Error: --wheel-dir requires a directory path." >&2 - exit 1 - fi - WHEEL_DIR="$2" - shift 2 - ;; - -h|--help) - echo "Usage: $0 [OPTIONS]" - echo "Options:" - echo " --ref REF Git reference to checkout (default: $DEEPGEMM_GIT_REF)" - echo " --cuda-version VER CUDA version (auto-detected if not provided)" - echo " --wheel-dir PATH If set, build wheel into PATH but do not install" - echo " -h, --help Show this help message" - exit 0 - ;; - *) - echo "Unknown option: $1" >&2 - exit 1 - ;; - esac -done - -# Auto-detect CUDA version if not provided -if [ -z "$CUDA_VERSION" ]; then - if command -v nvcc >/dev/null 2>&1; then - CUDA_VERSION=$(nvcc --version | grep "release" | sed -n 's/.*release \([0-9]\+\.[0-9]\+\).*/\1/p') - echo "Auto-detected CUDA version: $CUDA_VERSION" - else - echo "Warning: Could not auto-detect CUDA version. Please specify with --cuda-version" - exit 1 - fi -fi - -# Extract major and minor version numbers -CUDA_MAJOR="${CUDA_VERSION%%.*}" -CUDA_MINOR="${CUDA_VERSION#"${CUDA_MAJOR}".}" -CUDA_MINOR="${CUDA_MINOR%%.*}" -echo "CUDA version: $CUDA_VERSION (major: $CUDA_MAJOR, minor: $CUDA_MINOR)" - -# Check CUDA version requirement -if [ "$CUDA_MAJOR" -lt 12 ] || { [ "$CUDA_MAJOR" -eq 12 ] && [ "$CUDA_MINOR" -lt 8 ]; }; then - echo "Skipping DeepGEMM build/installation (requires CUDA 12.8+ but got ${CUDA_VERSION})" - exit 0 -fi - -echo "Preparing DeepGEMM build..." -echo "Repository: $DEEPGEMM_GIT_REPO" -echo "Reference: $DEEPGEMM_GIT_REF" - -# Create a temporary directory for the build -INSTALL_DIR=$(mktemp -d) -trap 'rm -rf "$INSTALL_DIR"' EXIT - -# Clone the repository -git clone --recursive --shallow-submodules "$DEEPGEMM_GIT_REPO" "$INSTALL_DIR/deepgemm" -pushd "$INSTALL_DIR/deepgemm" - -# Checkout the specific reference -git checkout "$DEEPGEMM_GIT_REF" -git submodule sync --recursive -git submodule update --init --recursive - -# Clean previous build artifacts -# (Based on https://github.com/deepseek-ai/DeepGEMM/blob/main/install.sh) -rm -rf -- build dist *.egg-info 2>/dev/null || true - -# Build wheel -echo "🏗️ Building DeepGEMM wheel..." -python3 setup.py bdist_wheel - -# If --wheel-dir was specified, copy wheels there and exit -if [ -n "$WHEEL_DIR" ]; then - mkdir -p "$WHEEL_DIR" - cp dist/*.whl "$WHEEL_DIR"/ - echo "✅ Wheel built and copied to $WHEEL_DIR" - popd - exit 0 -fi - -# Default behaviour: install built wheel -if command -v uv >/dev/null 2>&1; then - echo "Installing DeepGEMM wheel using uv..." - if [ -n "$VLLM_DOCKER_BUILD_CONTEXT" ]; then - uv pip install --system dist/*.whl - else - uv pip install dist/*.whl - fi -else - echo "Installing DeepGEMM wheel using pip..." - python3 -m pip install dist/*.whl -fi - -popd -echo "✅ DeepGEMM installation completed successfully" From cc2e4fecf2fca603939445faccf20dd466f1502e Mon Sep 17 00:00:00 2001 From: afriedri Date: Fri, 16 Oct 2026 15:19:07 -0500 Subject: [PATCH 19/49] [ROCm] [fusion] support ROCm aiter fusion of per tensor quant + rmsnorm (#52831) Signed-off-by: Andy Friedrich Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- tests/compile/passes/test_fusion.py | 16 ++ vllm/_aiter_ops.py | 94 ++++++++++++ .../passes/fusion/matcher_utils.py | 25 ++- .../passes/fusion/rocm_aiter_fusion.py | 145 ++++++++++++++++++ 4 files changed, 274 insertions(+), 6 deletions(-) diff --git a/tests/compile/passes/test_fusion.py b/tests/compile/passes/test_fusion.py index 934522a1b9..418d21f425 100644 --- a/tests/compile/passes/test_fusion.py +++ b/tests/compile/passes/test_fusion.py @@ -107,6 +107,8 @@ AITER_KERNEL_GROUPSHAPE_COMBINATIONS = [ # Per-token with ChannelWiseTorchFP8ScaledMMLinearKernel (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, True), (ChannelWiseTorchFP8ScaledMMLinearKernel, GroupShape.PER_TOKEN, False), + # Static per-tensor: only the aiter quant op is matched by the pass + (ROCmFP8ScaledMMLinearKernel, GroupShape.PER_TENSOR, True), # Blockwise (AiterFp8BlockScaledMMKernel, GroupShape(1, 128), True), ] @@ -197,6 +199,9 @@ class TestModel(torch.nn.Module): # Blockwise path if self.use_aiter_fusion and self.use_aiter_quant_op: return [rocm_aiter_ops.get_group_quant_op()] + elif self.group_shape.is_per_tensor(): + if self.use_aiter_quant_op: + return [rocm_aiter_ops.get_per_tensor_quant_op()] else: if self.use_aiter_quant_op: return [rocm_aiter_ops.get_per_token_quant_op()] @@ -221,6 +226,17 @@ class TestModel(torch.nn.Module): AiterFusedAddRMSFp8GroupQuantPattern.FUSED_OP, AiterRMSFp8GroupQuantPattern.FUSED_OP, ] + elif self.group_shape.is_per_tensor(): + # Static per-tensor aiter fusion + from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( + AiterFusedAddRMSNormStaticQuantPattern, + AiterRMSNormStaticQuantPattern, + ) + + return [ + AiterFusedAddRMSNormStaticQuantPattern.FUSED_OP, + AiterRMSNormStaticQuantPattern.FUSED_OP, + ] else: # Per-token aiter fusion from vllm.compilation.passes.fusion.rocm_aiter_fusion import ( diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index fadb053915..6a6c99840a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -1492,6 +1492,74 @@ def _rocm_aiter_rmsnorm_fp8_group_quant_fake( ) +def _rocm_aiter_fused_rms_fp8_per_tensor_static_quant_impl( + x: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + scale: torch.Tensor, + quant_dtype: torch.dtype, +) -> torch.Tensor: + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_tensor_static_quant + + x_quant, _, _, _ = fused_rms_fp8_per_tensor_static_quant( + x, + weight, + epsilon, + scale, + None, + None, + None, + dtype_quant=quant_dtype, + res1=None, + ) + return x_quant + + +def _rocm_aiter_fused_rms_fp8_per_tensor_static_quant_fake( + x: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + scale: torch.Tensor, + quant_dtype: torch.dtype, +) -> torch.Tensor: + return torch.empty_like(x, dtype=quant_dtype) + + +def _rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_impl( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + scale: torch.Tensor, + quant_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + from aiter.ops.triton.fused_fp8_quant import fused_rms_fp8_per_tensor_static_quant + + x_quant, _, _, residual_out = fused_rms_fp8_per_tensor_static_quant( + x, + weight, + epsilon, + scale, + None, + None, + None, + dtype_quant=quant_dtype, + res1=residual, + ) + return x_quant, residual_out + + +def _rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_fake( + x: torch.Tensor, + residual: torch.Tensor, + weight: torch.Tensor, + epsilon: float, + scale: torch.Tensor, + quant_dtype: torch.dtype, +) -> tuple[torch.Tensor, torch.Tensor]: + return torch.empty_like(x, dtype=quant_dtype), torch.empty_like(residual) + + def _rocm_aiter_fused_rms_gated_fp8_group_quant_impl( x: torch.Tensor, weight: torch.Tensor, @@ -3050,6 +3118,18 @@ class rocm_aiter_ops: fake_impl=_rocm_aiter_rmsnorm_fp8_group_quant_fake, ) + direct_register_custom_op( + op_name="rocm_aiter_fused_rms_fp8_per_tensor_static_quant", + op_func=_rocm_aiter_fused_rms_fp8_per_tensor_static_quant_impl, + fake_impl=_rocm_aiter_fused_rms_fp8_per_tensor_static_quant_fake, + ) + + direct_register_custom_op( + op_name="rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant", + op_func=_rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_impl, + fake_impl=_rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant_fake, + ) + direct_register_custom_op( op_name="rocm_aiter_fused_rms_gated_fp8_group_quant", op_func=_rocm_aiter_fused_rms_gated_fp8_group_quant_impl, @@ -3179,6 +3259,16 @@ class rocm_aiter_ops: def get_rmsnorm_group_fused_quant_op() -> OpOverload: return torch.ops.vllm.rocm_aiter_rmsnorm_fp8_group_quant.default + @staticmethod + def get_fused_rms_fp8_per_tensor_static_quant_op() -> OpOverload: + return torch.ops.vllm.rocm_aiter_fused_rms_fp8_per_tensor_static_quant.default + + @staticmethod + def get_fused_add_rms_fp8_per_tensor_static_quant_op() -> OpOverload: + return ( + torch.ops.vllm.rocm_aiter_fused_add_rms_fp8_per_tensor_static_quant.default + ) + @staticmethod def get_fused_rms_gated_fp8_group_quant_op() -> OpOverload: """Return the fused gated-RMSNorm + FP8 group quant custom op.""" @@ -3192,6 +3282,10 @@ class rocm_aiter_ops: def get_per_token_quant_op() -> OpOverload: return torch.ops.vllm.rocm_aiter_per_token_quant.default + @staticmethod + def get_per_tensor_quant_op() -> OpOverload: + return torch.ops.vllm.rocm_aiter_per_tensor_quant.default + @staticmethod def get_group_quant_op() -> OpOverload: return torch.ops.vllm.rocm_aiter_group_fp8_quant.default diff --git a/vllm/compilation/passes/fusion/matcher_utils.py b/vllm/compilation/passes/fusion/matcher_utils.py index 80be65a978..32cb0e35aa 100644 --- a/vllm/compilation/passes/fusion/matcher_utils.py +++ b/vllm/compilation/passes/fusion/matcher_utils.py @@ -394,10 +394,13 @@ class MatcherQuantFP8(MatcherCustomOp): self.is_tma_aligned = is_tma_aligned if match_rocm_aiter: - assert not quant_key.scale.group_shape.is_per_tensor(), ( - "ROCm aiter fusion pass does not support per tensor quantization" - ) - if quant_key.scale.group_shape.is_per_token(): + if quant_key.scale.group_shape.is_per_tensor(): + assert quant_key.scale.static, ( + "ROCm aiter fusion pass only supports static per-tensor " + "quantization" + ) + self.QUANT_OP = rocm_aiter_ops.get_per_tensor_quant_op() + elif quant_key.scale.group_shape.is_per_token(): self.QUANT_OP = rocm_aiter_ops.get_per_token_quant_op() else: assert quant_key.scale.group_shape.col == 128, ( @@ -432,14 +435,24 @@ class MatcherQuantFP8(MatcherCustomOp): scale: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: quant_key_group_shape = self.quant_key.scale.group_shape + if quant_key_group_shape == GroupShape.PER_TENSOR: + # rocm_aiter_per_tensor_quant mutates `out` and `scale` and returns + # None, so it has to be matched through auto_functionalized. + assert scale is not None + out = torch.empty_like(input, dtype=self.quant_key.dtype) + result = auto_functionalized( + self.QUANT_OP, out=out, x=input, scale=scale, is_dynamic=False + ) + # `scale` is mutable, so its functionalized output is a distinct + # value; return it to keep the quant node inside the match. + return result[1], result[2] if quant_key_group_shape == GroupShape.PER_TOKEN: return self.QUANT_OP( # type: ignore[no-any-return] x=input, quant_dtype=self.quant_key.dtype, scale=scale, ) - else: - return self.QUANT_OP(input, quant_key_group_shape.col) # type: ignore[no-any-return] + return self.QUANT_OP(input, quant_key_group_shape.col) # type: ignore[no-any-return] def forward_custom( self, diff --git a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py index 849cd5d67d..16515f24bc 100644 --- a/vllm/compilation/passes/fusion/rocm_aiter_fusion.py +++ b/vllm/compilation/passes/fusion/rocm_aiter_fusion.py @@ -19,6 +19,7 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import ( QuantKey, ScaleDesc, kFp8Dynamic128Sym, + kStaticTensorScale, ) from vllm.platforms import current_platform @@ -561,6 +562,140 @@ class AiterRMSNormGatedFp8GroupQuantPattern(AiterRMSNormQuantPattern): ) +class AiterRMSNormStaticQuantPattern(AiterRMSNormQuantPattern): + """This pattern fuses rms_norm & aiter static per-tensor fp8 quant + into an aiter fused_rms_fp8_per_tensor_static_quant op. + """ + + FUSED_OP = rocm_aiter_ops.get_fused_rms_fp8_per_tensor_static_quant_op() + + def __init__( + self, + epsilon: float, + quant_dtype: torch.dtype, + symmetric: bool = True, + ) -> None: + key = FusedRMSQuantKey( + fused_add=False, + quant=QuantKey( + dtype=quant_dtype, scale=kStaticTensorScale, symmetric=symmetric + ), + ) + + super().__init__(epsilon, key) + + def register(self, pm_pass: PatternMatcherPass) -> None: + def pattern( + input: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + result_rms = torch.ops.vllm_ir.rms_norm(input, weight, self.epsilon) + result, scale = self.quant_matcher(result_rms, scale) + + return result, scale + + def replacement( + input: torch.Tensor, + weight: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + result = self.FUSED_OP( + x=input, + weight=weight, + epsilon=self.epsilon, + scale=scale, + quant_dtype=self.quant_dtype, + ) + + # The static scale is an input, so it passes through unchanged. + return result, scale + + inputs = [ + self.empty(5, 16), # input + self.empty(16), # weight + self.empty_f32(1, 1), # scale + ] + + pm.register_replacement( + pattern, + replacement, + inputs, + pm.fwd_only, + pm_pass, + ) + + +class AiterFusedAddRMSNormStaticQuantPattern(AiterRMSNormQuantPattern): + """This pattern fuses fused_add_rms_norm & aiter static per-tensor fp8 + quant into an aiter fused_add_rms_fp8_per_tensor_static_quant op. + """ + + FUSED_OP = rocm_aiter_ops.get_fused_add_rms_fp8_per_tensor_static_quant_op() + + def __init__( + self, + epsilon: float, + quant_dtype: torch.dtype, + symmetric: bool = True, + ) -> None: + key = FusedRMSQuantKey( + fused_add=True, + quant=QuantKey( + dtype=quant_dtype, scale=kStaticTensorScale, symmetric=symmetric + ), + ) + + super().__init__(epsilon, key) + + def register(self, pm_pass: PatternMatcherPass) -> None: + def pattern( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + result_rms, residual_out = torch.ops.vllm_ir.fused_add_rms_norm( + input, residual, weight, self.epsilon + ) + result, scale = self.quant_matcher(result_rms, scale) + + return result, residual_out, scale + + def replacement( + input: torch.Tensor, + weight: torch.Tensor, + residual: torch.Tensor, + scale: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + result = self.FUSED_OP( + x=input, + residual=residual, + weight=weight, + epsilon=self.epsilon, + scale=scale, + quant_dtype=self.quant_dtype, + ) + + # The static scale is an input, so it passes through unchanged. + return result[0], result[1], scale + + inputs = [ + self.empty(5, 16), # input + self.empty(16), # weight + self.empty(5, 16), # residual + self.empty_f32(1, 1), # scale + ] + + pm.register_replacement( + pattern, + replacement, + inputs, + pm.fwd_only, + pm_pass, + ) + + class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): """This pass fuses aiter rms_norm & vllm/aiter quant custom ops into a fused rms_norm_quant op. @@ -630,6 +765,14 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): epsilon, FP8_DTYPE, GroupShape(1, 128), match_aiter_quant_op ).register(self.patterns) + # Fuse aiter rms_norm + aiter static per-tensor fp8 quant + AiterRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(self.patterns) + + # Fuse aiter fused_add_rms_norm + aiter static per-tensor fp8 quant + AiterFusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register( + self.patterns + ) + # When quant_fp8 custom ops are disabled, both AITER and native # quant matchers trace through QuantFP8's native implementation. # Registering both variants would create duplicate Inductor @@ -690,6 +833,8 @@ class RocmAiterRMSNormQuantFusionPass(VllmPatternMatcherPass): AiterFusedAddRMSNormDynamicQuantPattern, AiterRMSFp8GroupQuantPattern, AiterFusedAddRMSFp8GroupQuantPattern, + AiterRMSNormStaticQuantPattern, + AiterFusedAddRMSNormStaticQuantPattern, DoubleAiterRMSFp8GroupQuantPattern, DoubleAiterRMSFp8GroupQuantViewPattern, AiterRMSNormGatedFp8GroupQuantPattern, From 68ebea789a4652a6b16665d50d79805866b933ec Mon Sep 17 00:00:00 2001 From: Tyler Michael Smith Date: Fri, 16 Oct 2026 16:20:52 -0400 Subject: [PATCH 20/49] [Perf][MoE] Skip padded work in SM100 FP4 DeepGEMM experts (#59044) Signed-off-by: Tyler Michael Smith Co-authored-by: Codex Co-authored-by: Robert Shaw <114415538+robertgshaw2-redhat@users.noreply.github.com> --- ..._mul_per_token_group_quant_fp8_colmajor.py | 8 ++++-- .../layers/fused_moe/experts/deep_gemm_moe.py | 27 ++++++++++++++----- 2 files changed, 27 insertions(+), 8 deletions(-) diff --git a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py index b7dcb956c8..771c1f2440 100644 --- a/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py +++ b/tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py @@ -116,7 +116,7 @@ def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int): @pytest.mark.parametrize("scale_format", ["packed", "float32", "ceil_ue8m0"]) @pytest.mark.parametrize("hidden_size", [256, 2048]) @pytest.mark.parametrize("clamp_limit", [None, 10.0]) -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA graph test") +@pytest.mark.skipif(not current_platform.is_cuda_alike(), reason="CUDA/ROCm graph test") def test_silu_quant_skips_expert_padding_on_graph_replay( counts: list[int], group_size: int, @@ -135,7 +135,10 @@ def test_silu_quant_skips_expert_padding_on_graph_replay( x *= 16 clean_x = x.clone() ends = torch.empty(len(counts), device=DEVICE, dtype=torch.int32) - output = torch.empty((capacity, hidden_size), device=DEVICE, dtype=FLOAT8_DTYPE) + output_dtype = ( + FLOAT8_DTYPE if scale_format == "packed" else current_platform.fp8_dtype() + ) + output = torch.empty((capacity, hidden_size), device=DEVICE, dtype=output_dtype) def set_counts(values): positions = [] @@ -187,6 +190,7 @@ def test_silu_quant_skips_expert_padding_on_graph_replay( output.view(torch.uint8).fill_(91) graph.replay() reference_q, reference_s = quantize(x) + assert actual_q.dtype == reference_q.dtype torch.testing.assert_close( actual_q.view(torch.uint8)[live], reference_q.view(torch.uint8)[live], diff --git a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py index 0eef7a8335..cfa380f9b4 100644 --- a/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/deep_gemm_moe.py @@ -545,7 +545,12 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): return (workspace1, workspace2, output) def _act_mul_quant( - self, input: torch.Tensor, output: torch.Tensor, activation: MoEActivation + self, + input: torch.Tensor, + output: torch.Tensor, + activation: MoEActivation, + expert_ends: torch.Tensor | None = None, + expert_alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: block_k = self._ACT_BLOCK_K scale_fmt = DeepGemmQuantScaleFMT.from_oracle() @@ -561,6 +566,8 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): output_q=output, group_size=block_k, clamp_limit=self.gemm1_clamp_limit, + expert_ends=expert_ends, + expert_alignment=expert_alignment, ) use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 return silu_mul_per_token_group_quant_fp8_colmajor( @@ -634,7 +641,8 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): a1q_perm = _resize_cache( workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K) ) - a1q, a1q_scale, expert_ids, inv_perm, align_used = deepgemm_moe_permute( + use_psum_layout = current_platform.is_device_capability_family(100) + a1q, a1q_scale, grouped_layout, inv_perm, align_used = deepgemm_moe_permute( aq=a1q, aq_scale=a1q_scale, topk_ids=topk_ids, @@ -642,6 +650,7 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): expert_map=expert_map, expert_tokens_meta=expert_tokens_meta, aq_out=a1q_perm, + use_psum_layout=use_psum_layout, ) assert a1q.size(0) == M_sum @@ -655,18 +664,23 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): (a1q, a1q_scale), (w1.view(torch.int8), self.w1_scale), mm1_out, - expert_ids, + grouped_layout, + use_psum_layout=use_psum_layout, recipe_a=(1, self._ACT_BLOCK_K), recipe_b=(1, self._WEIGHT_BLOCK_K), ) - # SwiGLU activation + FP8 requant + # SiLU+mul and FP8 requant use the same live expert ranges as GEMM. activation_out_dim = self.adjust_N_for_activation(N, activation) quant_out = _resize_cache( workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, activation_out_dim) ) a2q, a2q_scale = self._act_mul_quant( - input=mm1_out.view(-1, N), output=quant_out, activation=activation + input=mm1_out.view(-1, N), + output=quant_out, + activation=activation, + expert_ends=grouped_layout if use_psum_layout else None, + expert_alignment=align_used if use_psum_layout else 0, ) # FC2: FP8 activations x FP4 weights @@ -675,7 +689,8 @@ class DeepGemmFP4Experts(mk.FusedMoEExpertsModular): (a2q, a2q_scale), (w2.view(torch.int8), self.w2_scale), mm2_out, - expert_ids, + grouped_layout, + use_psum_layout=use_psum_layout, recipe_a=(1, self._ACT_BLOCK_K), recipe_b=(1, self._WEIGHT_BLOCK_K), ) From 9731f30a030c545e7f29a8a4ed96b8562be12c88 Mon Sep 17 00:00:00 2001 From: Dan Blanaru <48605845+DanBlanaru@users.noreply.github.com> Date: Fri, 16 Oct 2026 22:59:48 +0200 Subject: [PATCH 21/49] [Bugfix] Keep VLLM_RAY_BUNDLE_INDICES out of the compile cache key (#61211) Signed-off-by: DanBlanaru <48605845+DanBlanaru@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- tests/test_envs.py | 8 ++++++++ vllm/envs.py | 1 + 2 files changed, 9 insertions(+) diff --git a/tests/test_envs.py b/tests/test_envs.py index 7356539bcd..c27a57ba4d 100644 --- a/tests/test_envs.py +++ b/tests/test_envs.py @@ -48,6 +48,14 @@ def test_api_key_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch): assert "VLLM_API_KEY" not in envs.compile_factors() +def test_ray_bundle_indices_is_not_compile_factor(monkeypatch: pytest.MonkeyPatch): + # Engines sharing a placement group get different bundle indices but + # compile the same graphs, so they must share a cache key. + monkeypatch.setenv("VLLM_RAY_BUNDLE_INDICES", "4,5,6,7") + + assert "VLLM_RAY_BUNDLE_INDICES" not in envs.compile_factors() + + def test_p2p_side_channel_defaults_and_override(monkeypatch: pytest.MonkeyPatch): monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_HOST", raising=False) monkeypatch.delenv("VLLM_P2P_SIDE_CHANNEL_PORT", raising=False) diff --git a/vllm/envs.py b/vllm/envs.py index 19e3250e2f..ab04f405d7 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -2366,6 +2366,7 @@ def compile_factors() -> dict[str, object]: "VLLM_RANDOMIZE_DP_DUMMY_INPUTS", "VLLM_MODEL_REDIRECT_PATH", "VLLM_HOST_IP", + "VLLM_RAY_BUNDLE_INDICES", "VLLM_ELASTIC_EP_SCALE_UP_LAUNCH", "VLLM_FORCE_AOT_LOAD", "S3_ACCESS_KEY_ID", From 153a17166e5d07bc38184dedeea542d2e8885e55 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 16 Oct 2026 14:08:29 -0700 Subject: [PATCH 22/49] [Frontend] Add the Beam reasoning and tool parsers (#62052) Signed-off-by: Woosuk Kwon Co-authored-by: Bugen Zhao Co-authored-by: Yoni Gozlan <74535834+yonigozlan@users.noreply.github.com> Co-authored-by: neilkg <33635204+neilkg@users.noreply.github.com> Co-authored-by: Cursor Co-authored-by: Claude Opus 5.5 --- docs/features/tool_calling.md | 24 + rust/src/chat/src/lib.rs | 4 +- rust/src/chat/src/output/default/mod.rs | 41 +- rust/src/chat/src/parser/unified.rs | 8 +- rust/src/parser/src/output_grammar.rs | 5 +- .../parser/src/output_grammar/arguments.rs | 97 +- rust/src/parser/src/unified/beam.rs | 543 +++++++++ rust/src/parser/src/unified/beam/tests.rs | 897 +++++++++++++++ rust/src/parser/src/unified/hf/content.rs | 79 +- rust/src/parser/src/unified/hf/mod.rs | 2 +- rust/src/parser/src/unified/hf/parser.rs | 44 + rust/src/parser/src/unified/hf/template.rs | 25 + rust/src/parser/src/unified/mod.rs | 4 +- tests/parser/test_beam_parser.py | 747 ++++++++++++ .../parser/test_response_template_grammar.py | 492 ++++++++ vllm/entrypoints/openai/responses/protocol.py | 11 + vllm/entrypoints/openai/responses/serving.py | 3 + vllm/envs.py | 6 + vllm/parser/beam.py | 262 +++++ vllm/parser/parser_manager.py | 65 +- vllm/parser/response_template.py | 166 ++- vllm/parser/response_template_grammar.py | 1000 +++++++++++++++++ vllm/reasoning/__init__.py | 4 + vllm/reasoning/beam_reasoning_parser.py | 6 + vllm/renderers/online_renderer.py | 11 +- vllm/tool_parsers/__init__.py | 4 + vllm/tool_parsers/beam_tool_parser.py | 6 + 27 files changed, 4464 insertions(+), 92 deletions(-) create mode 100644 rust/src/parser/src/unified/beam.rs create mode 100644 rust/src/parser/src/unified/beam/tests.rs create mode 100644 tests/parser/test_beam_parser.py create mode 100644 tests/parser/test_response_template_grammar.py create mode 100644 vllm/parser/beam.py create mode 100644 vllm/parser/response_template_grammar.py create mode 100644 vllm/reasoning/beam_reasoning_parser.py create mode 100644 vllm/tool_parsers/beam_tool_parser.py diff --git a/docs/features/tool_calling.md b/docs/features/tool_calling.md index dd56db0d57..9980429324 100644 --- a/docs/features/tool_calling.md +++ b/docs/features/tool_calling.md @@ -201,6 +201,30 @@ checkpoint's output format. !!! note With `tool_choice="auto"`, structural-tag constraints require both `VLLM_ENFORCE_STRICT_TOOL_CALLING=true` (the default) and at least one tool with `strict: true`, or a server-side floor set via `--tool-strict-level`. When these conditions are met and the selected parser supports structural tags, vLLM constrains the tool-call envelope and pins the argument schema of each tool that sets `strict: true` (or of every tool under `--tool-strict-level parameter`). Otherwise, vLLM extracts tool calls from raw text, so arguments may occasionally be malformed or violate the function's parameter schema. +### Beam Models (`beam`) + +Beam checkpoints describe their output format with `response_template` metadata, so the `beam` parser reads it like the `hf` parser does: + +```bash +vllm serve \ + --enable-auto-tool-choice \ + --tool-call-parser beam \ + --reasoning-parser beam +``` + +Unlike `hf`, the `beam` parser also constrains its output. For strict tools, required or named tool choice, and structured outputs, it derives a +structural-tag grammar from the template. The grammar follows the template's tool call format, pins the arguments of strict tools when their +encoding can be derived, allows a single call with `parallel_tool_calls=false`, and frames a structured output as the answer. Raw string arguments +keep the `format` or `pattern` of their schema, as XGrammar keeps them in JSON strings, unless a match could end the value early. With +`tool_choice="none"`, the grammar allows only reasoning and an answer, so a model that still sees the tools in its prompt cannot call them. When no +grammar can be derived, required and named tool choice are rejected, and strict tools under `tool_choice="auto"` are left unconstrained with a +warning. + +It also constrains the output of requests without strict tools, forced tool choice, or structured outputs to the template's format: tool calls with +free arguments or an answer, or just the rest of an answer the prompt already opened. This always-on grammar needs the `xgrammar` structured +outputs backend, the default, and is disabled with `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT=0`. From the message start, every `beam` grammar opens with +reasoning that holds visible text, unless `enable_thinking` is false. + ### Hermes Models (`hermes`) All Nous Research Hermes-series models newer than Hermes 2 Pro should be supported. diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index 1ffb95fe37..5abdd36cc9 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -434,7 +434,7 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); + expect_test::expect!["tool parser `definitely_missing_tool_parser` is not registered (choose from: beam, deepseek_v3, deepseek_v31, deepseek_v32, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, granite4, hermes, hf, hy_v3, hy_v4, inkling, internlm, kimi_k2, kimi_k3, llama3_json, llama4_json, mimo, minimax_m2, minimax_m3, mistral, phi4_mini_json, qwen3_coder, qwen3_json, qwen3_xml, seed_oss)"].assert_eq(&error.to_report_string()); } #[test] @@ -445,6 +445,6 @@ mod tests { ) .unwrap_err(); - expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); + expect_test::expect!["reasoning parser `definitely_missing_reasoning_parser` is not registered (choose from: beam, cohere_cmd, deepseek_r1, deepseek_v3, deepseek_v4, deepseek_v41, gemma4, glm45, glm47, hf, hy_v3, hy_v4, inkling, kimi, kimi_k2, kimi_k3, mimo, minimax_m2, minimax_m3, nemotron_v3, qwen3, seed_oss, step3, step3p5)"].assert_eq(&error.to_report_string()); } } diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 67a5a11eae..dbbedfd836 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -23,7 +23,8 @@ use crate::output::{ChatOutputProcessor, DynChatEventStream, DynDecodedTextEvent use crate::parser::reasoning::{ReasoningParser, ReasoningParserFactory}; use crate::parser::tool::{ToolParser, ToolParserFactory}; use crate::parser::unified::{ - HfTemplateError, HfUnifiedParser, ResponseTemplate, UnifiedParserFactory, names, + BeamUnifiedParser, HfTemplateError, HfUnifiedParser, ResponseTemplate, UnifiedParserFactory, + names, }; use crate::parser::{ParserSelection, ToolStrictLevel}; use crate::request::{ChatRequest, ChatTool}; @@ -206,17 +207,19 @@ impl DefaultChatOutputProcessor { }); } - // `hf` is built from the model's template, never by its registry constructor. - let parser = if parser_name == names::HF { - let template = - response_template.as_ref().map_err(|error| Error::ParserInitialization { - kind: "unified", - name: parser_name.to_string(), - error: Box::new(error.clone()), - })?; - Box::new(HfUnifiedParser::new(template.clone(), tools, tokenizer)) - } else { - factory.create(parser_name, tools, tokenizer)? + // Parsers driven by the model's template are built from it, never by + // their registry constructors. + let template = || { + response_template.clone().map_err(|error| Error::ParserInitialization { + kind: "unified", + name: parser_name.to_string(), + error: Box::new(error), + }) + }; + let parser: Box = match parser_name { + names::HF => Box::new(HfUnifiedParser::new(template()?, tools, tokenizer)), + names::BEAM => Box::new(BeamUnifiedParser::new(template()?, tools, tokenizer)), + _ => factory.create(parser_name, tools, tokenizer)?, }; UNIFIED_PARSER_LOG_ONCE.call_once(|| info!(parser_name, "using unified parser")); @@ -589,6 +592,20 @@ mod tests { .assert_eq(&error(invalid_template())); } + #[test] + fn explicit_beam_uses_the_model_response_template() { + let beam = ParserSelection::Explicit("beam".to_string()); + crate::parser::validate_parser_overrides(&beam, &beam).unwrap(); + assert!(keeps_special_tokens("other-model", &loaded_template(), &beam, &beam).unwrap()); + + let error = + keeps_special_tokens("other-model", &Err(HfTemplateError::Missing), &beam, &beam) + .unwrap_err() + .to_report_string(); + expect_test::expect!["failed to initialize unified parser `beam`: the model's tokenizer_config.json provides no response_template"] + .assert_eq(&error); + } + #[test] fn auto_never_selects_hf() { let auto = ParserSelection::Auto; diff --git a/rust/src/chat/src/parser/unified.rs b/rust/src/chat/src/parser/unified.rs index 7db093c696..705133bb0a 100644 --- a/rust/src/chat/src/parser/unified.rs +++ b/rust/src/chat/src/parser/unified.rs @@ -6,8 +6,8 @@ use std::sync::LazyLock; pub use vllm_parser::unified::{ - Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser, HyV4UnifiedParser, - InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser, + BeamUnifiedParser, Gemma4UnifiedParser, HfTemplateError, HfUnifiedParser, HyV3UnifiedParser, + HyV4UnifiedParser, InklingUnifiedParser, KimiK3UnifiedParser, ResponseTemplate, UnifiedParser, }; use vllm_tokenizer::DynTokenizer; @@ -26,6 +26,9 @@ pub mod names { /// name resolution only: it is built from the model's template, never by its /// registry constructor. pub const HF: &str = "hf"; + /// Beam checkpoints: [`HF`] with Beam's tool-call grammar, built from the + /// model's template the same way. + pub const BEAM: &str = "beam"; } /// Constructor signature for one registered unified parser implementation. @@ -55,6 +58,7 @@ impl UnifiedParserFactory { factory.register_parser::(names::KIMI_K3); // Opt-in only: no model patterns map to it. factory.register_parser::(names::HF); + factory.register_parser::(names::BEAM); factory .register_pattern("gemma-4", names::GEMMA4) diff --git a/rust/src/parser/src/output_grammar.rs b/rust/src/parser/src/output_grammar.rs index ba9b4968d5..0d05bb4dde 100644 --- a/rust/src/parser/src/output_grammar.rs +++ b/rust/src/parser/src/output_grammar.rs @@ -207,7 +207,8 @@ pub(crate) fn full_format_from_builder_for_test( ) } -fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec { +/// Tools as builder inputs, strict when the tool or the server floor asks for it. +pub(crate) fn tool_params(tools: &[Tool], strict_level: ToolStrictLevel) -> Vec { tools .iter() .map(|tool| { @@ -240,7 +241,7 @@ pub(crate) fn normalize_tool_choice( } /// Whether the request asks for a tool grammar at all. -fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { +pub(crate) fn tool_grammar_applies(ctx: &OutputGrammarContext<'_>) -> bool { if ctx.tools.is_empty() { return false; } diff --git a/rust/src/parser/src/output_grammar/arguments.rs b/rust/src/parser/src/output_grammar/arguments.rs index 0f621789a3..6e58c6956a 100644 --- a/rust/src/parser/src/output_grammar/arguments.rs +++ b/rust/src/parser/src/output_grammar/arguments.rs @@ -26,7 +26,7 @@ //! - `allOf` with more than one schema accepts any value. use serde_json::{Value, json}; -use xgrammar_structural_tag::format::{Format, JsonSchemaFormat}; +use xgrammar_structural_tag::format::{EndBoundary, Format, JsonSchemaFormat, TokenFormat}; pub use crate::schema::JsonType; use crate::schema::{self, OptionSource, SchemaOption, SchemaRoot}; @@ -60,7 +60,13 @@ impl ParameterKey<'_> { /// `prefix KEY suffix content end`: a tag for a declared key, or a /// sequence matching [`Self::FREE_PATTERN`] for a free key. - pub fn tag(&self, prefix: &str, suffix: &str, content: Format, end: &str) -> Format { + pub fn tag( + &self, + prefix: &str, + suffix: &str, + content: Format, + end: impl Into, + ) -> Format { match self { Self::Declared(key) => Format::tag(format!("{prefix}{key}{suffix}"), content, end), Self::Free => Format::sequence(vec![ @@ -68,7 +74,13 @@ impl ParameterKey<'_> { Format::regex(Self::FREE_PATTERN), Format::const_string(suffix), content, - Format::const_string(end), + match end.into() { + EndBoundary::Text(end) => Format::const_string(end), + EndBoundary::Texts(ends) => { + one_of(ends.into_iter().map(Format::const_string).collect()) + } + EndBoundary::Token(end) => Format::Token(TokenFormat { token: end.token }), + }, ]), } } @@ -183,6 +195,36 @@ impl ArgumentContext<'_> { /// Build the grammar of one call's arguments under `schema`, the tool's /// parameters schema, rendering each parameter with `syntax`. pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &ArgumentOptions) -> Format { + parameter_list( + options, + syntax.separator(), + parameters(schema, syntax, options), + ) +} + +/// The parameters of one call, each rendered by an [`ArgumentSyntax`], for a +/// syntax that lists them itself rather than through [`arguments`]. +pub struct Parameters { + /// The declared parameters, in declared order. + pub declared: Vec, + /// Any one additional parameter, if the schema allows them. + pub additional: Option, +} + +/// One declared parameter. +pub struct Parameter { + pub key: String, + pub format: Format, + pub required: bool, +} + +/// The parameters of one call under `schema`, the tool's parameters schema, +/// rendering each parameter with `syntax`. +pub fn parameters( + schema: &Value, + syntax: &dyn ArgumentSyntax, + options: &ArgumentOptions, +) -> Parameters { let cx = ArgumentContext { options, root: SchemaRoot::new(schema), @@ -190,8 +232,7 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument let parameter = |key: ParameterKey<'_>, schema: &Value| syntax.parameter(key, &cx.value_options(schema)); - let schema = cx.root.resolved(); - let (properties, additional) = match schema { + match cx.root.resolved() { Value::Object(schema) => { let additional = match schema.get("additionalProperties") { Some(Value::Bool(false)) => None, @@ -206,22 +247,33 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument .and_then(Value::as_array) .map(|required| required.iter().filter_map(Value::as_str).collect::>()) .unwrap_or_default(); - let properties = schema + let declared = schema .get("properties") .and_then(Value::as_object) .into_iter() .flatten() .filter_map(|(key, schema)| { - let parameter = parameter(ParameterKey::Declared(key), schema)?; - Some((parameter, required.contains(&key.as_str()))) + Some(Parameter { + key: key.clone(), + format: parameter(ParameterKey::Declared(key), schema)?, + required: required.contains(&key.as_str()), + }) }) .collect(); - (properties, additional) + Parameters { + declared, + additional, + } } - Value::Bool(false) => (vec![], None), - _ => (vec![], parameter(ParameterKey::Free, &ANY_SCHEMA)), - }; - parameter_list(options, syntax.separator(), properties, additional) + Value::Bool(false) => Parameters { + declared: vec![], + additional: None, + }, + _ => Parameters { + declared: vec![], + additional: parameter(ParameterKey::Free, &ANY_SCHEMA), + }, + } } /// Parameters in declared order, required ones mandatory, followed by any @@ -230,8 +282,10 @@ pub fn arguments(schema: &Value, syntax: &dyn ArgumentSyntax, options: &Argument fn parameter_list( options: &ArgumentOptions, separator: Option, - properties: Vec<(Format, bool)>, - additional: Option, + Parameters { + declared, + additional, + }: Parameters, ) -> Format { let separated = |parameter: Format| match &separator { Some(separator) => Format::sequence(vec![parameter, separator.clone()]), @@ -239,10 +293,10 @@ fn parameter_list( }; let mut elements = Vec::from_iter(separator.clone()); if options.any_order { - let any_required = properties.iter().any(|(_, required)| *required); - let parameters = properties + let any_required = declared.iter().any(|parameter| parameter.required); + let parameters = declared .into_iter() - .map(|(parameter, _)| parameter) + .map(|parameter| parameter.format) .chain(additional) .collect::>(); if !parameters.is_empty() { @@ -254,8 +308,11 @@ fn parameter_list( }); } } else { - for (parameter, required) in properties { - let parameter = separated(parameter); + for Parameter { + format, required, .. + } in declared + { + let parameter = separated(format); elements.push(if required { parameter } else { diff --git a/rust/src/parser/src/unified/beam.rs b/rust/src/parser/src/unified/beam.rs new file mode 100644 index 0000000000..fb3b904e4e --- /dev/null +++ b/rust/src/parser/src/unified/beam.rs @@ -0,0 +1,543 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +//! Unified parser for Beam checkpoints. +//! +//! Beam output is parsed by [`HfUnifiedParser`] with the checkpoint's +//! `response_template`; this parser adds Beam's output grammar. +//! A Beam message after `<|start|>assistant` is: +//! +//! ```text +//! <|channel|>analysis<|message|>REASONING<|end|><|start|>assistant +//! ( <|channel|>final<|message|>ANSWER<|return|> +//! | CALL (<|end|><|start|>assistant CALL)* <|call|> ) +//! ``` +//! +//! where reasoning may end at any closer of the template's reasoning field, +//! and `CALL` is +//! ` to=NAME<|channel|>analysis <|constrain|>xml<|message|>` followed by +//! `\nVALUE\n` parameters separated +//! by newlines. Where the template looks ahead for the next parameter, a value +//! may also end on a newline right before it. +//! +//! The grammar enforces the schema of every declared tool's arguments, in +//! declared or alphabetical order, since Beam drifts from schemas it does not +//! enforce. When the model may answer instead, calls may also name any other +//! tool with any arguments, which the parser passes on so a harness can tell +//! the model the tool does not exist. `ANSWER` follows the request's structured +//! output, if any. +//! +//! The chat template rejects a history turn without reasoning, so a message the +//! model opens holds non-empty reasoning: an empty one would come back as null +//! reasoning and fail the next request. +//! +//! Unless `VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT` turns it off, the grammar +//! applies to every message, including one the request does not constrain. + +use std::collections::BTreeMap; +use std::sync::{Arc, LazyLock}; + +use serde_json::Value; +use vllm_tokenizer::{DecodedText, DynTokenizer}; +use xgrammar_structural_tag::ToolChoice; +use xgrammar_structural_tag::format::{EndBoundary, Format}; +use xgrammar_structural_tag::tool::{ + BuilderToolChoice, NormalizedToolChoice, ToolChoiceValue, normalize_tool_choice, +}; + +use super::{ + HfPosition, HfUnifiedParser, ResponseTemplate, Result, UnifiedParser, UnifiedParserError, + UnifiedParserOutput, +}; +use crate::output_grammar::arguments::{ + self, ArgumentOptions, ArgumentSyntax, JsonType, ParameterKey, Parameters, ValueOption, one_of, +}; +use crate::output_grammar::{ + self, BuiltOutputGrammar, OutputGrammarContext, tool_grammar_applies, tool_params, +}; +use crate::tool::Tool; + +const ANALYSIS_OPEN: &str = "<|channel|>analysis<|message|>"; +const FINAL_OPEN: &str = "<|channel|>final<|message|>"; +const END: &str = "<|end|>"; +/// Closes one message and opens the next assistant message. +const NEXT_MESSAGE: &str = "<|end|><|start|>assistant"; +const RETURN: &str = "<|return|>"; +const CALL: &str = "<|call|>"; +/// Follows ` to=NAME` in a call. +const CALL_HEADER_END: &str = "<|channel|>analysis <|constrain|>xml<|message|>"; +const PARAMETER_OPEN: &str = "\n"; +const PARAMETER_CLOSE: &str = "\n"; +/// Ends a value where the template looks ahead for the next parameter. +const NEXT_PARAMETER: &str = "\n\n = LazyLock::new(|| { + std::env::var("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT").map_or(true, |value| { + matches!(value.to_lowercase().as_str(), "true" | "1") + }) +}); + +/// Beam parser: [`HfUnifiedParser`] with Beam's output grammar. +pub struct BeamUnifiedParser { + inner: HfUnifiedParser, + /// Whether a value may end where the next parameter opens. + lenient: bool, + /// What ends the reasoning and opens the next message. + reasoning_ends: Vec, +} + +impl BeamUnifiedParser { + /// Create a parser for one request from the checkpoint's `response_template`. + pub fn new(template: Arc, tools: &[Tool], tokenizer: DynTokenizer) -> Self { + Self { + lenient: template.argument_lookahead() == Some(PARAMETER_OPEN), + reasoning_ends: match template.reasoning_closers() { + [] => vec![NEXT_MESSAGE.to_string()], + closers => closers.to_vec(), + }, + inner: HfUnifiedParser::new(template, tools, tokenizer), + } + } +} + +impl UnifiedParser for BeamUnifiedParser { + fn create(_tools: &[Tool], _tokenizer: DynTokenizer) -> Result> { + Err(UnifiedParserError::NoNamedConstructor { + parser: "beam", + built_from: "the checkpoint's response_template", + }) + } + + fn initialize(&mut self, prompt_token_ids: &[u32]) -> Result<()> { + self.inner.initialize(prompt_token_ids) + } + + fn preserve_special_tokens(&self) -> bool { + self.inner.preserve_special_tokens() + } + + fn build_output_grammar( + &self, + ctx: &OutputGrammarContext<'_>, + ) -> output_grammar::Result> { + build_output_grammar( + self.inner.position(), + ctx, + *ALWAYS_CONSTRAIN_OUTPUT, + self.lenient, + &self.reasoning_ends, + ) + } + + fn tool_call_id(&self, tool_index: usize) -> Option<&str> { + self.inner.tool_call_id(tool_index) + } + + fn parse_into(&mut self, delta: DecodedText, output: &mut UnifiedParserOutput) -> Result<()> { + self.inner.parse_into(delta, output) + } + + fn finish(&mut self) -> Result { + self.inner.finish() + } + + fn reset(&mut self) -> String { + self.inner.reset() + } +} + +/// The grammar of a Beam message from the first generated token, or `None` +/// when the parser is neither at the start of the message nor inside its +/// reasoning, or when the request constrains nothing and `always_constrain` is +/// off. Values may end where the next parameter opens if `lenient`, and +/// reasoning ends at any of `reasoning_ends`. +fn build_output_grammar( + position: HfPosition, + ctx: &OutputGrammarContext<'_>, + always_constrain: bool, + lenient: bool, + reasoning_ends: &[String], +) -> output_grammar::Result> { + let forced = !ctx.tools.is_empty() + && matches!( + ctx.tool_choice, + ToolChoice::Value(ToolChoiceValue::Required) | ToolChoice::NamedFunction(_) + ); + let answer_format = ctx.answer.filter(|_| !forced); + // The prompt may still list the tools, so only the grammar keeps the model + // from calling them. + let tools_forbidden = !ctx.tools.is_empty() + && matches!(ctx.tool_choice, ToolChoice::Value(ToolChoiceValue::None)); + if !(always_constrain + || tool_grammar_applies(ctx) + || tools_forbidden + || answer_format.is_some()) + { + return Ok(None); + } + let text = Format::any_text(); + let reasoning_end = EndBoundary::from(reasoning_ends.to_vec()); + let reasoning = match position { + HfPosition::ResponseStart => Format::tag( + ANALYSIS_OPEN, + // At least one character, which cannot start Beam's closers. + Format::sequence(vec![Format::regex("[^<]"), text]), + reasoning_end, + ), + HfPosition::Reasoning => Format::tag("", text, reasoning_end), + // Inside an answer the prompt opened, only the always-on grammar + // applies, finishing a free answer. A tool choice that forces a call + // cannot be met there, and the request's own constraint applies as is. + HfPosition::Content => { + return Ok( + (always_constrain && !forced && answer_format.is_none()).then(|| { + BuiltOutputGrammar::from_token_zero(Format::sequence(answer_rest(None))) + }), + ); + } + HfPosition::Other => return Ok(None), + }; + let NormalizedToolChoice { + function_tools, + choice, + .. + } = normalize_tool_choice( + &tool_params(ctx.tools, ctx.tool_strict_level), + ctx.tool_choice.clone(), + )?; + if function_tools.is_empty() { + return Ok(match choice { + // `none`, or no tools to call: the model can only answer. + BuilderToolChoice::Auto => { + Some(BuiltOutputGrammar::from_token_zero(Format::sequence(vec![ + reasoning, + answer(answer_format), + ]))) + } + BuilderToolChoice::Required | BuilderToolChoice::Forced => None, + }); + } + + let mut calls = function_tools + .iter() + .map(|tool| { + let header = + Format::const_string(format!(" to={}{CALL_HEADER_END}", tool.function.name)); + let schema = tool.function.parameters.clone().unwrap_or(Value::Bool(true)); + Format::sequence(vec![header, arguments(&schema, lenient)]) + }) + .collect::>(); + if matches!(choice, BuilderToolChoice::Auto) { + let names = ctx.tools.iter().map(|tool| tool.name.as_str()).collect::>(); + let header = Format::sequence(vec![ + Format::const_string(" to="), + Format::regex(other_names(&names)), + Format::const_string(CALL_HEADER_END), + ]); + calls.push(Format::sequence(vec![ + header, + arguments(&Value::Bool(true), lenient), + ])); + } + let call = one_of(calls); + let mut calls = vec![call.clone()]; + if ctx.parallel_tool_calls && !choice.is_forced() { + calls.push(Format::star(Format::sequence(vec![ + Format::const_string(NEXT_MESSAGE), + call, + ]))); + } + // The last call ends with the `<|call|>` stop token. + calls.push(Format::optional(Format::const_string(CALL))); + let calls = Format::sequence(calls); + let output = match choice { + BuilderToolChoice::Auto => Format::or(vec![calls, answer(answer_format)]), + BuilderToolChoice::Required | BuilderToolChoice::Forced => calls, + }; + Ok(Some(BuiltOutputGrammar::from_token_zero(Format::sequence( + vec![reasoning, output], + )))) +} + +/// A final answer, ending with the `<|return|>` stop token. +fn answer(format: Option<&Format>) -> Format { + Format::sequence([vec![Format::const_string(FINAL_OPEN)], answer_rest(format)].concat()) +} + +/// The rest of an open final answer, holding `format` or else free text. +fn answer_rest(format: Option<&Format>) -> Vec { + vec![ + // Without `<|end|>`, no unconstrained message can follow the answer. + format.map_or_else(|| Format::any_text_excluding(&[END, RETURN]), Format::clone), + Format::optional(Format::const_string(RETURN)), + ] +} + +/// The arguments under `schema`: its declared parameters in declared or +/// alphabetical order, required ones mandatory, then any additional ones. If +/// `lenient`, each value but the last may end by opening the next parameter. +fn arguments(schema: &Value, lenient: bool) -> Format { + let options = ArgumentOptions { + excludes: CALL_TERMINATORS.map(str::to_string).to_vec(), + ..ArgumentOptions::default() + }; + let render = |syntax: BeamArguments| arguments::parameters(schema, &syntax, &options); + if lenient { + lenient_arguments(render(BeamArguments::Joined), render(BeamArguments::Last)) + } else { + closed_arguments(render(BeamArguments::Closed)) + } +} + +/// Arguments whose values all end with a closer. +fn closed_arguments( + Parameters { + declared, + additional, + }: Parameters, +) -> Format { + let additional = additional.map(Format::star); + let lists = orders(&declared, |parameter| ¶meter.key) + .into_iter() + .map(|order| { + let listed = order + .into_iter() + .map(|parameter| listed(parameter.format.clone(), parameter.required)); + listed.chain(additional.clone()).collect::>() + }) + .filter(|list| !list.is_empty()) + .map(Format::sequence) + .collect::>(); + if lists.is_empty() { + Format::const_string("") + } else { + one_of(lists) + } +} + +/// A declared parameter as one another parameter follows, `joined`, and as the +/// last one. +struct LenientParameter { + key: String, + joined: Format, + last: Format, + required: bool, +} + +impl LenientParameter { + fn joined(&self) -> Format { + listed(self.joined.clone(), self.required) + } +} + +/// Arguments whose values end by opening the next parameter, but for the last +/// one, from the parameters as [`BeamArguments::Joined`] and as +/// [`BeamArguments::Last`]. +fn lenient_arguments(joined: Parameters, last: Parameters) -> Format { + // Both forms share their terminators, so they keep the same parameters. + let declared = joined + .declared + .into_iter() + .zip(last.declared) + .map(|(joined, last)| LenientParameter { + key: joined.key, + joined: joined.format, + last: last.format, + required: joined.required, + }) + .collect::>(); + let additional = joined.additional.zip(last.additional); + + let mut lists = vec![]; + for order in orders(&declared, |parameter| ¶meter.key) { + // The last parameter is a declared one that follows every required + // one, or an additional one. + if !order.is_empty() { + let first = order.iter().rposition(|parameter| parameter.required).unwrap_or(0); + let mut list = + order[..first].iter().map(|parameter| parameter.joined()).collect::>(); + list.push(last_present(&order[first..])); + lists.push(Format::sequence(list)); + } + if let Some((joined, last)) = &additional { + let mut list = order.iter().map(|parameter| parameter.joined()).collect::>(); + list.extend([Format::star(joined.clone()), last.clone()]); + lists.push(Format::sequence(list)); + } + } + if lists.is_empty() { + return Format::const_string(""); + } + + // Each parameter starts past the opener, as the value before writes it. + let arguments = Format::sequence(vec![ + separator(), + Format::const_string(PARAMETER_OPEN), + one_of(lists), + ]); + if declared.iter().any(|parameter| parameter.required) { + arguments + } else { + Format::optional(arguments) + } +} + +/// `parameters`, at least one, in order and the last present one in its last +/// form. Only the first may be required. +fn last_present(parameters: &[&LenientParameter]) -> Format { + if let [only] = parameters { + return only.last.clone(); + } + // The last one is in the first half, or after all of it. Halving keeps the + // grammar O(n log n) in size and O(log n) deep. + let (head, tail) = parameters.split_at(parameters.len() / 2); + let mut after_head = head.iter().map(|parameter| parameter.joined()).collect::>(); + after_head.push(last_present(tail)); + Format::or(vec![last_present(head), Format::sequence(after_head)]) +} + +/// `parameters` in declared order, and in alphabetical order by `key` if that +/// differs, since Beam often writes them sorted. +fn orders

(parameters: &[P], key: impl Fn(&P) -> &str) -> Vec> { + let declared = parameters.iter().collect::>(); + let mut sorted = declared.clone(); + sorted.sort_by_key(|parameter| key(parameter)); + if sorted + .iter() + .map(|parameter| key(parameter)) + .eq(declared.iter().map(|parameter| key(parameter))) + { + vec![declared] + } else { + vec![declared, sorted] + } +} + +/// `parameter` as listed: optional unless `required`. +fn listed(parameter: Format, required: bool) -> Format { + if required { + parameter + } else { + Format::optional(parameter) + } +} + +fn separator() -> Format { + Format::optional(Format::const_string("\n")) +} + +/// A regex for the tool names the call header can spell (no whitespace or +/// `<`) besides `names`. +fn other_names(names: &[&str]) -> String { + #[derive(Default)] + struct Node { + children: BTreeMap, + end: bool, + } + + fn escape(char: char, meta: &str) -> String { + if meta.contains(char) { + format!("\\{char}") + } else { + char.to_string() + } + } + + /// A name character other than `chars`. + fn other<'a>(chars: impl Iterator) -> String { + let chars = chars.map(|&char| escape(char, r"\]^-[")).collect::(); + format!(r"[^\s<{chars}]") + } + + /// The rest of a name past `node` that ends on none of `names`. + fn rest(node: &Node, first: bool) -> String { + let any = other([].iter()); + let mut branches = vec![format!("{}{any}*", other(node.children.keys()))]; + branches.extend(node.children.iter().map(|(&char, child)| { + format!("{}{}", escape(char, r".^$*+?{}[]\|()"), rest(child, false)) + })); + let optional = if first || node.end { "" } else { "?" }; + format!("(?:{}){optional}", branches.join("|")) + } + + let mut root = Node::default(); + for name in names { + let mut node = &mut root; + for char in name.chars() { + node = node.children.entry(char).or_default(); + } + node.end = true; + } + rest(&root, true) +} + +/// How one parameter is written: `\nVALUE` followed +/// by its end. +/// +/// The parser reads every value as text and converts it by the tool schema, so +/// string options are raw text and every other option is JSON. +enum BeamArguments { + /// The whole parameter, ending with `\n`, after an optional + /// newline. + Closed, + /// The parameter past its opener, where the template looks ahead for the + /// next one: it ends by opening the next parameter, also right after the + /// value's newline, so raw text holds no such opener. + Joined, + /// The last parameter past its opener under such a template, ending with + /// `\n`. + Last, +} + +impl ArgumentSyntax for BeamArguments { + fn separator(&self) -> Option { + None + } + + fn parameter(&self, key: ParameterKey<'_>, options: &[ValueOption<'_>]) -> Option { + let terminators: &[&str] = match self { + Self::Closed => &[""], + Self::Joined | Self::Last => &["", NEXT_PARAMETER], + }; + let values = options + .iter() + .filter_map(|option| match option.ty { + JsonType::String => option.raw_string(terminators), + _ => Some(option.json()), + }) + .collect::>(); + (!values.is_empty()).then(|| { + let value = one_of(values); + match self { + Self::Closed => Format::sequence(vec![ + separator(), + key.tag(PARAMETER_OPEN, PARAMETER_NAME_CLOSE, value, PARAMETER_CLOSE), + ]), + Self::Joined => key.tag( + "", + PARAMETER_NAME_CLOSE, + value, + NEXT_PARAMETER_ENDS.to_vec(), + ), + Self::Last => key.tag("", PARAMETER_NAME_CLOSE, value, PARAMETER_CLOSE), + } + }) + } +} + +#[cfg(test)] +mod tests; diff --git a/rust/src/parser/src/unified/beam/tests.rs b/rust/src/parser/src/unified/beam/tests.rs new file mode 100644 index 0000000000..906a01f334 --- /dev/null +++ b/rust/src/parser/src/unified/beam/tests.rs @@ -0,0 +1,897 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +use std::sync::Arc; + +use serde_json::{Value, json}; +use vllm_tokenizer::test_utils::TestTokenizer; +use vllm_tokenizer::{DecodedText, Tokenizer as _}; +use xgrammar_structural_tag::ToolChoice; +use xgrammar_structural_tag::format::Format; + +use super::{BeamUnifiedParser, HfPosition, build_output_grammar, other_names}; +use crate::output_grammar::test_utils::outline; +use crate::output_grammar::{GrammarCoverage, OutputGrammarContext, ToolStrictLevel}; +use crate::tool::Tool; +use crate::unified::{ResponseTemplate, UnifiedParser, UnifiedParserEvent, UnifiedParserOutput}; + +/// The `tag_pattern` of the checkpoint of [`template`]. +const TAG_PATTERN: &str = + r#"\n(?P.*?)\n"#; +/// [`TAG_PATTERN`], but a value may also end where the next parameter opens. +const LENIENT_TAG_PATTERN: &str = r#"\n(?P.*?)\n(?:|(?= Arc { + template_with(tag_pattern, json!("<|end|><|start|>assistant")) +} + +/// [`template`] with `thinking_close` as the reasoning closer. +fn template_with(tag_pattern: &str, thinking_close: Value) -> Arc { + let template = json!({ + "defaults": {"role": "assistant"}, + "start_anchor_pattern": r"<\|start\|>assistant", + "fields": { + "thinking": { + "open": "<|channel|>analysis<|message|>", + "close": thinking_close, + "content": "text" + }, + "content": { + "open": "<|channel|>final<|message|>", + "close": ["<|return|>", "<|end|>", "<|call|>"], + "content": "text" + }, + "tool_calls": { + "open_pattern": r" to=(?P[^\s<]+)<\|channel\|>analysis <\|constrain\|>xml<\|message\|>", + "close": ["<|end|><|start|>assistant", "<|call|>"], + "repeats": true, + "content": "xml-inline", + "content_args": { + "tag_pattern": tag_pattern, + "value_parser": {"name": "text", "args": {"strip": false}} + }, + "transform": { + "type": "function", + "function": {"name": "{name}", "arguments": "{content}"} + } + } + } + }); + Arc::new(ResponseTemplate::from_json(&template).unwrap()) +} + +fn tools() -> Vec { + vec![ + Tool { + name: "get_weather".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": { + "location": {"type": "string"}, + "days": {"type": "integer"}, + "unit": {"enum": ["c", "f"]} + }, + "required": ["location"], + "additionalProperties": false + }), + strict: Some(true), + defer_loading: None, + }, + Tool { + name: "get_time".to_string(), + description: None, + parameters: json!({ + "type": "object", + "properties": {"zone": {"type": "string"}} + }), + strict: None, + defer_loading: None, + }, + ] +} + +/// A parser initialized from `prompt`, the rendered conversation. +fn parser(prompt: &str) -> BeamUnifiedParser { + parser_with(TAG_PATTERN, prompt) +} + +fn parser_with(tag_pattern: &str, prompt: &str) -> BeamUnifiedParser { + let tokenizer = Arc::new(TestTokenizer::new()); + let prompt_ids = tokenizer.encode(prompt, false).unwrap(); + let mut parser = BeamUnifiedParser::new(template(tag_pattern), &tools(), tokenizer); + parser.initialize(&prompt_ids).unwrap(); + parser +} + +const PROMPT: &str = "<|start|>user<|message|>Weather in Paris?<|end|><|start|>assistant"; + +fn grammar(position: HfPosition, tool_choice: ToolChoice, parallel: bool) -> Option { + grammar_with(position, &tools(), tool_choice, parallel, false, false) +} + +fn grammar_with( + position: HfPosition, + tools: &[Tool], + tool_choice: ToolChoice, + parallel: bool, + always_constrain: bool, + lenient: bool, +) -> Option { + let ctx = OutputGrammarContext { + tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: parallel, + answer: None, + }; + grammar_in(position, &ctx, always_constrain, lenient) +} + +fn grammar_in( + position: HfPosition, + ctx: &OutputGrammarContext<'_>, + always_constrain: bool, + lenient: bool, +) -> Option { + let grammar = build_output_grammar( + position, + ctx, + always_constrain, + lenient, + template(TAG_PATTERN).reasoning_closers(), + ) + .unwrap()?; + assert_eq!(grammar.coverage, GrammarCoverage::FromTokenZero); + Some(outline(&grammar.format)) +} + +fn constrained_grammar( + position: HfPosition, + tool_choice: ToolChoice, + answer: Format, + always_constrain: bool, +) -> Option { + let tools = tools(); + let ctx = OutputGrammarContext { + tools: &tools, + tool_choice: &tool_choice, + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: true, + answer: Some(&answer), + }; + grammar_in(position, &ctx, always_constrain, false) +} + +#[test] +fn position_follows_the_prompt_and_then_the_output() { + let position = |suffix: &str| parser(&format!("{PROMPT}{suffix}")).inner.position(); + + assert_eq!(position(""), HfPosition::ResponseStart); + assert_eq!( + position("<|channel|>analysis<|message|>Let me"), + HfPosition::Reasoning + ); + // The reasoning closer ends with the anchor, so the next message starts. + assert_eq!( + position("<|channel|>analysis<|message|>Done.<|end|><|start|>assistant"), + HfPosition::ResponseStart + ); + // A partial closer stays buffered. + assert_eq!( + position("<|channel|>analysis<|message|>Done.<|end|>"), + HfPosition::Other + ); + assert_eq!( + position("<|channel|>final<|message|>It is"), + HfPosition::Content + ); + // A partial boundary stays buffered. + assert_eq!(position("<|channel|>ana"), HfPosition::Other); + // Without the anchor, the parser parses the output from the start. + assert_eq!( + parser("no anchor").inner.position(), + HfPosition::ResponseStart + ); + + let mut parser = parser(PROMPT); + parser + .parse_into( + DecodedText::unattributed("<|channel|>analysis<|message|>Hmm"), + &mut UnifiedParserOutput::default(), + ) + .unwrap(); + assert_eq!(parser.inner.position(), HfPosition::Reasoning); +} + +#[test] +fn required_calls_any_tool_after_reasoning() { + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + sequence + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + star + sequence + `<|end|><|start|>assistant` + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional `<|call|>` + "#]].assert_eq( + &grammar(HfPosition::ResponseStart, ToolChoice::required(), true).unwrap(), + ); +} + +#[test] +fn auto_with_a_strict_tool_also_allows_an_answer() { + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + or + sequence + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + sequence + sequence + ` to=` + /(?:[^\sanalysis <|constrain|>xml<|message|>` + sequence + star + sequence + optional `\n` + sequence + `\n` + or + text excluding [`<|end|>`, `<|call|>`, ``] + json(number) excluding [`<|end|>`, `<|call|>`] + json(boolean) excluding [`<|end|>`, `<|call|>`] + json(null) excluding [`<|end|>`, `<|call|>`] + json({ ... }) excluding [`<|end|>`, `<|call|>`] + json(any[]) excluding [`<|end|>`, `<|call|>`] + `\n` + star + sequence + `<|end|><|start|>assistant` + or + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + sequence + sequence + ` to=` + /(?:[^\sanalysis <|constrain|>xml<|message|>` + sequence + star + sequence + optional `\n` + sequence + `\n` + or + text excluding [`<|end|>`, `<|call|>`, ``] + json(number) excluding [`<|end|>`, `<|call|>`] + json(boolean) excluding [`<|end|>`, `<|call|>`] + json(null) excluding [`<|end|>`, `<|call|>`] + json({ ... }) excluding [`<|end|>`, `<|call|>`] + json(any[]) excluding [`<|end|>`, `<|call|>`] + `\n` + optional `<|call|>` + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq(&grammar(HfPosition::ResponseStart, ToolChoice::auto(), true).unwrap()); +} + +#[test] +fn named_choice_calls_that_tool_once_after_prompted_reasoning() { + expect_test::expect![[r#" + sequence + tag `` text `<|end|><|start|>assistant` + sequence + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + or + sequence + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + sequence + optional + sequence + optional `\n` + tag `\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + optional + sequence + optional `\n` + tag `\n` .. `\n` + or + `c` + `f` + optional `<|call|>` + "#]].assert_eq( + &grammar(HfPosition::Reasoning, ToolChoice::function("get_weather"), true).unwrap(), + ); +} + +#[test] +fn no_grammar_without_a_constraining_request_or_known_position() { + let non_strict = vec![tools().pop().unwrap()]; + for tools in [&[][..], &non_strict] { + assert!( + grammar_with( + HfPosition::ResponseStart, + tools, + ToolChoice::auto(), + true, + false, + false + ) + .is_none() + ); + } + assert!( + grammar_with( + HfPosition::Other, + &tools(), + ToolChoice::required(), + true, + true, + false + ) + .is_none() + ); +} + +#[test] +fn none_with_tools_only_answers() { + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq(&grammar(HfPosition::ResponseStart, ToolChoice::none(), true).unwrap()); +} + +#[test] +fn always_constrain_answers_without_tools() { + expect_test::expect![[r#" + sequence + tag `` text `<|end|><|start|>assistant` + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq( + &grammar_with( + HfPosition::Reasoning, + &[], + ToolChoice::auto(), + true, + true, + false, + ) + .unwrap(), + ); +} + +#[test] +fn always_constrain_allows_non_strict_calls_or_an_answer() { + let non_strict = vec![tools().pop().unwrap()]; + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + or + sequence + or + sequence + ` to=get_time<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional + sequence + optional `\n` + tag `\n` text excluding [`<|end|>`, `<|call|>`, ``] `\n` + sequence + sequence + ` to=` + /(?:[^\sanalysis <|constrain|>xml<|message|>` + sequence + star + sequence + optional `\n` + sequence + `\n` + or + text excluding [`<|end|>`, `<|call|>`, ``] + json(number) excluding [`<|end|>`, `<|call|>`] + json(boolean) excluding [`<|end|>`, `<|call|>`] + json(null) excluding [`<|end|>`, `<|call|>`] + json({ ... }) excluding [`<|end|>`, `<|call|>`] + json(any[]) excluding [`<|end|>`, `<|call|>`] + `\n` + optional `<|call|>` + sequence + `<|channel|>final<|message|>` + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq( + &grammar_with( + HfPosition::ResponseStart, + &non_strict, + ToolChoice::auto(), + false, + true, + false, + ) + .unwrap(), + ); +} + +#[test] +fn always_constrain_finishes_an_open_answer() { + expect_test::expect![[r#" + sequence + text excluding [`<|end|>`, `<|return|>`] + optional `<|return|>` + "#]] + .assert_eq( + &grammar_with( + HfPosition::Content, + &tools(), + ToolChoice::auto(), + true, + true, + false, + ) + .unwrap(), + ); + assert!( + grammar_with( + HfPosition::Content, + &tools(), + ToolChoice::auto(), + true, + false, + false + ) + .is_none() + ); + assert!( + grammar_with( + HfPosition::Content, + &tools(), + ToolChoice::required(), + true, + true, + false + ) + .is_none() + ); +} + +#[test] +fn answer_holds_the_requested_format() { + let json = || Format::json_schema(json!({"type": "object"})); + let answer = constrained_grammar(HfPosition::ResponseStart, ToolChoice::none(), json(), false); + expect_test::expect![[r#" + sequence + tag `<|channel|>analysis<|message|>` .. `<|end|><|start|>assistant` + sequence + /[^<]/ + text + sequence + `<|channel|>final<|message|>` + json(object) + optional `<|return|>` + "#]] + .assert_eq(&answer.unwrap()); + + // Under `auto`, the answer that may replace calls holds it too. + let auto = + constrained_grammar(HfPosition::ResponseStart, ToolChoice::auto(), json(), false).unwrap(); + assert!(auto.ends_with( + "`<|channel|>final<|message|>`\n json(object)\n optional `<|return|>`\n" + )); + + // A forced call leaves no answer to constrain. + assert_eq!( + constrained_grammar( + HfPosition::ResponseStart, + ToolChoice::required(), + json(), + false + ), + grammar(HfPosition::ResponseStart, ToolChoice::required(), true) + ); +} + +#[test] +fn request_constraint_applies_alone_in_an_open_answer() { + let json = Format::json_schema(json!({"type": "object"})); + assert!(constrained_grammar(HfPosition::Content, ToolChoice::auto(), json, true).is_none()); +} + +/// Parse `output` after the prompt under `tag_pattern`, returning the calls as +/// `(name, arguments)`. +fn calls(tag_pattern: &str, output: &str) -> Vec<(String, Value)> { + let mut parser = parser_with(tag_pattern, PROMPT); + let mut events = UnifiedParserOutput::default(); + parser.parse_into(DecodedText::unattributed(output), &mut events).unwrap(); + events.append(parser.finish().unwrap()); + let mut calls: Vec<(String, String)> = Vec::new(); + for event in events.events { + if let UnifiedParserEvent::ToolCall(delta) = event { + if let Some(name) = delta.name { + calls.push((name, String::new())); + } + calls.last_mut().unwrap().1.push_str(&delta.arguments); + } + } + calls + .into_iter() + .map(|(name, arguments)| (name, serde_json::from_str(&arguments).unwrap())) + .collect() +} + +/// Output in the grammar's canonical spelling parses back to the values the +/// grammar admitted. +#[test] +fn parser_reads_back_the_grammar_spelling() { + let parameter = |key: &str, value: &str| { + format!("\n{value}\n") + }; + let call = |name: &str, parameters: &[String]| { + format!( + " to={name}<|channel|>analysis <|constrain|>xml<|message|>{}", + parameters.join("\n") + ) + }; + let output = format!( + "<|channel|>analysis<|message|>Two tools.<|end|><|start|>assistant{}<|end|><|start|>assistant{}", + call( + "get_weather", + &[ + parameter("location", "Paris"), + parameter("days", "2"), + parameter("unit", "c") + ] + ), + call("get_time", &[parameter("zone", "multi\nline")]), + ); + + for tag_pattern in [TAG_PATTERN, LENIENT_TAG_PATTERN] { + assert_eq!( + calls(tag_pattern, &output), + vec![ + ( + "get_weather".to_string(), + json!({"location": "Paris", "days": 2, "unit": "c"}) + ), + ("get_time".to_string(), json!({"zone": "multi\nline"})), + ] + ); + } +} + +#[test] +fn values_end_before_the_next_parameter_where_the_template_looks_ahead() { + let output = " to=get_weather<|channel|>analysis <|constrain|>xml<|message|>\ + \n2\n\ + \nParis\n"; + + assert!(parser_with(LENIENT_TAG_PATTERN, PROMPT).lenient); + assert_eq!( + calls(LENIENT_TAG_PATTERN, output), + vec![( + "get_weather".to_string(), + json!({"days": 2, "location": "Paris"}) + )] + ); + assert!(!parser(PROMPT).lenient); +} + +#[test] +fn reasoning_ends_at_every_closer_of_the_template() { + let closers = [ + "<|end|><|start|>assistant", + "<|constrain|><|start|>assistant", + "<|start|>assistant", + ]; + let tools = tools(); + let tokenizer = Arc::new(TestTokenizer::new()); + let mut parser = BeamUnifiedParser::new( + template_with(TAG_PATTERN, json!(closers)), + &tools, + tokenizer.clone(), + ); + parser.initialize(&tokenizer.encode(PROMPT, false).unwrap()).unwrap(); + let grammar = parser + .build_output_grammar(&OutputGrammarContext { + tools: &tools, + tool_choice: &ToolChoice::required(), + tool_strict_level: ToolStrictLevel::Auto, + parallel_tool_calls: false, + answer: None, + }) + .unwrap() + .unwrap(); + + assert_eq!( + outline(&grammar.format).lines().nth(1).unwrap(), + " tag `<|channel|>analysis<|message|>` .. `<|constrain|><|start|>assistant` | \ + `<|end|><|start|>assistant` | `<|start|>assistant`" + ); +} + +#[test] +fn lenient_arguments_come_in_declared_or_sorted_order() { + expect_test::expect![[r#" + sequence + tag `` text `<|end|><|start|>assistant` + sequence + sequence + ` to=get_weather<|channel|>analysis <|constrain|>xml<|message|>` + sequence + optional `\n` + `\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n` + sequence + optional tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n\n` .. `\n` + or + `c` + `f` + sequence + optional tag `days">\n` json(integer) excluding [`<|end|>`, `<|call|>`] `\n\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n` text excluding [`<|end|>`, `<|call|>`, ``, `\n\n\n` .. `\n` + or + `c` + `f` + optional `<|call|>` + "#]].assert_eq( + &grammar_with( + HfPosition::Reasoning, + &tools(), + ToolChoice::function("get_weather"), + true, + false, + true, + ) + .unwrap(), + ); +} + +#[test] +fn other_names_match_every_name_but_the_given_ones() { + let regex = regex_automata::meta::Regex::new(&format!( + "^(?:{})$", + other_names(&["get", "get_time", "a.b"]) + )) + .unwrap(); + + for name in [ + "g", + "ge", + "gets", + "get_", + "get_t", + "get_times", + "x", + "a", + "a.", + "axb", + "a.bc", + ] { + assert!(regex.is_match(name), "{name}"); + } + for name in ["get", "get_time", "a.b", "", "get time", "get<"] { + assert!(!regex.is_match(name), "{name}"); + } +} diff --git a/rust/src/parser/src/unified/hf/content.rs b/rust/src/parser/src/unified/hf/content.rs index d62ba9305e..598c6597e1 100644 --- a/rust/src/parser/src/unified/hf/content.rs +++ b/rust/src/parser/src/unified/hf/content.rs @@ -6,8 +6,12 @@ //! Port of `content_parsers.py`. Each parser takes a chunk of captured text and //! parses it into a single key in the output message dictionary. +use std::borrow::Cow; + use regex_automata::meta::Regex; +use regex_automata::util::iter::Searcher; use regex_automata::util::syntax; +use regex_automata::{Input, Match}; use serde::Deserialize; use serde_json::{Map, Value}; use thiserror_ext::AsReport as _; @@ -22,6 +26,10 @@ use crate::unified::{self, parsing_failed}; const LAX_OPEN: char = '\u{1}'; const LAX_CLOSE: char = '\u{2}'; +/// Name of the group that stands for the lookahead ending an `xml-inline` tag +/// pattern. +const AHEAD_GROUP: &str = "lookahead"; + /// A configured content parser (`content` plus `content_args` of one field). #[derive(Debug, Clone)] pub(super) enum ContentParser { @@ -66,6 +74,9 @@ pub(super) struct XmlInlineArgs { key_group: usize, /// Index of the `value` group, if the pattern declares one. value_group: Option, + /// Index of the group standing for the lookahead that ends the pattern, + /// and the lookahead's own pattern, if any. + lookahead: Option<(usize, String)>, value_parser: Option>, /// Collect duplicate keys into a list. merge_duplicates: bool, @@ -143,7 +154,8 @@ impl ContentParser { } ContentKind::XmlInline => { let spec: XmlInlineSpec = deserialize(&args, &context)?; - let tag_pattern = compile_regex(scope, &spec.tag_pattern)?; + let (pattern, lookahead) = final_lookahead(&spec.tag_pattern); + let tag_pattern = compile_regex(scope, &pattern)?; let group_info = tag_pattern.group_info(); let Some(key_group) = group_info.to_index(Default::default(), "key") else { return Err(invalid!( @@ -152,10 +164,14 @@ impl ContentParser { )); }; let value_group = group_info.to_index(Default::default(), "value"); + let lookahead = group_info + .to_index(Default::default(), AHEAD_GROUP) + .zip(lookahead.map(str::to_string)); Self::XmlInline(XmlInlineArgs { tag_pattern, key_group, value_group, + lookahead, value_parser: ValueParserSpec::compile(scope, spec.value_parser)?, merge_duplicates: spec.merge_duplicates, }) @@ -193,6 +209,32 @@ impl ContentParser { Self::KvLines(args) => args.parse(text), } } + + /// The pattern of the lookahead that ends an `xml-inline` tag pattern, if + /// any. + pub fn lookahead(&self) -> Option<&str> { + match self { + Self::XmlInline(args) => args.lookahead.as_ref().map(|(_, ahead)| ahead.as_str()), + _ => None, + } + } +} + +/// `pattern` with a lookahead that ends it matched as the group +/// [`AHEAD_GROUP`] instead, as the Rust regex engine has no lookahead, and +/// the lookahead's own pattern. The engine rejects any other lookahead. +fn final_lookahead(pattern: &str) -> (Cow<'_, str>, Option<&str>) { + if let Some((head, rest)) = pattern.split_once("(?=") { + // Only closing parentheses may follow a lookahead that ends the pattern. + let ahead = rest.trim_end_matches(')'); + if ahead.len() < rest.len() && !ahead.contains(['(', ')']) { + return ( + format!("{head}(?P<{AHEAD_GROUP}>{rest}").into(), + Some(ahead), + ); + } + } + (pattern.into(), None) } impl ValueParserSpec { @@ -258,7 +300,21 @@ impl XmlInlineArgs { /// Parse shallow XML-ish tags into a dict. fn parse(&self, text: &str) -> unified::Result { let mut out = Map::new(); - for captures in self.tag_pattern.captures_iter(text) { + let mut captures = self.tag_pattern.create_captures(); + let mut tags = Searcher::new(Input::new(text)); + while tags + .advance(|input| { + self.tag_pattern.search_captures(input, &mut captures); + // The lookahead's text is not part of the tag. + let ahead = + self.lookahead.as_ref().and_then(|(group, _)| captures.get_group(*group)); + Ok(captures.get_match().map(|tag| match ahead { + Some(ahead) => Match::new(tag.pattern(), tag.start()..ahead.start), + None => tag, + })) + }) + .is_some() + { let key = captures.get_group(self.key_group).map(|span| &text[span.range()]).ok_or_else( || parsing_failed!("xml-inline: named group 'key' did not participate"), @@ -490,6 +546,25 @@ mod tests { ); } + #[test] + fn xml_inline_tags_may_end_where_the_pattern_looks_ahead() { + let xml = parser( + "xml-inline", + json!({ + "tag_pattern": r#"(?P.*?)\n(?:|(?=1\nx y\n") + .unwrap(); + assert_eq!(value, json!({"a": "1", "b": "x y"})); + assert_eq!(xml.lookahead(), Some("\w+)(?=:)(?P.*)"}); + assert!(ContentParser::new("test", kind("xml-inline"), args.as_object().unwrap()).is_err()); + } + #[test] fn kv_lines_split_and_strip() { let kv = parser("kv-lines", json!({"value_parser": {"name": "int"}})); diff --git a/rust/src/parser/src/unified/hf/mod.rs b/rust/src/parser/src/unified/hf/mod.rs index 568a8371d0..9d9bd3bc06 100644 --- a/rust/src/parser/src/unified/hf/mod.rs +++ b/rust/src/parser/src/unified/hf/mod.rs @@ -25,7 +25,7 @@ mod transform; #[cfg(test)] mod tests; -pub use parser::HfUnifiedParser; +pub use parser::{HfPosition, HfUnifiedParser}; pub use template::ResponseTemplate; use thiserror::Error; use thiserror_ext::Macro; diff --git a/rust/src/parser/src/unified/hf/parser.rs b/rust/src/parser/src/unified/hf/parser.rs index 5d5065f109..febdd47175 100644 --- a/rust/src/parser/src/unified/hf/parser.rs +++ b/rust/src/parser/src/unified/hf/parser.rs @@ -59,6 +59,20 @@ pub struct HfUnifiedParser { emitted_tool_count: usize, } +/// Where the parser is in the response; see [`HfUnifiedParser::position`]. +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum HfPosition { + /// Nothing parsed since the start anchor, or since the start of the stream + /// when there is none. + ResponseStart, + /// Inside an open reasoning field. + Reasoning, + /// Inside an open content field. + Content, + /// Anywhere else. + Other, +} + /// Parser mode: the region currently receiving text. enum HfMode { /// No implicit field is declared: text outside regions is discarded. @@ -156,6 +170,36 @@ impl HfUnifiedParser { parser } + /// Where the parser is in the response. Between + /// [`UnifiedParser::initialize`] and the first generated text, this is where + /// the prompt leaves the response. + /// + /// Text outside every region, which a template without an implicit field + /// discards, leaves no state behind and so does not move the parser from + /// [`HfPosition::ResponseStart`]. + pub fn position(&self) -> HfPosition { + if !self.buffer.text.is_empty() { + return HfPosition::Other; + } + match &self.mode { + HfMode::Implicit(occurrence) | HfMode::Explicit(occurrence) if occurrence.opened => { + match &self.template.regions[occurrence.region].kind { + RegionKind::Text(TextRegion { + role: TextRole::Reasoning, + .. + }) => HfPosition::Reasoning, + RegionKind::Text(TextRegion { + role: TextRole::Content, + .. + }) => HfPosition::Content, + _ => HfPosition::Other, + } + } + _ if self.closed.contains(&true) || self.emitted_tool_count > 0 => HfPosition::Other, + _ => HfPosition::ResponseStart, + } + } + /// Clear all per-stream state. fn reset_state(&mut self) { self.buffer.clear(); diff --git a/rust/src/parser/src/unified/hf/template.rs b/rust/src/parser/src/unified/hf/template.rs index 4889f6f9b5..29c096ea32 100644 --- a/rust/src/parser/src/unified/hf/template.rs +++ b/rust/src/parser/src/unified/hf/template.rs @@ -191,6 +191,31 @@ impl ResponseTemplate { }) } + /// The pattern of the lookahead that ends the `xml-inline` tag pattern of + /// tool-call arguments, such as for the opener of the next argument. + pub fn argument_lookahead(&self) -> Option<&str> { + self.regions.iter().find_map(|region| match ®ion.kind { + RegionKind::ToolCalls(calls) => calls.content.lookahead(), + _ => None, + }) + } + + /// The literal closers of the reasoning field, longest first, or none if it + /// has no literal close. + pub fn reasoning_closers(&self) -> &[String] { + self.regions + .iter() + .find_map(|region| match (®ion.kind, ®ion.close) { + (RegionKind::Text(text), Some(Boundary::Literals(closers))) + if text.role == TextRole::Reasoning => + { + Some(closers.as_slice()) + } + _ => None, + }) + .unwrap_or_default() + } + /// Return the start and end offsets of the last start-anchor match in /// `prompt`. pub(super) fn last_start_anchor(&self, prompt: &str) -> Option<(usize, usize)> { diff --git a/rust/src/parser/src/unified/mod.rs b/rust/src/parser/src/unified/mod.rs index 74bef8e1f4..585f28cfde 100644 --- a/rust/src/parser/src/unified/mod.rs +++ b/rust/src/parser/src/unified/mod.rs @@ -3,6 +3,7 @@ //! Unified parser interface for reasoning and tool-call deltas. +mod beam; mod combined; mod gemma4; mod hf; @@ -10,9 +11,10 @@ mod hy; mod inkling; mod kimi_k3; +pub use beam::BeamUnifiedParser; pub use combined::CombinedParser; pub use gemma4::Gemma4UnifiedParser; -pub use hf::{HfTemplateError, HfUnifiedParser, ResponseTemplate}; +pub use hf::{HfPosition, HfTemplateError, HfUnifiedParser, ResponseTemplate}; pub use hy::{HyV3UnifiedParser, HyV4UnifiedParser}; pub use inkling::InklingUnifiedParser; pub use kimi_k3::KimiK3UnifiedParser; diff --git a/tests/parser/test_beam_parser.py b/tests/parser/test_beam_parser.py new file mode 100644 index 0000000000..3673b36090 --- /dev/null +++ b/tests/parser/test_beam_parser.py @@ -0,0 +1,747 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from copy import deepcopy +from types import SimpleNamespace +from typing import Any + +import pytest +from openai.types.responses import FunctionTool, ToolChoiceFunction +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string + +from vllm.config.reasoning import ReasoningConfig +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionRequest, + ChatCompletionToolsParam, +) +from vllm.entrypoints.openai.responses.protocol import ( + ResponsesRequest, + ResponsesResponse, +) +from vllm.exceptions import VLLMValidationError +from vllm.parser.beam import BeamParser, BeamReasoningParser, BeamToolParser +from vllm.parser.parser_manager import ParserManager +from vllm.parser.response_template import ResponseTemplateParser +from vllm.reasoning import ReasoningParserManager +from vllm.tool_parsers import ToolParserManager + +# The `response_template` of the Beam checkpoints. +BEAM_RESPONSE_TEMPLATE: dict[str, Any] = { + "defaults": {"role": "assistant"}, + "start_anchor_pattern": r"<\|start\|>assistant", + "fields": { + "thinking": { + "open": "<|channel|>analysis<|message|>", + "close": "<|end|><|start|>assistant", + "content": "text", + }, + "content": { + "open": "<|channel|>final<|message|>", + "close": ["<|return|>", "<|end|>", "<|call|>"], + "content": "text", + }, + "tool_calls": { + "open_pattern": ( + r" to=(?P[^\s<]+)" + r"<\|channel\|>analysis <\|constrain\|>xml<\|message\|>" + ), + "close": ["<|end|><|start|>assistant", "<|call|>"], + "repeats": True, + "content": "xml-inline", + "content_args": { + "tag_pattern": ( + r'' + r'\n(?P.*?)\n(?:|(?=": 200, + "<|channel|>": 201, + "<|message|>": 202, + "<|end|>": 203, + "<|call|>": 204, + "<|return|>": 205, + "<|constrain|>": 206, +} +TEXT_TOKENS = {100: "assistant", 101: "analysis", 102: "final", 103: "plan"} +PROMPT_TOKEN_IDS = [200, 100] +THOUGHT = "<|channel|>analysis<|message|>plan<|end|><|start|>assistant" +FINAL = "<|channel|>final<|message|>" +NEXT_MESSAGE = "<|end|><|start|>assistant" + + +def _tool(name, properties, *, strict=None): + return ChatCompletionToolsParam( + type="function", + function={ + "name": name, + "parameters": {"type": "object", "properties": properties}, + "strict": strict, + }, + ) + + +TOOLS = [ + _tool("get_weather", {"city": {"type": "string"}}), + _tool("set_count", {"count": {"type": "integer"}}), +] +STRICT_TOOLS = [ + _tool("get_weather", {"city": {"type": "string"}}), + _tool("set_count", {"count": {"type": "integer"}}, strict=True), +] + + +class BeamTokenizer: + def __init__(self, response_template=BEAM_RESPONSE_TEMPLATE): + self.response_template = response_template + self.init_kwargs = {} + self.name_or_path = "" + self.all_special_tokens = list(VOCAB) + self.all_special_ids = list(VOCAB.values()) + + def decode(self, token_ids, skip_special_tokens=False, **_kwargs): + inverse_vocab = {token_id: token for token, token_id in VOCAB.items()} + return "".join( + TEXT_TOKENS.get(token_id, inverse_vocab.get(token_id, "")) + for token_id in token_ids + ) + + def get_vocab(self): + return VOCAB + + def encode(self, text, add_special_tokens=False): + del add_special_tokens + return list(text.encode()) + + +def request(**overrides): + values = dict( + tool_choice="auto", + include_reasoning=True, + tools=TOOLS, + parallel_tool_calls=True, + ) + values.update(overrides) + return SimpleNamespace(**values) + + +def parameter(name, value): + return f'\n{value}\n' + + +def call(name, body, *, terminator="<|call|>"): + return ( + f" to={name}<|channel|>analysis <|constrain|>xml<|message|>{body}{terminator}" + ) + + +def make_parser(prompt_token_ids=PROMPT_TOKEN_IDS, tools=TOOLS, **kwargs): + parser_cls = ParserManager.get_parser( + tool_parser_name="beam", + reasoning_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(), + ) + assert parser_cls is not None + parser = parser_cls(BeamTokenizer(), tools, **kwargs) + parser.set_prompt_token_ids(prompt_token_ids) + return parser + + +def parse(text, prompt_token_ids=PROMPT_TOKEN_IDS): + reasoning, content, calls = make_parser(prompt_token_ids).parse( + text, request(), enable_auto_tools=True + ) + return reasoning, content, normalize_calls(calls) + + +def normalize_calls(calls): + return [(call.name, json.loads(call.arguments)) for call in calls or []] + + +def test_beam_parsers_are_registered(): + assert ToolParserManager.get_tool_parser("beam") is BeamToolParser + assert ReasoningParserManager.get_reasoning_parser("beam") is BeamReasoningParser + + +def test_parser_manager_selects_beam_parser(): + parser_cls = ParserManager.get_parser( + tool_parser_name="beam", + reasoning_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(), + ) + + assert parser_cls is not None + assert issubclass(parser_cls, BeamParser) + assert parser_cls.reasoning_parser_cls is BeamReasoningParser + assert parser_cls.tool_parser_cls is BeamToolParser + assert parser_cls._enable_auto_tools + + +@pytest.mark.parametrize("other", ["hermes", "hf"]) +def test_beam_parser_rejects_other_parsers(other): + with pytest.raises(TypeError, match="cannot be combined"): + ParserManager.get_parser( + tool_parser_name=other, + reasoning_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(), + ) + + +def test_beam_parser_requires_checkpoint_metadata_at_startup(): + with pytest.raises( + TypeError, match="--reasoning-parser beam requires `response_template`" + ): + ParserManager.get_parser( + reasoning_parser_name="beam", + tokenizer=BeamTokenizer(response_template=None), + ) + + template = deepcopy(BEAM_RESPONSE_TEMPLATE) + del template["fields"]["tool_calls"] + with pytest.raises(TypeError, match="--tool-call-parser beam requires"): + ParserManager.get_parser( + tool_parser_name="beam", + enable_auto_tools=True, + tokenizer=BeamTokenizer(template), + ) + + +def test_reasoning_and_final_answer(): + assert parse(THOUGHT + FINAL + "Sunny.<|return|>") == ("plan", "Sunny.", []) + + +def test_prefilled_reasoning_continues_from_prompt(): + prompt = PROMPT_TOKEN_IDS + [201, 101, 202, 103] + text = " more.<|end|><|start|>assistant" + FINAL + "Done.<|return|>" + + assert parse(text, prompt) == (" more.", "Done.", []) + + +def test_parallel_calls_are_coerced_to_their_schemas(): + text = ( + THOUGHT + + call( + "get_weather", + parameter("city", "Paris"), + terminator="<|end|><|start|>assistant", + ) + + call("set_count", parameter("count", "7")) + ) + + assert parse(text) == ( + "plan", + None, + [("get_weather", {"city": "Paris"}), ("set_count", {"count": 7})], + ) + + +def test_call_stop_token_removed_from_the_text_still_ends_the_call(): + text = THOUGHT + call("get_weather", parameter("city", "Paris"), terminator="") + + assert parse(text) == ("plan", None, [("get_weather", {"city": "Paris"})]) + + +@pytest.mark.parametrize("end", ["<|end|>", "<|end|><|start|>"]) +def test_output_cut_before_the_next_message_start_ends_its_region(end): + """Output that stops after `<|end|>`, such as at `max_tokens`, has ended + its reasoning or call and lacks only the next message's start.""" + weather = call("get_weather", parameter("city", "Paris"), terminator=end) + + assert parse("<|channel|>analysis<|message|>plan" + end) == ("plan", None, []) + assert parse(THOUGHT + weather) == ( + "plan", + None, + [("get_weather", {"city": "Paris"})], + ) + + +THOUGHT_IDS = [201, 101, 202, 103, 203, 200, 100] +FINAL_IDS = [201, 102, 202] + + +@pytest.mark.parametrize( + ("prompt_token_ids", "token_ids", "count"), + [ + (PROMPT_TOKEN_IDS, [*THOUGHT_IDS, *FINAL_IDS, 205], 1), + (PROMPT_TOKEN_IDS + [201, 101, 202], [103, 103, 203, 200, 100, *FINAL_IDS], 2), + (PROMPT_TOKEN_IDS, [201, 101, 202, 103, 103, 103], 3), + ], + ids=["reasoning_then_answer", "prefilled_reasoning", "cut_off"], +) +def test_reasoning_tokens_are_counted(prompt_token_ids, token_ids, count): + parser = make_parser(prompt_token_ids) + + assert parser.count_reasoning_tokens(token_ids) == count + + +def test_reasoning_tokens_are_counted_as_the_generation_grows(): + token_ids = [*THOUGHT_IDS, *FINAL_IDS, 205] + parser = make_parser() + + assert [ + parser.count_reasoning_tokens(token_ids[:end]) + for end in range(1, len(token_ids) + 1) + ] == [ + make_parser().count_reasoning_tokens(token_ids[:end]) + for end in range(1, len(token_ids) + 1) + ] + assert parser.count_reasoning_tokens(FINAL_IDS) == 0 + + +def test_call_stop_token_ends_the_answer(): + assert parse(FINAL + "done<|call|>") == (None, "done", []) + + +@pytest.mark.parametrize("chunk_size", [1, 7, 10_000]) +def test_streaming_matches_non_streaming(chunk_size): + text = ( + THOUGHT + + call( + "get_weather", + parameter("city", "Paris"), + terminator="<|end|><|start|>assistant", + ) + + call("set_count", parameter("count", "7")) + ) + parser = make_parser() + reasoning, calls = [], [] + for start in range(0, len(text), chunk_size): + end = min(start + chunk_size, len(text)) + delta = parser.parse_delta( + text[start:end], [], request(), finished=end == len(text) + ) + if delta is not None: + reasoning.append(delta.reasoning or "") + calls.extend(delta.tool_calls) + + assert "".join(reasoning) == "plan" + assert [ + (call.function.name, json.loads(call.function.arguments)) for call in calls + ] == parse(text)[2] + + +WEATHER_CALL = call("get_weather", parameter("city", "Paris")) +WEATHER = [("get_weather", {"city": "Paris"})] + + +def chat_request(tools=STRICT_TOOLS, **body): + return ChatCompletionRequest( + model="beam", + messages=[{"role": "user", "content": "Weather?"}], + tools=tools and [tool.model_dump() for tool in tools], + **body, + ) + + +def named(name): + return {"type": "function", "function": {"name": name}} + + +def adjust(req, prompt_token_ids=PROMPT_TOKEN_IDS, **kwargs): + make_parser(prompt_token_ids, req.tools, **kwargs).adjust_request(req) + return req + + +def accepts(req, text): + grammar = Grammar.from_structural_tag( + req.extract_structured_outputs().structural_tag + ) + return _is_grammar_accept_string(grammar, text) + + +@pytest.mark.parametrize( + ("tool_choice", "text", "calls"), + [ + ("required", THOUGHT + WEATHER_CALL, WEATHER), + ("required", WEATHER_CALL, None), + ( + "required", + THOUGHT + + call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE) + + call("set_count", parameter("count", "7")), + [*WEATHER, ("set_count", {"count": 7})], + ), + ("required", THOUGHT + call("set_count", parameter("count", "x")), None), + ( + "required", + THOUGHT + call("get_weather", parameter("city", "a<|call|>b")), + None, + ), + ("required", THOUGHT + FINAL + "Paris<|return|>", None), + (named("set_count"), THOUGHT + WEATHER_CALL, None), + (named("get_weather"), THOUGHT + WEATHER_CALL, WEATHER), + ( + named("get_weather"), + THOUGHT + + call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE) + + WEATHER_CALL, + None, + ), + ], +) +def test_forced_tool_calls_follow_the_template_grammar(tool_choice, text, calls): + """From the message start, the grammar opens with the reasoning, and + accepts only calls the parser reads back.""" + req = adjust(chat_request(tool_choice=tool_choice)) + + assert req._grammar_from_parser + assert accepts(req, text) is (calls is not None) + if calls is not None: + assert parse(text)[2] == calls + + +def test_reasoning_may_be_skipped_when_thinking_is_disabled(): + req = adjust( + chat_request(tool_choice="required"), + chat_template_kwargs={"enable_thinking": False}, + ) + + assert accepts(req, WEATHER_CALL) + assert accepts(req, THOUGHT + WEATHER_CALL) + + +def test_parallel_tool_calls_false_allows_a_single_call(): + req = adjust(chat_request(tool_choice="required", parallel_tool_calls=False)) + calls = call("get_weather", parameter("city", "Paris"), terminator=NEXT_MESSAGE) + + assert accepts(req, THOUGHT + WEATHER_CALL) + assert not accepts(req, THOUGHT + calls + WEATHER_CALL) + + +@pytest.mark.parametrize("strict", [False, True]) +@pytest.mark.parametrize( + ("separator", "accepted"), [("", True), ("\n", True), ("\n\n", False), (" ", False)] +) +def test_parameters_are_separated_by_at_most_a_newline(strict, separator, accepted): + """A model that skips a parameter's opener must not write whitespace until + `max_tokens`.""" + write = _tool( + "write", + {"path": {"type": "string"}, "content": {"type": "string"}}, + strict=strict, + ) + req = adjust(chat_request([write], tool_choice="required")) + body = parameter("path", "a.py") + separator + parameter("content", "x") + + assert accepts(req, THOUGHT + call("write", body)) is accepted + assert accepts(req, THOUGHT + call("write", separator + body)) is accepted + + +def test_forced_calls_in_prefilled_reasoning_start_after_it(): + req = adjust( + chat_request(tool_choice="required"), PROMPT_TOKEN_IDS + [201, 101, 202, 103] + ) + + assert not req._grammar_from_parser + assert accepts(req, WEATHER_CALL) + + +def test_forced_calls_cannot_follow_a_prefilled_answer(): + with pytest.raises(VLLMValidationError, match="required tool choice"): + adjust(chat_request(tool_choice="required"), PROMPT_TOKEN_IDS + [201, 102, 202]) + + +def test_prefilled_answer_can_only_be_finished(): + req = adjust(chat_request(TOOLS), PROMPT_TOKEN_IDS + [201, 102, 202]) + + assert accepts(req, "Sunny.<|return|>") + assert not accepts(req, "Sunny." + NEXT_MESSAGE + WEATHER_CALL) + + +def test_forced_call_with_unencodable_tool_name_is_rejected(): + tools = [_tool("get weather", {"city": {"type": "string"}})] + + with pytest.raises(VLLMValidationError, match="tool name"): + adjust(chat_request(tools, tool_choice="required")) + + +def test_string_formats_and_patterns_are_held_in_raw_text(): + """Like XGrammar, a strict string keeps its format or pattern, unless a + match could end the value early.""" + contact = _tool( + "set_contact", + { + "email": {"type": "string", "format": "email"}, + "code": {"type": "string", "pattern": "^[A-Z]{3}$"}, + "n": {"type": "integer"}, + }, + strict=True, + ) + contact.function.parameters["required"] = ["email", "code", "n"] + req = adjust(chat_request([contact], tool_choice="required")) + + def contact_call(email, code, *rest): + body = parameter("email", email) + parameter("code", code) + "".join(rest) + return THOUGHT + call("set_contact", body) + + text = contact_call("a.b@x.co", "ABC", parameter("n", "1")) + calls = make_parser(tools=[contact]).parse(text, req, enable_auto_tools=True)[2] + assert accepts(req, text) + assert normalize_calls(calls) == [ + ("set_contact", {"email": "a.b@x.co", "code": "ABC", "n": 1}) + ] + assert not accepts(req, contact_call("ab", "ABC", parameter("n", "1"))) + assert not accepts(req, contact_call("a.b@x.co", "ABCD", parameter("n", "1"))) + assert not accepts(req, contact_call("a.b@x.co", "ABC")) + # Line breaks could end the value. + dotted = _tool( + "set_code", {"code": {"type": "string", "pattern": "^a.b$"}}, strict=True + ) + with pytest.raises(VLLMValidationError, match="required tool choice"): + adjust(chat_request([dotted], tool_choice="required")) + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + (THOUGHT + WEATHER_CALL, True), + (THOUGHT + FINAL + "Sunny.<|return|>", True), + (THOUGHT + FINAL + "Sunny." + NEXT_MESSAGE + WEATHER_CALL, False), + (THOUGHT + call("set_count", parameter("count", "x")), False), + ], +) +def test_strict_tools_constrain_automatic_tool_choice(text, accepted): + assert accepts(adjust(chat_request()), text) is accepted + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + (THOUGHT + FINAL + "Sunny.<|return|>", True), + (FINAL + "Sunny.<|return|>", False), + (WEATHER_CALL, False), + (THOUGHT + WEATHER_CALL, False), + (FINAL + "Sunny." + NEXT_MESSAGE + WEATHER_CALL, False), + ], +) +def test_tool_choice_none_allows_only_an_answer(text, accepted): + req = adjust(chat_request(TOOLS, tool_choice="none")) + + assert req._grammar_from_parser + assert accepts(req, text) is accepted + + +@pytest.mark.parametrize(("tools", "tool_choice"), [(TOOLS, "auto"), (None, None)]) +def test_always_constrained_output_follows_the_template_framing(tools, tool_choice): + req = adjust(chat_request(tools, tool_choice=tool_choice)) + count_call = THOUGHT + call("set_count", parameter("count", "3")) + + assert accepts(req, THOUGHT + FINAL + "Sunny.<|return|>") + assert not accepts(req, FINAL + "Sunny.<|return|>") + assert not accepts(req, THOUGHT.replace("plan", "") + FINAL + "Sunny.<|return|>") + assert not accepts(req, THOUGHT + "Sunny.") + assert accepts(req, count_call) is (tool_choice == "auto") + + +@pytest.mark.parametrize("response_format", [None, {"type": "json_object"}]) +@pytest.mark.parametrize("strict", [None, True]) +@pytest.mark.parametrize("tool_choice", ["auto", "required"]) +def test_automatic_choice_lets_calls_name_undeclared_tools( + strict, tool_choice, response_format +): + """Unless a forced choice limits calls to the declared tools. The parser + passes such calls on, so the harness can tell the model the tool does not + exist.""" + bash = _tool("bash", {"command": {"type": "string"}}, strict=strict) + req = adjust( + chat_request([bash], tool_choice=tool_choice, response_format=response_format) + ) + undeclared = THOUGHT + call("read_file", parameter("path", "a.py")) + parser = make_parser(tools=[bash]) + + assert accepts(req, THOUGHT + call("bash", parameter("command", "ls"))) + assert accepts(req, undeclared) is (tool_choice == "auto") + calls = parser.parse(undeclared, req, enable_auto_tools=True)[2] + assert normalize_calls(calls) == [("read_file", {"path": "a.py"})] + + +@pytest.mark.parametrize( + ("name", "body", "accepted"), + [ + ("set_count", parameter("count", "1"), True), + ("set_count", parameter("count", "x"), False), + ("get_weather", parameter("town", "Paris"), False), + # The grammar cannot enforce `minProperties`. + ("find", parameter("x", "1") + parameter("x", "y"), True), + ("search_files", parameter("x", "1") + parameter("x", "y"), True), + ("set_coun", parameter("x", "1") + parameter("x", "y"), True), + ("set_counts", parameter("x", "1") + parameter("x", "y"), True), + ], +) +def test_automatic_calls_enforce_declared_arguments_beside_any_other_name( + name, body, accepted +): + find = _tool("find", {"q": {"type": "string"}}) + find.function.parameters["minProperties"] = 1 + req = adjust(chat_request([*TOOLS, find])) + + assert accepts(req, THOUGHT + call(name, body)) is accepted + + +def test_value_ends_before_an_opener_on_the_next_line(): + write = _tool("write", {"path": {"type": "string"}, "content": {"type": "string"}}) + req = adjust(chat_request([write], tool_choice="required")) + body = '\nprint(1)\n\n' + parameter("path", "a.py") + text = THOUGHT + call("write", body) + + assert accepts(req, text) + calls = make_parser(tools=[write]).parse(text, req, enable_auto_tools=True)[2] + assert normalize_calls(calls) == [ + ("write", {"content": "print(1)\n", "path": "a.py"}) + ] + + +def test_always_constrained_output_can_be_disabled(monkeypatch): + monkeypatch.setenv("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT", "0") + + assert adjust(chat_request(TOOLS)).structured_outputs is None + + +def test_strict_tools_the_template_cannot_hold_leave_automatic_choice_unconstrained(): + dotted = _tool( + "set_code", {"code": {"type": "string", "pattern": "^a.b$"}}, strict=True + ) + + assert adjust(chat_request([dotted])).structured_outputs is None + + +def test_generic_parser_does_not_derive_the_template_grammar(): + req = chat_request(TOOLS, tool_choice="required") + parser = ResponseTemplateParser(BeamTokenizer(), req.tools, enable_auto_tools=True) + + with pytest.raises(VLLMValidationError, match="required tool choice"): + parser.adjust_request(req) + + +@pytest.mark.parametrize("tools", [None, TOOLS]) +def test_structured_output_is_framed_as_the_answer(tools): + schema = {"type": "object", "properties": {"a": {"type": "integer"}}} + req = adjust( + chat_request( + tools, + response_format={ + "type": "json_schema", + "json_schema": {"name": "answer", "schema": schema}, + }, + ) + ) + + assert req.response_format is None + assert accepts(req, THOUGHT + FINAL + '{"a": 1}<|return|>') + assert not accepts(req, THOUGHT + FINAL + "Sunny.<|return|>") + assert accepts(req, THOUGHT + WEATHER_CALL) is (tools is not None) + + +def test_responses_named_tool_choice_follows_the_template_grammar(): + req = adjust( + ResponsesRequest( + input="Weather?", + tools=[ + FunctionTool( + type="function", + name="get_weather", + parameters={ + "type": "object", + "properties": {"city": {"type": "string"}}, + }, + strict=True, + ) + ], + tool_choice=ToolChoiceFunction(type="function", name="get_weather"), + text={"format": {"type": "json_object"}}, + ) + ) + + assert req._grammar_from_parser + assert req.text.format.type == "json_object" + assert accepts(req, THOUGHT + WEATHER_CALL) + assert not accepts(req, FINAL + "{}<|return|>") + + +@pytest.mark.parametrize("status", ["in_progress", "completed"]) +@pytest.mark.parametrize("inside_reasoning", [False, True]) +def test_responses_preserves_text_config_with_template_grammar( + status, inside_reasoning +): + text = { + "format": { + "type": "json_schema", + "name": "answer", + "strict": True, + "schema": { + "type": "object", + "properties": {"a": {"type": "integer"}}, + "required": ["a"], + "additionalProperties": False, + }, + }, + "verbosity": "low", + } + req = ResponsesRequest(input="Answer", text=text) + expected_text = req.text.model_dump(by_alias=True) + prompt = PROMPT_TOKEN_IDS + ([201, 101, 202] if inside_reasoning else []) + adjust(req, prompt_token_ids=prompt) + params = req.to_sampling_params(default_max_tokens=128) + grammar = Grammar.from_structural_tag(params.structured_outputs.structural_tag) + prefix = "" if inside_reasoning else THOUGHT + assert _is_grammar_accept_string(grammar, prefix + FINAL + '{"a": 1}<|return|>') + assert not _is_grammar_accept_string(grammar, prefix + FINAL + "Sunny.<|return|>") + response = ResponsesResponse.from_request(req, params, "beam", 0, [], status) + assert response.model_dump(by_alias=True)["text"] == expected_text + + +@pytest.mark.parametrize("reverse_alternatives", [False, True]) +def test_reasoning_budget_initializes_from_checkpoint_pattern( + monkeypatch, reverse_alternatives +): + template = deepcopy(BEAM_RESPONSE_TEMPLATE) + alternatives = [ + r"^<\|channel\|>analysis<\|message\|>", + r"<\|start\|>assistant<\|channel\|>analysis<\|message\|>", + ] + if reverse_alternatives: + alternatives.reverse() + template["fields"]["thinking"] = { + "open_pattern": "(?:" + "|".join(alternatives) + ")", + "close": "<|end|>", + "content": "text", + } + tokenizer = BeamTokenizer(template) + monkeypatch.setattr( + "vllm.config.reasoning.cached_tokenizer_from_config", lambda **_: tokenizer + ) + config = ReasoningConfig(reasoning_parser="beam") + config.initialize_token_ids(SimpleNamespace()) + + assert config.enabled + # The marker must work both after the prompt's assistant header and in + # subsequent complete assistant messages, regardless of regex branch order. + assert config.reasoning_start_token_ids == tokenizer.encode( + "<|channel|>analysis<|message|>" + ) + assert config.reasoning_end_token_ids == tokenizer.encode("<|end|>") + assert config.natural_reasoning_end_token_ids == tokenizer.encode("<|end|>") + + +@pytest.mark.parametrize("pattern", [r"analysis.*", r"(?:analysis|thinking)"]) +def test_reasoning_budget_does_not_guess_a_marker_for_other_patterns(pattern): + template = deepcopy(BEAM_RESPONSE_TEMPLATE) + thinking = template["fields"]["thinking"] + del thinking["open"] + thinking["open_pattern"] = pattern + + assert BeamReasoningParser(BeamTokenizer(template)).reasoning_start_str is None diff --git a/tests/parser/test_response_template_grammar.py b/tests/parser/test_response_template_grammar.py new file mode 100644 index 0000000000..29ae227262 --- /dev/null +++ b/tests/parser/test_response_template_grammar.py @@ -0,0 +1,492 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json +from copy import deepcopy +from typing import Any + +import pytest +from xgrammar import Grammar +from xgrammar.testing import _is_grammar_accept_string + +from vllm.parser.chat_parsing.response_templates import load_response_template +from vllm.parser.response_template_grammar import ( + answer_rest_structural_tag, + tool_call_structural_tag, +) + +TEMPLATE: dict[str, Any] = { + "start_anchor": "", + "fields": { + "thinking": {"open": "", "close": "", "content": "text"}, + "content": {"content": "text"}, + "tool_calls": { + "open_pattern": r'', + "close": "", + "repeats": True, + "content": "xml-inline", + "content_args": { + "tag_pattern": r'(?P.*?)' + }, + "transform": { + "type": "function", + "function": {"name": "{name}", "arguments": "{content}"}, + }, + }, + }, +} +ANSWER_TEMPLATE = deepcopy(TEMPLATE) +ANSWER_TEMPLATE["fields"]["content"] = { + "open": "", + "close": "", + "content": "text", +} + + +def tool(parameters, name="f", strict=True) -> dict[str, Any]: + return {"name": name, "parameters": parameters, "strict": strict} + + +WEATHER = tool( + { + "type": "object", + "properties": {"city": {"type": "string"}, "days": {"type": "integer"}}, + "required": ["city"], + }, + name="get.weather", +) +PING = {"name": "ping"} +CALL = 'Paris' + + +def build( + template=TEMPLATE, + functions=(WEATHER, PING), + *, + parallel=True, + thinking=True, + answer=False, + answer_format=None, + sorted_arguments=False, +): + fields = load_response_template(template).fields + return tool_call_structural_tag( + fields["tool_calls"], + list(functions), + parallel_tool_calls=parallel, + thinking=fields["thinking"] if thinking else None, + answer=fields["content"] if answer else None, + answer_format=answer_format, + sorted_arguments=sorted_arguments, + ) + + +def accepts(tag, text): + return _is_grammar_accept_string(Grammar.from_structural_tag(json.dumps(tag)), text) + + +def call(*arguments): + body = "".join(f'{value}' for key, value in arguments) + return f'{body}' + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + ("plan" + CALL, True), + (CALL, True), + (CALL.removesuffix(""), True), + ( + 'Paris\n' + '3', + True, + ), + ( + 'Paris' + 'soon', + False, + ), + ('3', False), + (CALL + 'x', True), + ('', False), + ("planThe answer.", False), + ], +) +def test_forced_calls_follow_the_template_format(text, accepted): + assert accepts(build(), text) is accepted + + +def test_single_call_when_parallel_calls_are_disabled(): + tag = build(parallel=False) + + assert accepts(tag, CALL) + assert not accepts(tag, CALL + CALL) + + +def test_thinking_is_rejected_when_not_requested(): + assert not accepts(build(thinking=False), "plan" + CALL) + + +@pytest.mark.parametrize( + ("text", "accepted"), + [ + ("planSunny.", True), + ("Sunny.", True), + ("plan" + CALL, True), + ("Sunny." + CALL, False), + ('soon', False), + ], +) +def test_answer_field_lets_the_model_reply_instead(text, accepted): + assert accepts(build(ANSWER_TEMPLATE, answer=True), text) is accepted + + +@pytest.mark.parametrize("required", [False, True]) +@pytest.mark.parametrize( + ("text", "optional_only"), + [ + ("plan" + CALL, False), + ("\nplanSunny.", False), + (" \nSunny.", False), + (CALL, True), + ("Sunny.", True), + ("Sunny.", True), + ("<Sunny.", True), + ], +) +def test_required_reasoning_opens_the_message_with_nonempty_reasoning( + required, text, optional_only +): + fields = load_response_template(ANSWER_TEMPLATE).fields + tag = tool_call_structural_tag( + fields["tool_calls"], + [WEATHER], + parallel_tool_calls=True, + thinking=fields["thinking"], + answer=fields["content"], + message_start=[""], + requires_reasoning=required, + ) + + assert accepts(tag, text) is (not required or not optional_only) + + +def test_answer_rest_finishes_the_open_answer(): + answer = load_response_template(ANSWER_TEMPLATE).fields["content"] + tag = answer_rest_structural_tag(answer) + + assert accepts(tag, "Sunny.") + assert accepts(tag, "Sunny.") + assert not accepts(tag, "Sunny." + CALL) + + +def test_answer_without_literal_opener_is_not_derived(): + assert build(answer=True) is None + + +def test_answer_holds_the_answer_format(): + schema = {"type": "json_schema", "json_schema": {"type": "object"}} + with_tools = build(ANSWER_TEMPLATE, answer=True, answer_format=schema) + answer_only = build(ANSWER_TEMPLATE, (), answer=True, answer_format=schema) + + for tag in (with_tools, answer_only): + assert accepts(tag, 'plan{"a": 1}') + assert not accepts(tag, "Sunny.") + assert accepts(with_tools, CALL) + assert not accepts(answer_only, CALL) + + +def test_json_arguments_follow_the_schema(): + template = deepcopy(TEMPLATE) + template["fields"]["tool_calls"].update(content="json", content_args={}) + tag = build(template, [WEATHER]) + + assert accepts(tag, '{"city": "Paris"}') + assert not accepts(tag, '{"days": 3}') + + +END_X = {"tag_pattern": r'(?P.*?)X'} + + +@pytest.mark.parametrize( + "change", + [ + {"content": "json", "content_args": {"unquoted_keys": True}}, + { + "transform": { + "type": "function", + "function": {"name": "fixed", "arguments": "{content}"}, + } + }, + {"content": "kv-lines", "content_args": {}}, + # Framing in which argument text can spell a closer, where the parser + # would end the call. + {"close": ["}X", ""], "content_args": END_X}, + { + "close": ["}X", ""], + "content_args": {**END_X, "value_parser": {"name": "json"}}, + }, + {"close": ["7X", ""], "content_args": END_X}, + {"close": [">X", ""]}, + {"close": [">", ""]}, + {"close": [">"]}, + {"close": ["}", ""], "content": "json", "content_args": {}}, + ], +) +def test_underivable_tool_formats_are_not_derived(change): + template = deepcopy(TEMPLATE) + template["fields"]["tool_calls"].update(change) + + assert build(template) is None + + +def test_non_literal_thinking_delimiters_are_not_derived(): + template = deepcopy(TEMPLATE) + del template["fields"]["thinking"]["open"] + template["fields"]["thinking"]["open_pattern"] = r"" + + assert build(template) is None + assert build(template, thinking=False) is not None + + +LEVEL = {"type": "object", "properties": {"level": {"$ref": "#/$defs/level"}}} + + +@pytest.mark.parametrize( + ("schema", "value", "accepted"), + [ + ({"type": "integer", "nullable": True}, "null", True), + ({"type": "integer", "nullable": True}, "soon", False), + ({"type": ["integer", "null"]}, "null", True), + ({"anyOf": [{"type": "integer"}, {"type": "boolean"}]}, "true", True), + ({"anyOf": [{"type": "integer"}, {"type": "boolean"}]}, "soon", False), + ({"type": ["string", "integer"]}, "soon", True), + ({"enum": ["celsius", "fahrenheit"]}, "celsius", True), + ({"enum": ["celsius", "fahrenheit"]}, "kelvin", False), + ({"type": "string", "const": "on"}, "off", False), + (LEVEL, '{"level": 2}', True), + (LEVEL, '{"level": "high"}', False), + ], +) +def test_argument_values_follow_how_the_parser_reads_them(schema, value, accepted): + parameters = { + "type": "object", + "properties": {"v": schema}, + "required": ["v"], + "$defs": {"level": {"type": "integer"}}, + } + + assert accepts(build(functions=[tool(parameters)]), call(("v", value))) is accepted + + +@pytest.mark.parametrize( + ("schema", "strict", "derivable"), + [ + # The parser reads values without a type as text. + ({"enum": [1, 2]}, True, False), + ({"$ref": "#/$defs/level"}, True, False), + ({"$ref": "#/$defs/level"}, False, True), + # Raw text cannot hold its closer. + ({"enum": ["ab"]}, True, False), + ({"type": "object", "properties": {"x": {"$ref": "#/x"}}}, True, False), + (False, True, False), + # A match could write the start of the value's end. + ({"type": "string", "pattern": "^[a-z<]+$"}, True, False), + ({"type": "string", "pattern": "^[a-z<]+$"}, False, True), + ({"type": "string", "pattern": "^a.b$"}, True, False), + # XGrammar holds length bounds only in JSON. + ({"type": "string", "minLength": 3}, True, False), + ({"type": "string", "format": "ipv6"}, True, False), + ], +) +def test_schemas_text_values_cannot_hold_are_not_derived(schema, strict, derivable): + parameters = { + "type": "object", + "properties": {"v": schema}, + "$defs": {"level": {"type": "integer"}}, + } + + assert (build(functions=[tool(parameters, strict=strict)]) is not None) is derivable + + +@pytest.mark.parametrize( + ("schema", "value", "accepted"), + [ + ({"type": "string", "format": "email"}, "a.b@x.co", True), + ({"type": "string", "format": "email"}, "ab", False), + ({"type": "string", "format": "date-time"}, "2026-10-06T12:00:00Z", True), + ({"type": "string", "format": "date-time"}, "2026-13-06T12:00:00Z", False), + ({"type": "string", "pattern": "^[A-Z]{3}$"}, "ABC", True), + ({"type": "string", "pattern": "^[A-Z]{3}$"}, "ABCD", False), + # Like XGrammar, the format overrides the pattern and the length bounds. + ( + {"type": "string", "format": "uuid", "pattern": "^x$", "maxLength": 1}, + "123e4567-e89b-12d3-a456-426614174000", + True, + ), + ], +) +def test_text_values_hold_string_formats_and_patterns(schema, value, accepted): + parameters = {"type": "object", "properties": {"v": schema}, "required": ["v"]} + + assert accepts(build(functions=[tool(parameters)]), call(("v", value))) is accepted + + +def test_values_cannot_hold_a_closer(): + template = deepcopy(TEMPLATE) + template["fields"]["tool_calls"]["close"] = ["", "STOP"] + text = {"type": "object", "properties": {"v": {"type": "string"}}} + strings = {"type": "array", "items": {"type": "string"}} + json_text = {"type": "object", "properties": {"v": strings}} + + assert accepts(build(template, [tool(text)]), call(("v", "ab"))) + assert not accepts(build(template, [tool(text)]), call(("v", "aSTOPb"))) + assert not accepts( + build(template, [tool(None, strict=False)]), call(("v", "aSTOPb")) + ) + assert accepts(build(template, [tool(json_text)]), call(("v", '["ab"]'))) + assert not accepts(build(template, [tool(json_text)]), call(("v", '["aSTOPb"]'))) + + +def test_extra_arguments_need_additional_properties(): + parameters = {"type": "object", "properties": {"city": {"type": "string"}}} + extra = call(("city", "Paris"), ("zip", "75001")) + + assert not accepts(build(functions=[tool(parameters)]), extra) + # Next to declared keys, an extra one could repeat them unchecked. + allowed = {**parameters, "additionalProperties": True} + assert build(functions=[tool(allowed)]) is None + free_form = {"type": "object", "additionalProperties": True} + assert accepts(build(functions=[tool(free_form)]), extra) + + +def test_only_strict_tools_follow_their_schema(): + off_schema = call(("days", "soon")) + parameters = WEATHER["parameters"] + + assert accepts(build(functions=[tool(parameters, strict=False)]), off_schema) + assert not accepts(build(functions=[tool(parameters)]), off_schema) + assert accepts(build(functions=[tool(None)]), call()) + assert not accepts(build(functions=[tool(None)]), call(("a", "x"))) + + +def test_annotations_keep_strict_arguments_to_their_schema(): + parameters = { + **WEATHER["parameters"], + "$schema": "https://json-schema.org/draft/2020-12/schema", + "$id": "weather", + } + tag = build(functions=[tool(parameters)]) + + assert accepts(tag, call(("city", "Paris"), ("days", "3"))) + assert not accepts(tag, call(("city", "Paris"), ("days", "soon"))) + + +@pytest.mark.parametrize( + ("arguments", "accepted", "accepted_sorted"), + [ + ((("path", "a"), ("content", "b")), True, True), + ((("content", "b"), ("path", "a")), False, True), + ((("path", "a"),), True, True), + ((("content", "b"),), False, False), + ((("path", "a"), ("path", "c")), False, False), + ((("path", "a"), ("mode", "c")), False, False), + ], +) +def test_strict_arguments_follow_declared_or_sorted_order( + arguments, accepted, accepted_sorted +): + parameters = { + "properties": {"path": {"type": "string"}, "content": {"type": "string"}}, + "required": ["path"], + } + tags = [ + build(functions=[tool(parameters)], sorted_arguments=sorted_arguments) + for sorted_arguments in (False, True) + ] + + assert [accepts(tag, call(*arguments)) for tag in tags] == [ + accepted, + accepted_sorted, + ] + + +LOOKAHEAD_TEMPLATE = deepcopy(TEMPLATE) +LOOKAHEAD_TEMPLATE["fields"]["tool_calls"]["content_args"]["tag_pattern"] = ( + r'(?P.*?)\n(?:|(?=Paris\n\n3\n', True, True), + ('Paris\n3\n', True, True), + ('3\nParis\n', True, False), + ('Paris\nLyon\n', True, False), + ('Paris\n3\n', True, False), + ('Paris\nParis\na c\n', True, True), + ], +) +def test_values_end_where_a_lookahead_end_does(body, accepted_loose, accepted_strict): + parameters = WEATHER["parameters"] + tags = [ + build(LOOKAHEAD_TEMPLATE, [tool(parameters, strict=strict)]) + for strict in (False, True) + ] + text = f'{body}' + + assert [accepts(tag, text) for tag in tags] == [accepted_loose, accepted_strict] + + +@pytest.mark.parametrize("lookahead", [False, True]) +@pytest.mark.parametrize( + ("keys", "accepted"), + [ + ("b", True), + ("ab", True), + ("bc", True), + ("bd", True), + ("be", True), + ("abde", True), + ("bcde", True), + ("", False), + ("a", False), + ("acd", False), + ("bdc", False), + ("bcc", False), + ], +) +def test_last_parameter_follows_every_required_one(keys, accepted, lookahead): + parameters = { + "properties": {key: {"type": "string"} for key in "abcde"}, + "required": ["b"], + } + tag = build(LOOKAHEAD_TEMPLATE, [tool(parameters)]) + ends = ["\n" if lookahead else "\n"] * (len(keys) - 1) + ["\n"] + body = "".join(f'x{e}' for k, e in zip(keys, ends)) + + assert accepts(tag, f'{body}') is accepted + + +def test_required_argument_without_schema_is_still_required(): + parameters = {"type": "object", "properties": {}, "required": ["a"]} + tag = build(functions=[tool(parameters)]) + + assert accepts(tag, call(("a", "x"))) + assert not accepts(tag, call()) + assert build(functions=[tool({**parameters, "required": ["a", 1]})]) is None + + +@pytest.mark.parametrize( + ("function", "match"), + [ + (tool({}, name="get weather"), "tool name"), + (tool({"properties": {'a"b': {"type": "string"}}}), "parameter name"), + ], +) +def test_unencodable_names_are_rejected(function, match): + with pytest.raises(ValueError, match=match): + build(functions=[function]) diff --git a/vllm/entrypoints/openai/responses/protocol.py b/vllm/entrypoints/openai/responses/protocol.py index 3ddf81b288..39d5ff9961 100644 --- a/vllm/entrypoints/openai/responses/protocol.py +++ b/vllm/entrypoints/openai/responses/protocol.py @@ -53,6 +53,7 @@ from openai_harmony import Message as OpenAIHarmonyMessage from pydantic import ( BeforeValidator, Field, + PrivateAttr, ValidationError, field_serializer, model_validator, @@ -469,6 +470,14 @@ class ResponsesRequest(OpenAIBaseModel): max_output_tokens_param="max_output_tokens", ) + _grammar_from_parser: bool = PrivateAttr(default=False) + """CAUTION: Should only be set by the parser's adjust_request.""" + + _parser_structured_outputs: StructuredOutputsParams | None = PrivateAttr( + default=None + ) + """Parser-derived constraints that preserve the public request's text metadata.""" + _DEFAULT_SAMPLING_PARAMS = { "temperature": 1.0, "top_p": 1.0, @@ -478,6 +487,8 @@ class ResponsesRequest(OpenAIBaseModel): def extract_structured_outputs(self) -> StructuredOutputsParams | None: """Normalize request constraints into ``StructuredOutputsParams``.""" + if self._parser_structured_outputs is not None: + return self._parser_structured_outputs if self.text is None or self.text.format is None: return self.structured_outputs diff --git a/vllm/entrypoints/openai/responses/serving.py b/vllm/entrypoints/openai/responses/serving.py index dd154f968b..ba9d95b428 100644 --- a/vllm/entrypoints/openai/responses/serving.py +++ b/vllm/entrypoints/openai/responses/serving.py @@ -526,6 +526,7 @@ class OpenAIServingResponses(GenerateBaseServing): trace_headers=trace_headers, session_id=session_id, reasoning_parser_kwargs=reasoning_parser_kwargs, + reasoning_ended=True if request._grammar_from_parser else None, ) generators.append(generator) @@ -647,6 +648,7 @@ class OpenAIServingResponses(GenerateBaseServing): trace_headers: Mapping[str, str] | None = None, session_id: str | None = None, reasoning_parser_kwargs: dict[str, Any] | None = None, + reasoning_ended: bool | None = None, ): max_model_len = self.model_config.max_model_len cache_salt = cast(str | None, engine_input.get("cache_salt")) @@ -673,6 +675,7 @@ class OpenAIServingResponses(GenerateBaseServing): priority=priority, session_id=session_id, reasoning_parser_kwargs=reasoning_parser_kwargs, + reasoning_ended=reasoning_ended, ) async for res in generator: diff --git a/vllm/envs.py b/vllm/envs.py index ab04f405d7..5cd1d3eecb 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -246,6 +246,7 @@ if TYPE_CHECKING: VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE: int = 163840 VLLM_TOOL_PARSE_REGEX_TIMEOUT_SECONDS: int = 1 VLLM_ENFORCE_STRICT_TOOL_CALLING: bool = True + VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT: bool = True VLLM_MQ_MAX_CHUNK_BYTES_MB: int = 16 VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS: int = 300 VLLM_WORKER_SHUTDOWN_TIMEOUT_SECONDS: int = 5 @@ -1844,6 +1845,11 @@ environment_variables: dict[str, Callable[[], Any]] = { "VLLM_ENFORCE_STRICT_TOOL_CALLING": lambda: ( os.getenv("VLLM_ENFORCE_STRICT_TOOL_CALLING", "True").lower() in ("true", "1") ), + # Constrain every `beam` output to its response template's format, also + # without tool or structured output constraints. + "VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT": lambda: ( + os.getenv("VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT", "True").lower() in ("true", "1") + ), # Control the max chunk bytes (in MB) for the rpc message queue. # Object larger than this threshold will be broadcast to worker # processes via zmq. diff --git a/vllm/parser/beam.py b/vllm/parser/beam.py new file mode 100644 index 0000000000..9f99e321f1 --- /dev/null +++ b/vllm/parser/beam.py @@ -0,0 +1,262 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Parser for Beam checkpoints.""" + +from __future__ import annotations + +import json +from typing import TYPE_CHECKING, Any, cast + +from openai.types.responses import ToolChoiceFunction + +from vllm import envs +from vllm.entrypoints.openai.chat_completion.protocol import ( + ChatCompletionNamedToolChoiceParam, + ChatCompletionRequest, +) +from vllm.exceptions import VLLMValidationError +from vllm.logger import init_logger +from vllm.parser.abstract_parser import ( + _xgrammar_supports, + structured_outputs_to_format, +) +from vllm.parser.chat_parsing import ResponseParser +from vllm.parser.chat_parsing.response_templates import ResponseTemplate +from vllm.parser.response_template import ( + CONTENT_FIELD, + THINKING_FIELD, + TOOL_FIELD, + ResponseTemplateParser, + ResponseTemplateReasoningParser, + ResponseTemplateToolParser, + _as_tool_dicts, + _is_named_tool_choice, + _tool_parsing_enabled, +) +from vllm.parser.response_template_grammar import ( + _openers, + answer_rest_structural_tag, + message_start_anchors, + tool_call_structural_tag, +) +from vllm.sampling_params import StructuredOutputsParams +from vllm.tool_parsers.structural_tag_registry import resolve_tool_strictness + +if TYPE_CHECKING: + from vllm.entrypoints.openai.responses.protocol import ResponsesRequest + +logger = init_logger(__name__) + + +def _named_tool(request: ChatCompletionRequest | ResponsesRequest) -> str: + choice = request.tool_choice + if isinstance(choice, ToolChoiceFunction): + return choice.name + return cast(ChatCompletionNamedToolChoiceParam, choice).function.name + + +def _starts_message(template: ResponseTemplate, prefix: str) -> bool: + """Whether output after `prefix` starts the assistant message.""" + ends = [match.end() for match in template.start_anchor_re.finditer(prefix)] + return bool(ends) and ends[-1] == len(prefix) + + +def _open_field(template: ResponseTemplate, prefix: str) -> str | None: + """The field that `prefix` leaves open with its explicit opener, if any.""" + field = None + for event in ResponseParser(template, prefix=prefix).initial_events: + if event["type"] == "region_open" and event["end"] > event["start"]: + field = event["field"] + elif event["type"] in {"region_close", "region_malformed"}: + field = None + return field + + +class BeamReasoningParser(ResponseTemplateReasoningParser): + """Selects `BeamParser` as the reasoning parser.""" + + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + return BeamParser + + @property + def reasoning_start_str(self) -> str | None: + if literal := super().reasoning_start_str: + return literal + thinking = self.response_template.fields.get(THINKING_FIELD) + if thinking is None: + return None + openers = _openers(thinking, message_start_anchors(self.response_template)) + if not openers: + return None + # Budget tracking needs one marker for both a prompt continuation + # and a complete assistant message in the checkpoint's pattern. + marker = min(openers, key=len) + return marker if all(opener.endswith(marker) for opener in openers) else None + + +class BeamToolParser(ResponseTemplateToolParser): + """Selects `BeamParser` as the tool parser.""" + + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + return BeamParser + + +class BeamParser(ResponseTemplateParser): + """Parse Beam output with the checkpoint's `response_template` metadata, + and constrain it with a grammar derived from that template.""" + + reasoning_parser_cls = BeamReasoningParser + tool_parser_cls = BeamToolParser + + def __init__(self, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + chat_template_kwargs = kwargs.get("chat_template_kwargs") or {} + # From the message start, the grammar opens the message with non-empty + # reasoning, unless the request turns thinking off. + self._requires_reasoning = bool( + chat_template_kwargs.get("enable_thinking", True) + ) + + def _constrain_request( + self, + request: ChatCompletionRequest | ResponsesRequest, + ) -> ChatCompletionRequest | ResponsesRequest: + message_start = _starts_message(self.response_template, self._prefix) + tag = self._template_tag(request, message_start=message_start) + if tag is not None: + grammar = StructuredOutputsParams(structural_tag=json.dumps(tag)) + if isinstance(request, ChatCompletionRequest): + request.structured_outputs = grammar + request.response_format = None + else: + request._parser_structured_outputs = grammar + # From the message start, the tag also covers the reasoning. + request._grammar_from_parser = message_start + return request + + def _template_tag( + self, + request: ChatCompletionRequest | ResponsesRequest, + *, + message_start: bool, + ) -> dict[str, Any] | None: + """The template grammar for `request`, or `None` to leave its output + to the request's own constraint, if any.""" + fields = self.response_template.fields + named = _is_named_tool_choice(request) + tools_enabled = ( + self._tool_parser is not None + and bool(request.tools) + and _tool_parsing_enabled( + request, enable_auto_tools=self._enable_auto_tools + ) + ) + forced = tools_enabled and (request.tool_choice == "required" or named) + # The prompt may still list the tools, so only the grammar keeps the + # model from calling them. + tools_forbidden = ( + bool(request.tools) + and request.tool_choice == "none" + and TOOL_FIELD in fields + ) + constrained_tools = None + if tools_enabled and (forced or envs.VLLM_ENFORCE_STRICT_TOOL_CALLING): + assert self._tool_parser is not None + strict_level = self.tool_strict_level + if strict_level is None: + strict_level = self._tool_parser.default_tool_strict_level + constrained_tools = resolve_tool_strictness( + request.tools or [], request.tool_choice, strict_level + ) + output = request.extract_structured_outputs() + if forced and output is not None: + logger.warning_once( + "structured outputs are ignored because tool_choice forces tool call." + ) + output = None + answer = fields.get(CONTENT_FIELD) + if output is not None and ( + output.structural_tag is not None + or answer is None + or answer.open_re is None + or not _xgrammar_supports(output) + ): + if constrained_tools is not None: + logger.warning_once( + "response_template parsing cannot enforce strict tools " + "alongside this structured output." + ) + return None + # Without a tool or structured output constraint, the output still + # follows the template. + always_constrain = envs.VLLM_BEAM_ALWAYS_CONSTRAIN_OUTPUT and not forced + if ( + constrained_tools is None + and output is None + and not tools_forbidden + and not always_constrain + ): + return None + + # Inside an open answer, only the always-on grammar applies, finishing a + # free answer. Inside open reasoning, the grammar needs the reasoning + # gate to defer it to the reasoning's end. + open_field = _open_field(self.response_template, self._prefix) + reasoning_gated = ( + self._reasoning_parser is not None and request.include_reasoning + ) + tag = None + if open_field == CONTENT_FIELD: + if always_constrain and output is None: + tag = answer_rest_structural_tag(fields[CONTENT_FIELD]) + elif open_field != THINKING_FIELD or reasoning_gated: + functions = [] + if tools_enabled: + tools = _as_tool_dicts(constrained_tools or request.tools) + functions = [tool.get("function", tool) for tool in tools] + if named: + name = _named_tool(request) + functions = [f for f in functions if f.get("name") == name] + answer_format = ( + None if output is None else structured_outputs_to_format(output) + ) + try: + tag = tool_call_structural_tag( + fields.get(TOOL_FIELD), + functions, + parallel_tool_calls=( + not named and request.parallel_tool_calls is not False + ), + thinking=fields.get(THINKING_FIELD), + answer=None if forced else answer, + answer_format=( + None if answer_format is None else answer_format.model_dump() + ), + message_start=( + message_start_anchors(self.response_template) + if message_start + else None + ), + requires_reasoning=self._requires_reasoning, + # Beam drifts from schemas the grammar does not enforce, + # and also writes arguments sorted by name. + sorted_arguments=True, + strict_arguments=True, + ) + except ValueError as exc: + raise VLLMValidationError(str(exc), parameter="tools") from exc + if tag is None and forced: + kind = "named" if named else "required" + raise VLLMValidationError( + f"response_template parsing cannot guarantee {kind} tool choice " + "for this request", + parameter="tool_choice", + ) + if tag is None and constrained_tools is not None: + logger.warning_once( + "response_template parsing cannot enforce strict tools for this " + "request." + ) + return tag diff --git a/vllm/parser/parser_manager.py b/vllm/parser/parser_manager.py index d6fdbce19c..b6f00cacd8 100644 --- a/vllm/parser/parser_manager.py +++ b/vllm/parser/parser_manager.py @@ -16,8 +16,6 @@ if TYPE_CHECKING: logger = init_logger(__name__) -HF_PARSER = "hf" - class ParserManager: """Provides a unified Parser by composing reasoning and tool parser adapters.""" @@ -122,8 +120,10 @@ class ParserManager: tags (``--tool-strict-level``), or ``None`` for the tool parser's default level. tokenizer: Tokenizer the composed parser is validated against - at startup (see :meth:`validate_parser_tokenizer`). ``None`` - skips the check (e.g. ``skip_tokenizer_init``). + at startup (see :meth:`validate_parser_tokenizer`), including + its `response_template` metadata when a response-template + parser is selected. ``None`` skips the check (e.g. + ``skip_tokenizer_init``). Returns: A Parser class, or None if neither parser is specified. @@ -184,39 +184,40 @@ class ParserManager: HarmonyParser.tool_strict_level = strict_level return HarmonyParser - if HF_PARSER in (reasoning_parser_name, tool_parser_name): - if {reasoning_parser_name, tool_parser_name} - { - HF_PARSER, - None, - "", - }: - raise TypeError( - "The hf parser cannot be combined with other " - "reasoning or tool call parsers" - ) - from vllm.parser.response_template import ( - ResponseTemplateParser, - validate_tokenizer_response_template, - ) + from vllm.parser.response_template import ( + response_template_parser_cls, + validate_tokenizer_response_template, + ) + template_parsers = { + response_template_parser_cls(parser_cls) + for parser_cls in (reasoning_parser_cls, tool_parser_cls) + if parser_cls is not None + } + template_parser_cls = next(iter(template_parsers - {None}), None) + if template_parser_cls is not None: + if len(template_parsers) > 1: + raise TypeError( + f"--reasoning-parser {reasoning_parser_name} and " + f"--tool-call-parser {tool_parser_name} cannot be combined" + ) if tokenizer is not None: validate_tokenizer_response_template( tokenizer, - reasoning=reasoning_parser_cls is not None, - tools=tool_parser_cls is not None, + reasoning_parser=reasoning_parser_name, + # Without auto tool choice, no tool parser is selected. + tool_parser=tool_parser_name if tool_parser_cls else None, ) - - r_cls = reasoning_parser_cls - t_cls = tool_parser_cls - auto_tools = enable_auto_tools - - class _ResponseTemplateParser(ResponseTemplateParser): - reasoning_parser_cls = r_cls - tool_parser_cls = t_cls - tool_strict_level = strict_level - _enable_auto_tools = auto_tools - - return _ResponseTemplateParser + return type( + template_parser_cls.__name__, + (template_parser_cls,), + { + "reasoning_parser_cls": reasoning_parser_cls, + "tool_parser_cls": tool_parser_cls, + "tool_strict_level": strict_level, + "_enable_auto_tools": enable_auto_tools, + }, + ) if reasoning_parser_name == "kimi_k3" or tool_parser_name == "kimi_k3": from vllm.parser.kimi_k3 import KimiK3Parser diff --git a/vllm/parser/response_template.py b/vllm/parser/response_template.py index b1665f1772..ea95567305 100644 --- a/vllm/parser/response_template.py +++ b/vllm/parser/response_template.py @@ -7,6 +7,7 @@ from __future__ import annotations import copy import json from collections.abc import Callable, Iterable, Sequence +from math import inf from typing import TYPE_CHECKING, Any, cast from openai.types.responses import ToolChoiceCustom, ToolChoiceFunction @@ -28,6 +29,7 @@ from vllm.parser.chat_parsing.response_templates import ( ResponseTemplate, load_response_template, ) +from vllm.parser.response_template_grammar import message_start_anchors from vllm.reasoning.abs_reasoning_parsers import ReasoningParser from vllm.renderers.chat_utils import make_tool_call_id from vllm.tool_parsers.abstract_tool_parser import ToolParser @@ -135,27 +137,38 @@ def validate_response_template_for_serving( def validate_tokenizer_response_template( tokenizer: Any, *, - reasoning: bool, - tools: bool, + reasoning_parser: str | None, + tool_parser: str | None, ) -> None: - """Fail at startup if the tokenizer cannot serve the requested parsing.""" + """Fail at startup if the tokenizer cannot serve the selected parsers, + named as given on the command line.""" template = resolve_response_template(tokenizer) if template is None: + selected = " ".join( + f"{flag} {name}" + for flag, name in ( + ("--reasoning-parser", reasoning_parser), + ("--tool-call-parser", tool_parser), + ) + if name + ) raise TypeError( - "The hf parser requires `response_template` metadata " + f"{selected} requires `response_template` metadata " "in the tokenizer configuration" ) try: fields = validate_response_template_for_serving(template).fields except (TypeError, ValueError) as exc: raise TypeError(f"Invalid response_template metadata: {exc}") from exc - if reasoning and THINKING_FIELD not in fields: + if reasoning_parser and THINKING_FIELD not in fields: raise TypeError( - "--reasoning-parser hf requires a response_template with a thinking field" + f"--reasoning-parser {reasoning_parser} requires a response_template " + "with a thinking field" ) - if tools and TOOL_FIELD not in fields: + if tool_parser and TOOL_FIELD not in fields: raise TypeError( - "--tool-call-parser hf requires a response_template with a tool_calls field" + f"--tool-call-parser {tool_parser} requires a response_template " + "with a tool_calls field" ) @@ -292,7 +305,7 @@ def _serving_template( template = resolve_response_template(tokenizer, template) if template is None: raise ValueError( - "The hf parser requires `response_template` " + "Parsing with a response template requires `response_template` " "metadata in the tokenizer configuration" ) validate_response_template_for_serving(template) @@ -319,12 +332,17 @@ def _to_tool_calls(value: Any) -> list[tuple[str, str]]: class ResponseTemplateReasoningParser(ReasoningParser): - """Reasoning end detection for `--reasoning-parser hf`. + """Reasoning end detection for a response-template parser. `ResponseTemplateParser` parses the output; this parser tells it and the structured-output engine where reasoning ends. """ + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + """The unified parser that selecting this reasoning parser stands for.""" + return ResponseTemplateParser + def __init__( self, tokenizer: TokenizerLike, @@ -400,8 +418,26 @@ class ResponseTemplateReasoningParser(ReasoningParser): class ResponseTemplateToolParser(ToolParser): - """Registers `--tool-call-parser hf`; `ResponseTemplateParser` extracts the - tool calls.""" + """Selects a response-template parser as the tool parser; + `ResponseTemplateParser` extracts the tool calls.""" + + @classmethod + def unified_parser_cls(cls) -> type[ResponseTemplateParser]: + """The unified parser that selecting this tool parser stands for.""" + return ResponseTemplateParser + + +def response_template_parser_cls( + split_parser_cls: type[ReasoningParser] | type[ToolParser], +) -> type[ResponseTemplateParser] | None: + """The response-template parser that a registered reasoning or tool parser + stands for, or `None` for any other parser.""" + if issubclass( + split_parser_cls, + (ResponseTemplateReasoningParser, ResponseTemplateToolParser), + ): + return split_parser_cls.unified_parser_cls() + return None class _DeltaBuilder: @@ -441,23 +477,50 @@ class _ResponseStream: self.include_reasoning = include_reasoning self.parse_tools = parse_tools self.new_tool_call_id = new_tool_call_id + self.template = template self.parser = ResponseParser(template, prefix=prefix, tools=tools) self._prefix_end = len(self.parser.input_text) self._next_tool_index = 0 self._finalized = False + self._open_region: dict[str, Any] | None = None + self._track(self.parser.initial_events) def feed(self, text: str, *, finished: bool) -> DeltaMessage | None: if self._finalized: return None delta = _DeltaBuilder() - self._route(self.parser.feed(text), delta) + self._route(self._track(self.parser.feed(text)), delta) if finished: self._finalized = True if len(self.parser.input_text) > self._prefix_end: - _, events = self.parser.finalize() - self._route(events, delta) + events = self.parser.feed(self._closer_rest()) + self._route(events + self.parser.finalize()[1], delta) return delta.build() + def _track(self, events: list[dict[str, Any]]) -> list[dict[str, Any]]: + for event in events: + if event["type"] == "region_open": + self._open_region = event + elif event["type"] in ("region_close", "region_malformed"): + self._open_region = None + return events + + def _closer_rest(self) -> str: + """The rest of a closer the output stopped in after ending its region, + when only the next message's start anchor is missing.""" + if self._open_region is None: + return "" + field = self.template.fields[self._open_region["field"]] + text = self.parser.input_text[self._open_region["end"] :] + for closer in field.close_literals or (): + for anchor in message_start_anchors(self.template): + if not anchor or not closer.endswith(anchor) or closer == anchor: + continue + for cut in range(len(closer) - 1, len(closer) - len(anchor) - 1, -1): + if text.endswith(closer[:cut]): + return closer[cut:] + return "" + def _route(self, events: Sequence[dict[str, Any]], delta: _DeltaBuilder) -> None: for event in events: field = event.get("field") @@ -499,6 +562,54 @@ class _ResponseStream: self._next_tool_index += 1 +class _ReasoningTokenCounter: + """Counts the reasoning tokens of one growing generation, parsing each + token once.""" + + def __init__( + self, template: ResponseTemplate, prefix: str, tokenizer: TokenizerLike + ) -> None: + self.tokenizer = tokenizer + self.parser = ResponseParser(template, prefix=prefix) + self.start = len(self.parser.input_text) + self.token_ids: list[int] = [] + self.token_ends: list[int] = [] + # Input spans between the delimiters of thinking regions. + self.spans: list[tuple[int, float]] = [] + self.open: int | None = None + self._track(self.parser.initial_events) + + def continues(self, token_ids: Sequence[int]) -> bool: + return list(token_ids[: len(self.token_ids)]) == self.token_ids + + def count(self, token_ids: Sequence[int]) -> int: + new_ids = list(token_ids[len(self.token_ids) :]) + texts = [_decode(self.tokenizer, [token_id]) for token_id in new_ids] + end = self.token_ends[-1] if self.token_ends else self.start + for text in texts: + end += len(text) + self.token_ends.append(end) + self.token_ids.extend(new_ids) + self._track(self.parser.feed("".join(texts))) + # Text held back as a possible closer counts until a closer arrives. + spans = self.spans if self.open is None else [*self.spans, (self.open, inf)] + count, start = 0, self.start + for end in self.token_ends: + count += any(begin < end and start < close for begin, close in spans) + start = end + return count + + def _track(self, events: Iterable[dict[str, Any]]) -> None: + for event in events: + if event["field"] != THINKING_FIELD or event["type"] == "region_chunk": + continue + if event["type"] == "region_open": + self.open = event["end"] + elif self.open is not None: + self.spans.append((self.open, event["start"])) + self.open = None + + class ResponseTemplateParser(Parser): """Parse reasoning, content, and tool calls with one response template.""" @@ -526,6 +637,7 @@ class ResponseTemplateParser(Parser): self._prompt_token_ids: list[int] | None = None self._prefix = "" self._stream: _ResponseStream | None = None + self._reasoning_counter: _ReasoningTokenCounter | None = None def set_prompt_token_ids(self, prompt_token_ids: Sequence[int]) -> None: prompt_token_ids = list(prompt_token_ids) @@ -534,6 +646,7 @@ class ResponseTemplateParser(Parser): self._prompt_token_ids = prompt_token_ids self._prefix = _decode(self.model_tokenizer, prompt_token_ids) self._stream = None + self._reasoning_counter = None def adjust_request( self, @@ -547,6 +660,14 @@ class ResponseTemplateParser(Parser): request.skip_special_tokens = False if hasattr(request, "spaces_between_special_tokens"): request.spaces_between_special_tokens = False + return self._constrain_request(request) + + def _constrain_request( + self, + request: ChatCompletionRequest | ResponsesRequest, + ) -> ChatCompletionRequest | ResponsesRequest: + """Constrain `request` to the tool guarantees it asks for, or reject + those this parser cannot enforce.""" if ( self._tool_parser is not None and request.tools @@ -567,6 +688,21 @@ class ResponseTemplateParser(Parser): or self._reasoning_parser.is_reasoning_end(input_ids) ) + def count_reasoning_tokens(self, token_ids: Sequence[int]) -> int: + """Count the generated tokens whose text the template parses as + reasoning, continuing reasoning left open by the prompt.""" + if ( + self._reasoning_parser is None + or THINKING_FIELD not in self.response_template.fields + ): + return 0 + counter = self._reasoning_counter + if counter is None or not counter.continues(token_ids): + counter = self._reasoning_counter = _ReasoningTokenCounter( + self.response_template, self._prefix, self.model_tokenizer + ) + return counter.count(token_ids) + def parse( self, model_output: str, diff --git a/vllm/parser/response_template_grammar.py b/vllm/parser/response_template_grammar.py new file mode 100644 index 0000000000..40317023f4 --- /dev/null +++ b/vllm/parser/response_template_grammar.py @@ -0,0 +1,1000 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Structural tags for tool calls and answers, derived from a response template.""" + +from __future__ import annotations + +import json +from collections.abc import Callable +from functools import lru_cache +from typing import Any + +import regex as re + +try: + from re import _parser # type: ignore[attr-defined] +except ImportError: # Python 3.10 + import sre_parse as _parser + +from vllm.parser.chat_parsing.content_parsers import parse_content +from vllm.parser.chat_parsing.response_parser import _schema_types +from vllm.parser.chat_parsing.response_templates import ( + ResponseTemplate, + ResponseTemplateField, +) + +_sre_parse = _parser +_REGEX_META = re.compile(r"[.^$*+?{}\[\]\\|()]") +_LITERAL_ESCAPES = {"n": "\n", "r": "\r", "t": "\t"} +# Stands for `^` while expanding a delimiter pattern. +_START = "\0" +# Precedes the text of a final lookahead while expanding a delimiter pattern. +_AHEAD = "\1" +# Arbitrary bound on the strings expanded from one delimiter pattern. +_MAX_LITERALS = 16 +# Templates write at most a newline between parameters. Unbounded whitespace +# would let a model that skipped a parameter's opener write it until `max_tokens`. +_SEPARATOR = {"type": "regex", "pattern": r"\n?"} +_ANY_TEXT = {"type": "any_text"} +# Character classes in XGrammar regexes take whitespace only as escapes. +_BLANKS = {" ": " ", "\t": r"\t", "\n": r"\n", "\r": r"\r"} +# Keywords that describe a schema without constraining its values. +_ANNOTATIONS = frozenset( + { + "title", + "description", + "default", + "examples", + "deprecated", + "$comment", + "$schema", + "$id", + } +) +_OBJECT_KEYWORDS = frozenset( + {"type", "properties", "required", "additionalProperties", "$defs", "definitions"} +) +# Characters JSON writes outside strings. +_JSON_SYNTAX = frozenset(' \t\n\r{}[],:"-+.0123456789eEtrufalsn') +# Pieces of the text of a call for `_spells_closer`. Text and JSON hold no +# closer, and JSON ends with a character it writes outside strings. +_TEXT_PIECE = (r"[\s\S]*", r"[\s\S]+", False) +_JSON_PIECE = ( + r"[\s\S]*", + r"[\s\S]*[" + "".join(map(re.escape, sorted(_JSON_SYNTAX))) + "]", + False, +) +_REPEATS = tuple( + getattr(_sre_parse, name) + for name in ("MAX_REPEAT", "MIN_REPEAT", "POSSESSIVE_REPEAT") + if hasattr(_sre_parse, name) +) +_CATEGORIES = { + _sre_parse.CATEGORY_DIGIT: r"\d", + _sre_parse.CATEGORY_NOT_DIGIT: r"\D", + _sre_parse.CATEGORY_SPACE: r"\s", + _sre_parse.CATEGORY_NOT_SPACE: r"\S", + _sre_parse.CATEGORY_WORD: r"\w", + _sre_parse.CATEGORY_NOT_WORD: r"\W", +} +_HOST_LABEL = r"[a-z0-9](?:[a-z0-9-]*[a-z0-9])?" +_DATE = r"[0-9]{4}-(?:0[1-9]|1[0-2])-(?:0[1-9]|[12][0-9]|3[01])" +_TIME = ( + r"(?:[01][0-9]|2[0-3]):[0-5][0-9]:(?:[0-5][0-9]|60)(?:\.[0-9]+)?" + r"(?:Z|[+-](?:[01][0-9]|2[0-3]):[0-5][0-9])" +) +_OCTET = r"(?:25[0-5]|2[0-4][0-9]|[01]?[0-9]?[0-9])" +_ATOM = r"[A-Za-z0-9_!#$%&'*+/=?^`{|}~-]+" +_URI_CHARS = r"A-Za-z0-9_.~!$&'()*+,;=" +_ESCAPED = r"%[0-9A-Fa-f]{2}" +_PATH = rf"(?:[{_URI_CHARS}:@-]|{_ESCAPED})" +_QUERY = rf"(?:[{_URI_CHARS}:@/?-]|{_ESCAPED})*" +# Raw string forms of the formats XGrammar holds in JSON strings. Emails leave +# out quoted local parts, whose quotes XGrammar writes as JSON escapes. +_STRING_FORMATS = { + "date": _DATE, + "time": _TIME, + "date-time": f"{_DATE}T{_TIME}", + "email": ( + rf"{_ATOM}(?:\.{_ATOM})*@[A-Za-z0-9](?:[A-Za-z0-9-]*[A-Za-z0-9])?" + r"(?:\.[A-Za-z0-9][A-Za-z0-9-]*[A-Za-z0-9])*" + ), + "hostname": rf"{_HOST_LABEL}(?:\.{_HOST_LABEL})*", + "ipv4": rf"(?:{_OCTET}\.){{3}}{_OCTET}", + "uuid": r"[0-9A-Fa-f]{8}(?:-[0-9A-Fa-f]{4}){3}-[0-9A-Fa-f]{12}", + "uri": ( + rf"[A-Za-z][A-Za-z+.-]*:(?://(?:(?:[{_URI_CHARS}:-]|{_ESCAPED})*@)?" + rf"(?:[{_URI_CHARS}-]|{_ESCAPED})*(?::[0-9]*)?(?:/{_PATH}*)*" + rf"|/?(?:{_PATH}+(?:/{_PATH}*)*)?)(?:\?{_QUERY})?(?:#{_QUERY})?" + ), +} + + +def _split_group(pattern: str, name: str) -> tuple[str, str, str] | None: + """Split a named regex group into its prefix, body and suffix.""" + opener = f"(?P<{name}>" + start = pattern.find(opener) + if start < 0: + return None + depth = 0 + in_class = False + index = start + while index < len(pattern): + char = pattern[index] + if char == "\\": + index += 2 + continue + if in_class: + in_class = char != "]" + elif char == "[": + in_class = True + elif char == "(": + depth += 1 + elif char == ")": + depth -= 1 + if depth == 0: + body = pattern[start + len(opener) : index] + return pattern[:start], body, pattern[index + 1 :] + index += 1 + return None + + +def _literals( + pattern: str, *, message_start: bool = False, lookahead: bool = False +) -> list[str] | None: + """Expand literal alternatives, allowing `^` only at the message start and + a final lookahead only with `lookahead`.""" + literals, index = _alternatives(pattern, 0) + if literals is None or index != len(pattern): + return None + if any(literal.count(_AHEAD) > lookahead for literal in literals): + return None + if not message_start: + return None if any(_START in literal for literal in literals) else literals + return [ + literal.removeprefix(_START) + for literal in literals + if _START not in literal[1:] and literal != _START + ] + + +def _alternatives(pattern: str, index: int) -> tuple[list[str] | None, int]: + literals = [] + while True: + branch, index = _concatenation(pattern, index) + if branch is None: + return None, index + literals.extend(branch) + if not pattern.startswith("|", index): + break + index += 1 + literals = list(dict.fromkeys(literals)) + return (literals if len(literals) <= _MAX_LITERALS else None), index + + +def _concatenation(pattern: str, index: int) -> tuple[list[str] | None, int]: + literals = [""] + while index < len(pattern) and pattern[index] not in "|)": + if pattern.startswith(("(?:", "(?="), index): + ahead = pattern[index + 2] == "=" + atom, index = _alternatives(pattern, index + 3) + if atom is None or not pattern.startswith(")", index): + return None, index + index += 1 + if ahead: + atom = [_AHEAD + literal for literal in atom] + elif pattern[index] == "\\" and index + 1 < len(pattern): + escaped = pattern[index + 1] + if escaped.isalnum() and escaped not in _LITERAL_ESCAPES: + return None, index + atom = [_LITERAL_ESCAPES.get(escaped, escaped)] + index += 2 + elif pattern[index] == "^": + atom = [_START] + index += 1 + elif _REGEX_META.fullmatch(pattern[index]): + return None, index + else: + atom = [pattern[index]] + index += 1 + if pattern.startswith("?", index): + atom = ["", *atom] + index += 1 + # A lookahead reads as a literal only at the end of the pattern. + if any(_AHEAD in prefix for prefix in literals): + return None, index + literals = [prefix + suffix for prefix in literals for suffix in atom] + if len(literals) > _MAX_LITERALS: + return None, index + return literals, index + + +def _openers( + field: ResponseTemplateField, message_start: list[str] | None = None +) -> list[str] | None: + """Derive field openers, including continuations of message-start anchors.""" + if field.open_literals or field.open_re is None: + openers = field.open_literals + else: + openers = _literals( + field.open_re.pattern, message_start=message_start is not None + ) + return _continuing(openers, message_start, field.open_re) + + +def _continuing( + literals: list[str] | None, anchors: list[str] | None, opener: re.Pattern +) -> list[str] | None: + """Keep anchor continuations only where the parser finds the opener.""" + if not literals or not anchors: + return literals + rests = [ + literal[len(anchor) :] + for literal in literals + for anchor in anchors + if len(literal) > len(anchor) and literal.startswith(anchor) + ] + rests = [ + rest + for rest in rests + if any(opener.fullmatch(rest, index) for index in range(len(rest))) + ] + return list(dict.fromkeys([*literals, *rests])) + + +def _escape(text: str) -> str: + return _REGEX_META.sub(lambda match: "\\" + match.group(), text) + + +def _call_opener( + opener: tuple[str, str, str], name: str, message_start: list[str] | None +) -> str: + before, _, after = opener + if message_start is not None and before: + literals = _literals(before, message_start=True) + # Text before a name inside a group is no pattern of its own. + if literals: + literals = _continuing(literals, message_start, re.compile(before)) + before = "(?:" + "|".join(map(_escape, literals or ())) + ")" + return before + name + after + + +def _other_names(pattern: str, names: list[str]) -> str | None: + """A regex for the names `pattern` matches besides `names`, or `None` if + `pattern` is no negatable character class repeated.""" + if not names: + return pattern + match = re.fullmatch(r"\\([wS])\+|\[\^((?:\\.|[^\]\\])+)\]\+", pattern) + if match is None: + return None + negated = {"w": r"\W", "S": r"\s"}[match[1]] if match[1] else match[2] + + def other(chars: set[str]) -> str: + return f"[^{negated}{''.join(map(_class_escape, sorted(chars)))}]" + + trie: dict[str, dict] = {} + for name in names: + node = trie + for char in name: + node = node.setdefault(char, {}) + node[""] = {} + + def rest(node: dict[str, dict], first: bool) -> str: + """The rest of a name past `node` that ends on none of `names`.""" + branches = [other(set(node) - {""}) + other(set()) + "*"] + branches += [ + _escape(char) + rest(child, False) for char, child in node.items() if char + ] + return f"(?:{'|'.join(branches)})" + ("" if first or "" in node else "?") + + return rest(trie, True) + + +def _thinking_format( + thinking: ResponseTemplateField, + message_start: list[str] | None, + nonempty: bool = False, +) -> dict[str, Any] | None: + """The reasoning, which holds at least a character if `nonempty`.""" + openers, closers = _openers(thinking, message_start), _closers(thinking) + if not openers or not closers: + return None + text = _sequence(_nonempty_start(closers), _ANY_TEXT) if nonempty else _ANY_TEXT + return _one_of([_tag(opener, text, closers) for opener in openers]) + + +def _nonempty_start(closers: list[str]) -> dict[str, Any]: + """The start of text that is not empty and does not begin with one of the + `closers`. Each branch leaves every closer's spelling before completing it, + on a character that starts no closer, so the text after it cannot hide a + closer's start and a closer written as one token cannot pass for the text.""" + closers = [closer for closer in closers if closer] + starts = {closer[0] for closer in closers} + prefixes = dict.fromkeys( + [""] + [closer[:end] for closer in closers for end in range(1, len(closer))] + ) + branches = [] + for prefix in prefixes: + # A prefix holding a whole closer has already closed the text. + if any(closer in prefix for closer in closers): + continue + nexts = { + closer[len(prefix)] + for closer in closers + if len(closer) > len(prefix) and closer.startswith(prefix) + } + excluded = "".join(map(_class_escape, sorted(starts | nexts))) + branches.append(f"{_escape(prefix)}[^{excluded}]") + return {"type": "regex", "pattern": "|".join(branches)} + + +def _class_escape(char: str) -> str: + if char in _BLANKS: + return _BLANKS[char] + return "\\" + char if char in "\\]^-[" else char + + +def message_start_anchors(template: ResponseTemplate) -> list[str]: + """The strings the assistant message starts after, if literal.""" + if template.start_anchor_literals: + return template.start_anchor_literals + return _literals(template.start_anchor_re.pattern) or [] + + +def _closers(field: ResponseTemplateField) -> list[str] | None: + """Return literal closers, an empty list if absent, or `None` if underivable.""" + if field.close_re is None: + return [] + return field.close_literals or _literals(field.close_re.pattern) + + +# Every request derives the tags for its tools anew. +@lru_cache(maxsize=512) +def _endings(literals: tuple[str, ...]) -> str: + return "|".join( + re.escape(literal[index:]) + for literal in literals + for index in range(len(literal)) + ) + + +def _literal_piece(literals: list[str], before: str = "") -> tuple[str, str, bool]: + either = before + "(?:" + "|".join(map(re.escape, literals)) + ")" + return either, f"{either}|{_endings(tuple(literals))}", True + + +@lru_cache(maxsize=512) +def _spells_closer( + closers: tuple[str, ...], + pieces: tuple[tuple[str, str, bool], ...], + *, + repeats: bool = False, +) -> bool: + """Check whether argument framing can spell a call closer across pieces.""" + texts = [text for text, _, _ in pieces] + loop = f"(?:{''.join(texts)})*" if repeats else "" + # Another call or an answer may follow the call's own closer. + closed = "(?:" + "|".join(map(re.escape, closers)) + r")[\s\S]*" + for index, (_, endings, holds) in enumerate(pieces): + ends = re.compile(endings) + rest = re.compile("".join(texts[index + 1 :]) + loop + closed, re.DOTALL) + for closer in closers: + if holds and ends.fullmatch(closer, partial=True): + return True + # A closer that starts in this piece and ends in the text after it. + for split in range(1, len(closer)): + if ends.fullmatch(closer[:split]) and rest.fullmatch( + closer[split:], partial=True + ): + return True + return False + + +def _const_strings(values: list[str]) -> dict[str, Any]: + return _one_of([{"type": "const_string", "value": value} for value in values]) + + +def _sequence(*elements: dict[str, Any]) -> dict[str, Any]: + return {"type": "sequence", "elements": list(elements)} + + +def _one_of(formats: list[dict[str, Any]]) -> dict[str, Any]: + return formats[0] if len(formats) == 1 else {"type": "or", "elements": formats} + + +def _optional(content: dict[str, Any]) -> dict[str, Any]: + return {"type": "optional", "content": content} + + +def _tag(begin: str, content: dict[str, Any], end: str | list[str]) -> dict[str, Any]: + return {"type": "tag", "begin": begin, "content": content, "end": end} + + +def _json_value_schema(schema: Any, parameters: dict[str, Any]) -> dict[str, Any]: + if not isinstance(schema, dict): + return {} + if schema.get("nullable"): + # XGrammar ignores `nullable`. + schema = {"anyOf": [schema, {"type": "null"}]} + definitions = { + key: parameters[key] for key in ("$defs", "definitions") if key in parameters + } + return {**definitions, **schema} + + +def _string_choices(schema: Any) -> list[str] | None: + if not isinstance(schema, dict): + return None + choices = [schema["const"]] if "const" in schema else schema.get("enum") + if not isinstance(choices, list) or not choices: + return None + return choices if all(isinstance(choice, str) for choice in choices) else None + + +def _typed_by_parser(schema: Any) -> bool: + """Whether the parser reads schema values with their JSON type.""" + if schema is True or _schema_types(schema): + return True + return isinstance(schema, dict) and ( + not set(schema) - _ANNOTATIONS or _string_choices(schema) is not None + ) + + +def _text_format(schema: Any, excludes: list[str]) -> dict[str, Any] | None: + """Constrain raw text, which the `excludes` end, without changing its + parsed value.""" + if schema is True: + return _ANY_TEXT + if isinstance(schema, dict) and {"format", "pattern"} & set(schema): + return _string_regex(schema, excludes) + if not isinstance(schema, dict) or set(schema) - _ANNOTATIONS - { + "type", + "nullable", + "enum", + "const", + "anyOf", + "oneOf", + }: + return None + branches = [*(schema.get("anyOf") or []), *(schema.get("oneOf") or [])] + if "enum" in schema or "const" in schema: + choices = _string_choices(schema) + if choices is None or branches or set(_schema_types(schema)) - {"string"}: + return None + return _const_strings(choices) + if any(_text_format(branch, excludes) is not _ANY_TEXT for branch in branches): + return None + return _ANY_TEXT + + +def _string_regex(schema: dict[str, Any], excludes: list[str]) -> dict[str, Any] | None: + """Hold the `format` or `pattern` of a raw string. Like XGrammar, a format + overrides a pattern, either one drops length bounds, and the pattern + matches the whole string.""" + if set(schema) - _ANNOTATIONS - { + "type", + "format", + "pattern", + "minLength", + "maxLength", + } or _schema_types(schema) != ("string",): + return None + pattern = ( + _STRING_FORMATS.get(schema["format"]) + if "format" in schema + else schema["pattern"] + ) + if ( + not isinstance(pattern, str) + or _may_write(pattern, {text[0] for text in excludes if text}) + or not _regex_compiles(pattern) + ): + return None + return {"type": "regex", "pattern": pattern} + + +def _may_write(pattern: str, chars: set[str]) -> bool: + """Whether a match of `pattern` may hold one of `chars`, as any pattern may + that Python cannot parse.""" + try: + parsed = _sre_parse.parse(pattern) + except (_sre_parse.error, OverflowError, RecursionError): + return True + if parsed.state.flags & _sre_parse.SRE_FLAG_IGNORECASE: + return True + + def may(items: Any) -> bool: + for op, arg in items: + if op is _sre_parse.LITERAL: + written = {chr(arg)} + elif op is _sre_parse.IN: + written = {char for char in chars if _in_class(arg, char)} + elif op in _REPEATS: + written = chars if arg[1] and may(arg[2]) else set() + elif op is _sre_parse.SUBPATTERN: + if arg[1] & _sre_parse.SRE_FLAG_IGNORECASE: + return True + written = chars if may(arg[3]) else set() + elif op is _sre_parse.BRANCH: + written = chars if any(map(may, arg[1])) else set() + elif op is _sre_parse.AT: + continue + else: + # Such as `.`, which XGrammar lets match line breaks. + return True + if written & chars: + return True + return False + + return may(parsed) + + +def _in_class(items: list[tuple[Any, Any]], char: str) -> bool: + negated = bool(items) and items[0][0] is _sre_parse.NEGATE + matched = False + for op, arg in items[negated:]: + if op is _sre_parse.LITERAL: + matched |= ord(char) == arg + elif op is _sre_parse.RANGE: + matched |= arg[0] <= ord(char) <= arg[1] + elif op is _sre_parse.CATEGORY and arg in _CATEGORIES: + matched |= bool(re.fullmatch(_CATEGORIES[arg], char)) + else: + return True + return matched != negated + + +# Every request derives the tags for its tools anew. +@lru_cache(maxsize=512) +def _regex_compiles(pattern: str) -> bool: + from xgrammar import Grammar + + try: + Grammar.from_regex(pattern) + except RuntimeError: + return False + return True + + +def _value_format( + schema: Any, parameters: dict[str, Any], json_values: bool, closers: list[str] +) -> dict[str, Any] | None: + """Use JSON for decoded values and raw text for string values.""" + if schema is False: + return None + types = _schema_types(schema) + if json_values or (types and "string" not in types): + return _json_format(_json_value_schema(schema, parameters), closers) + return _text_format(schema, closers) + + +def _json_format(schema: Any, closers: list[str]) -> dict[str, Any] | None: + """Use JSON unless a closer is ambiguous with its syntax or the schema is + unsupported. Strings hold no closer, as the parser ends the call there.""" + if closers and ( + any(set(closer) <= _JSON_SYNTAX for closer in closers) + or not _json_compiles(json.dumps(schema)) + ): + return None + return {"type": "json_schema", "json_schema": schema, "excludes": closers} + + +# Every request derives the tags for its tools anew. +@lru_cache(maxsize=512) +def _json_compiles(schema: str) -> bool: + from xgrammar import Grammar + + try: + Grammar.from_json_schema(schema) + except RuntimeError: + return False + return True + + +def _xml_arguments_format( + content_args: dict[str, Any], + parameters: dict[str, Any], + *, + strict: bool, + closers: list[str], + sorted_arguments: bool, +) -> dict[str, Any] | None: + if set(content_args) - {"tag_pattern", "value_parser", "merge_duplicates"}: + return None + value_parser = content_args.get("value_parser") or {"name": "text"} + if value_parser.get("name") not in ("text", "json") or ( + value_parser.get("name") == "json" and value_parser.get("args") + ): + return None + key_split = _split_group(content_args["tag_pattern"], "key") + value_split = key_split and _split_group(key_split[2], "value") + if not key_split or not value_split: + return None + key_pattern = key_split[1] + begins, middles = _literals(key_split[0]), _literals(value_split[0]) + ends = _literals(value_split[2], lookahead=True) + if not begins or not middles or not ends: + return None + aheads = [end.split(_AHEAD) for end in ends if _AHEAD in end] + ends = [end for end in ends if _AHEAD not in end] + # A value may also end where the next parameter opens, if the template + # looks ahead for its only opener. + if not ends or ( + aheads and (len(begins) > 1 or {a for _, a in aheads} != {*begins}) + ): + return None + follows = [before + ahead for before, ahead in aheads] + properties = parameters.get("properties") + json_values = value_parser.get("name") == "json" + typed = properties is None or ( + isinstance(properties, dict) + and (json_values or all(map(_typed_by_parser, properties.values()))) + ) + # Other schema shapes cannot be typed by the XML parser. + if strict and (set(parameters) - _ANNOTATIONS - _OBJECT_KEYWORDS or not typed): + return None + required = parameters.get("required") or [] + names: tuple[str, ...] = () + if strict and properties is not None: + # Required entries that are no names leave the tool to the JSON schema. + if not isinstance(required, list) or not all( + isinstance(key, str) for key in required + ): + return None + names = (*properties, *(key for key in required if key not in properties)) + pieces = ( + _literal_piece(begins, before=_SEPARATOR["pattern"]), + (f"(?:{key_pattern})", _endings(names), False), + _literal_piece(middles), + # Values of every format are some text. + _TEXT_PIECE, + _literal_piece(ends), + ) + if _spells_closer(tuple(closers), pieces, repeats=True): + return None + + # Raw text has no escaping for a call closer, nor, where a value may end + # before the next opener, for any value end. + excludes = closers + follows + ends if follows else closers + text = {"type": "any_text", "excludes": excludes} if excludes else _ANY_TEXT + # Each value but the last then ends by opening the next parameter, whose + # tag starts after the opener. + starts = [""] if follows else begins + + def tag( + begin: str, value: dict[str, Any], fits: list[str], last: bool + ) -> dict[str, Any]: + if follows and not last: + fits = [e + sep + begins[0] for e in fits for sep in ("", "\n")] + follows + return _tag(begin, value, fits) + + def separated(tags: list[dict[str, Any]]) -> dict[str, Any]: + return _one_of(tags) if follows else _sequence(_SEPARATOR, _one_of(tags)) + + def arguments(body: dict[str, Any], optional: bool) -> dict[str, Any]: + """`body` after the opener of its first parameter.""" + if not follows: + return body + body = _sequence(_SEPARATOR, _const_strings(begins), body) + return _optional(body) if optional else body + + def any_parameter(last: bool) -> dict[str, Any]: + name = {"type": "regex", "pattern": key_pattern} + content = _sequence(name, _const_strings(middles), text) + return separated([tag(start, content, ends, last) for start in starts]) + + def any_parameters() -> dict[str, Any]: + more = {"type": "star", "content": any_parameter(False)} + return ( + arguments(_sequence(more, any_parameter(True)), True) if follows else more + ) + + if not strict: + return any_parameters() + if properties is None: + extra = parameters.get("additionalProperties", True) + if extra is False: + return {"type": "const_string", "value": ""} + return any_parameters() if extra is True or extra == {} else None + # A declared key repeated as an extra one would bypass its schema. + if parameters.get("additionalProperties", False) is not False: + return None + + def read_back(written: str) -> dict[str, Any] | None: + """The arguments the parser reads from `written`, unless a closer ends + the call in it.""" + if any(closer in written for closer in closers): + return None + try: + return parse_content(written, "xml-inline", content_args) + except ValueError: + return None + + def parameter( + key: str, value: dict[str, Any], choices: list[str] | None, last: bool + ) -> dict[str, Any] | None: + if not re.fullmatch(key_pattern, key) or any( + list(read_back(b + key + m + "0" + e) or ()) != [key] + for b in begins + for m in middles + for e in ends + ): + raise ValueError(f"parameter name {key!r} cannot be encoded in tool calls") + if choices is None: + return separated( + [tag(s + key + m, value, ends, last) for s in starts for m in middles] + ) + # Optional framing whitespace must not change a text choice. + tags, framed = [], set() + for b, s in zip(begins, starts): + for m in middles: + by_ends: dict[tuple[str, ...], list[str]] = {} + for choice in choices: + fits = tuple( + e + for e in ends + if read_back(b + key + m + choice + e) == {key: choice} + ) + by_ends.setdefault(fits, []).append(choice) + for fits, group in by_ends.items(): + if fits: + tags.append( + tag(s + key + m, _const_strings(group), list(fits), last) + ) + framed.update(group) + if framed != set(choices): + return None + return separated(tags) + + if not names: + return {"type": "const_string", "value": ""} + more: dict[str, dict[str, Any]] = {} + final: dict[str, dict[str, Any]] = {} + for key in names: + # A required name without a schema takes any value. + schema = properties.get(key, True) + value = _value_format(schema, parameters, json_values, closers + follows + ends) + if value is None: + return None + choices = None + if value["type"] in ("const_string", "or"): + choices = _string_choices(schema) + value = text if value == _ANY_TEXT else value + joined = parameter(key, value, choices, last=False) + ending = parameter(key, value, choices, last=True) + if joined is None or ending is None: + return None + more[key], final[key] = joined, ending + + def ordered(keys: tuple[str, ...]) -> dict[str, Any]: + before = [more[k] if k in required else _optional(more[k]) for k in keys] + if not follows: + return _sequence(*before) + + def last(start: int, stop: int) -> dict[str, Any]: + """Parameters of `keys[start:stop]` that end at one of them.""" + if stop - start == 1: + return final[keys[start]] + # Halving keeps the tag O(n log n) in size and O(log n) deep. + middle = (start + stop) // 2 + return _one_of( + [ + last(start, middle), + _sequence(*before[start:middle], last(middle, stop)), + ] + ) + + # The last parameter follows every required one. + first = max((i for i, k in enumerate(keys) if k in required), default=0) + return _sequence(*before[:first], last(first, len(keys))) + + orders = [names] + if sorted_arguments and sorted(names) != list(names): + orders.append(tuple(sorted(names))) + return arguments(_one_of([ordered(keys) for keys in orders]), not required) + + +def _arguments_format( + field: ResponseTemplateField, + parameters: dict[str, Any], + *, + strict: bool, + sorted_arguments: bool = False, +) -> dict[str, Any] | None: + closers = _closers(field) or [] + if field.content == "json" and not field.content_args: + if _spells_closer(tuple(closers), (_JSON_PIECE,)): + return None + schema = parameters if strict else {"type": "object"} + return _json_format(schema, closers) + if field.content == "xml-inline": + return _xml_arguments_format( + field.content_args, + parameters, + strict=strict, + closers=closers, + sorted_arguments=sorted_arguments, + ) + return None + + +def _names_function_by_opener(field: ResponseTemplateField) -> bool: + function = field.transform.get("function") if field.transform else None + return ( + field.open_re is not None + and "name" in field.open_re.groupindex + and not field.transform_each + and isinstance(function, dict) + and function.get("name") == "{name}" + and function.get("arguments") == "{content}" + ) + + +def _reply_format( + answer: ResponseTemplateField, + message_start: list[str] | None, + content: dict[str, Any] | None, +) -> dict[str, Any] | None: + """The `answer` field holding `content`, or free text if not given, or + `None` if it cannot be derived.""" + openers, closers = _openers(answer, message_start), _closers(answer) + if not openers or closers is None: + return None + return _sequence(_const_strings(openers), *_answer_text(closers, content)) + + +def _answer_text( + closers: list[str], content: dict[str, Any] | None = None +) -> list[dict[str, Any]]: + # Text cannot close the answer early, so no unconstrained call can follow + # it. Like a call, the answer may end at a stop token instead of a closer. + elements = [content or {"type": "any_text", "excludes": closers}] + if closers: + elements.append(_optional(_const_strings(closers))) + return elements + + +def answer_rest_structural_tag(answer: ResponseTemplateField) -> dict[str, Any] | None: + """The rest of an `answer` that is already open, or `None` if underivable.""" + closers = _closers(answer) + if closers is None: + return None + return {"type": "structural_tag", "format": _sequence(*_answer_text(closers))} + + +def tool_call_structural_tag( + field: ResponseTemplateField | None, + functions: list[dict[str, Any]], + *, + parallel_tool_calls: bool, + thinking: ResponseTemplateField | None = None, + answer: ResponseTemplateField | None = None, + answer_format: dict[str, Any] | None = None, + message_start: list[str] | None = None, + requires_reasoning: bool = False, + sorted_arguments: bool = False, + strict_arguments: bool = False, +) -> dict[str, Any] | None: + """Derive native calls to `functions` in the format of the tool call + `field`, respecting strict arguments, count and reasoning. + + `functions` are OpenAI function definitions. `strict_arguments` enforces + the schema of every function's arguments where the template can. Strict + `xml-inline` arguments come in declared order, or also alphabetical with + `sorted_arguments`. An `answer` field lets the model answer instead, in + `answer_format` if given; without `functions`, only the answer is derived. + When the model may answer, calls may also name any other function the + template can spell, with any arguments, so a harness can tell the model + the function does not exist. At the `message_start`, + the tag covers the output from its first token, which opens with non-empty + reasoning if `requires_reasoning`. Returns `None` if the template format + cannot be derived, and raises `ValueError` for names it cannot encode.""" + if not functions and answer is not None: + + def reply(start: list[str] | None) -> list[dict[str, Any]] | None: + format = _reply_format(answer, start, answer_format) + return None if format is None else [format] + + return _structural_tag(reply, thinking, message_start, requires_reasoning) + if not functions or field is None or not _names_function_by_opener(field): + return None + opener = _split_group(field.open_re.pattern, "name") + closer_literals = _closers(field) + if opener is None or closer_literals is None: + return None + arguments, enforced = [], [] + for function in functions: + # Like the generic JSON schema, a function without parameters takes none. + parameters = function.get("parameters") or {"type": "object", "properties": {}} + strict = bool(function.get("strict")) + function_arguments = None + if strict or strict_arguments: + try: + function_arguments = _arguments_format( + field, parameters, strict=True, sorted_arguments=sorted_arguments + ) + except ValueError: + if strict: + raise + enforced.append(function_arguments is not None) + if function_arguments is None and not strict: + function_arguments = _arguments_format(field, parameters, strict=False) + if function_arguments is None: + return None + arguments.append(function_arguments) + names = [] + for function in functions: + name = function["name"] + if not re.fullmatch(opener[1], name): + raise ValueError(f"tool name {name!r} cannot be encoded in tool calls") + names.append(_escape(name)) + if answer is not None: + others = _other_names( + opener[1], [f["name"] for f, e in zip(functions, enforced) if e] + ) + loose = _arguments_format(field, {}, strict=False) + if others is not None and loose is not None: + # Calls to any function whose schema is not enforced take any arguments. + names = [name for name, e in zip(names, enforced) if e] + [f"(?:{others})"] + arguments = [args for args, e in zip(arguments, enforced) if e] + [loose] + closers = _const_strings(closer_literals) if closer_literals else None + + def call(start: list[str] | None) -> dict[str, Any]: + return _one_of( + [ + _sequence( + {"type": "regex", "pattern": _call_opener(opener, name, start)}, + function_arguments, + ) + for name, function_arguments in zip(names, arguments) + ] + ) + + def output(start: list[str] | None) -> list[dict[str, Any]] | None: + elements = [call(start)] + if parallel_tool_calls and field.repeats: + repeated = _sequence(closers, call(None)) if closers else call(None) + elements.append({"type": "star", "content": repeated}) + if closers: + # EOS can terminate a call without writing its closer. + elements.append(_optional(closers)) + if answer is not None: + reply = _reply_format(answer, start, answer_format) + if reply is None: + return None + elements = [_one_of([_sequence(*elements), reply])] + return elements + + return _structural_tag(output, thinking, message_start, requires_reasoning) + + +def _structural_tag( + output: Callable[[list[str] | None], list[dict[str, Any]] | None], + thinking: ResponseTemplateField | None, + message_start: list[str] | None, + requires_reasoning: bool = False, +) -> dict[str, Any] | None: + """Frame output after reasoning, from the first token at the + `message_start`. The reasoning is optional unless `requires_reasoning` + makes the message open with non-empty reasoning.""" + rest = output(None) + if rest is None: + return None + think = None + if thinking is not None: + think = _thinking_format(thinking, message_start, requires_reasoning) + if think is None: + return None + if message_start is None: + if think is not None: + rest.insert(0, _optional(think)) + return {"type": "structural_tag", "format": _sequence(*rest)} + if think is not None and requires_reasoning: + return {"type": "structural_tag", "format": _sequence(think, *rest)} + start = output(message_start) + if start is None: + return None + format = _sequence(*start) + if think is not None: + format = _one_of([_sequence(think, *rest), format]) + return {"type": "structural_tag", "format": format} diff --git a/vllm/reasoning/__init__.py b/vllm/reasoning/__init__.py index 9e75e9b77e..796fc3ba3f 100644 --- a/vllm/reasoning/__init__.py +++ b/vllm/reasoning/__init__.py @@ -84,6 +84,10 @@ _REASONING_PARSERS_TO_REGISTER = { "response_template_reasoning_parser", "ResponseTemplateReasoningParser", ), + "beam": ( + "beam_reasoning_parser", + "BeamReasoningParser", + ), "holo2": ( "deepseek_v3_reasoning_parser", "DeepSeekV3ReasoningWithThinkingParser", diff --git a/vllm/reasoning/beam_reasoning_parser.py b/vllm/reasoning/beam_reasoning_parser.py new file mode 100644 index 0000000000..b5d87170f0 --- /dev/null +++ b/vllm/reasoning/beam_reasoning_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.beam import BeamReasoningParser + +__all__ = ["BeamReasoningParser"] diff --git a/vllm/renderers/online_renderer.py b/vllm/renderers/online_renderer.py index babd0ff6e5..65771211d0 100644 --- a/vllm/renderers/online_renderer.py +++ b/vllm/renderers/online_renderer.py @@ -60,6 +60,7 @@ from vllm.renderers.chat_utils import ( has_non_text_content, ) from vllm.renderers.inputs.preprocess import ( + extract_prompt_components, parse_model_prompt, prompt_to_seq, ) @@ -844,13 +845,17 @@ class OnlineRenderer: f"but got {type(request).__name__}" ) raise NotImplementedError(msg) - request = parser( + request_parser = parser( tokenizer, request.tools, model_config=self.model_config, chat_template_kwargs=chat_params.chat_template_kwargs, - ).adjust_request( - request=request, ) + prompt_token_ids = extract_prompt_components( + self.model_config, engine_input + ).token_ids + if prompt_token_ids is not None: + request_parser.set_prompt_token_ids(prompt_token_ids) + request = request_parser.adjust_request(request=request) return conversation, [engine_input] diff --git a/vllm/tool_parsers/__init__.py b/vllm/tool_parsers/__init__.py index 4b9cb33502..9e2e4a6d9f 100644 --- a/vllm/tool_parsers/__init__.py +++ b/vllm/tool_parsers/__init__.py @@ -90,6 +90,10 @@ _TOOL_PARSERS_TO_REGISTER = { "response_template_tool_parser", "ResponseTemplateToolParser", ), + "beam": ( + "beam_tool_parser", + "BeamToolParser", + ), "poolside_v1": ( "poolside_v1_tool_parser", "PoolsideV1ToolParser", diff --git a/vllm/tool_parsers/beam_tool_parser.py b/vllm/tool_parsers/beam_tool_parser.py new file mode 100644 index 0000000000..dbfd820688 --- /dev/null +++ b/vllm/tool_parsers/beam_tool_parser.py @@ -0,0 +1,6 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from vllm.parser.beam import BeamToolParser + +__all__ = ["BeamToolParser"] From 596e53067548fab65de6fc3abdcd020f752e7ae7 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Fri, 16 Oct 2026 15:39:46 -0700 Subject: [PATCH 23/49] [Kernel] Add router GEMM kernels for (5120, 256) and BF16 output to ll_bf16_gemm (#62042) Signed-off-by: Woosuk Kwon Co-authored-by: Lain Co-authored-by: Codex Co-authored-by: Claude Opus 5.5 Co-authored-by: vllm-agent --- csrc/fp32_router_gemm.cu | 54 ++++++++++-- csrc/fp32_router_gemm_entry.cu | 19 +++-- tests/kernels/moe/test_routing.py | 14 +++ tests/kernels/test_fp32_router_gemm.py | 61 ++++++++++--- tests/kernels/test_ll_bf16_gemm.py | 26 +++++- vllm/_custom_ops.py | 5 ++ .../linear/cute_dsl/_ll_bf16_dotprod.py | 12 +-- .../linear/cute_dsl/_ll_bf16_splitk.py | 4 +- .../linear/cute_dsl/_ll_router_common.py | 12 ++- .../kernels/linear/cute_dsl/ll_bf16.py | 85 +++++++++++++++---- .../router/bf16x3_router_gemm_cutedsl.py | 5 +- .../layers/fused_moe/router/gate_linear.py | 25 ++++-- 12 files changed, 258 insertions(+), 64 deletions(-) diff --git a/csrc/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu index 9719adb513..832874f914 100644 --- a/csrc/fp32_router_gemm.cu +++ b/csrc/fp32_router_gemm.cu @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project // -// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=32, for the +// Router GEMM: activation(T) x weight(fp32) -> fp32, M<=48, for the // supported (E, H) pairs listed at the bottom of this file. // Supports bf16 or fp32 activation; weight is always fp32. @@ -214,8 +214,26 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, // production-fidelity harness (CUDA-graph replay, per-layer cold weights). // Only enabled on the Blackwell family where it was validated; Hopper and // other shapes / fp32 activation keep the legacy geometry. - if constexpr (std::is_same_v && kNumExperts == 128 && - kHiddenDim == 6144) { + if constexpr (std::is_same_v && kNumExperts == 256 && + kHiddenDim == 5120) { + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens > 16 && kNumTokens % 8 == 0) { + constexpr int block_size = kNumTokens == 24 ? 160 : 128; + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else { + constexpr int block_size = kNumTokens == 1 ? 64 : 128; + constexpr int token_groups = kNumTokens % 2 == 0 ? 2 : 1; + launchFp32RouterGemm(output, mat_a, mat_b, + stream); + } + } else if constexpr (std::is_same_v && + kNumExperts == 128 && kHiddenDim == 6144) { // MiniMax-M3. Legacy 128/1 only fills 128 blocks and pays the same // accumulator register cliffs; B300 sweep: // even M in [6, 10] : BS=384, EPB=1, 2 token groups (1.26-1.43x) @@ -268,8 +286,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, } // --------------------------------------------------------------------------- -// Explicit instantiations: M=1..32, for both input types, for the supported -// (E, H) pairs: (256, 3072) [MiniMax-M2/M2.5] and (128, 6144) [MiniMax-M3]. +// Explicit instantiations for both input types. +// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120). // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -315,5 +333,31 @@ INSTANTIATE_ALL(__nv_bfloat16, 256, 3072) INSTANTIATE_ALL(float, 128, 6144) INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) +INSTANTIATE_ALL(float, 256, 5120) +INSTANTIATE_ALL(__nv_bfloat16, 256, 5120) + +#define INSTANTIATE_LARGE(T, E, H) \ + INSTANTIATE(T, 33, E, H) \ + INSTANTIATE(T, 34, E, H) \ + INSTANTIATE(T, 35, E, H) \ + INSTANTIATE(T, 36, E, H) \ + INSTANTIATE(T, 37, E, H) \ + INSTANTIATE(T, 38, E, H) \ + INSTANTIATE(T, 39, E, H) \ + INSTANTIATE(T, 40, E, H) \ + INSTANTIATE(T, 41, E, H) \ + INSTANTIATE(T, 42, E, H) \ + INSTANTIATE(T, 43, E, H) \ + INSTANTIATE(T, 44, E, H) \ + INSTANTIATE(T, 45, E, H) \ + INSTANTIATE(T, 46, E, H) \ + INSTANTIATE(T, 47, E, H) \ + INSTANTIATE(T, 48, E, H) + +INSTANTIATE_LARGE(float, 256, 5120) +INSTANTIATE_LARGE(__nv_bfloat16, 256, 5120) + +#undef INSTANTIATE_LARGE + #undef INSTANTIATE_ALL #undef INSTANTIATE diff --git a/csrc/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu index e55ac7e721..a6efc32831 100644 --- a/csrc/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -25,10 +25,11 @@ inline int getSMVersion() { static constexpr int FP32_MAX_TOKENS = 32; // Supported (hidden_dim, num_experts) pairs (must match the instantiations in -// fp32_router_gemm.cu): (3072, 256) for MiniMax-M2/M2.5, (6144, 128) for M3. +// fp32_router_gemm.cu). static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) { return (hidden_dim == 3072 && num_experts == 256) || - (hidden_dim == 6144 && num_experts == 128); + (hidden_dim == 6144 && num_experts == 128) || + (hidden_dim == 5120 && num_experts == 256); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -60,8 +61,7 @@ struct Fp32LoopUnroller { invokeFp32RouterGemm( output, mat_a, mat_b, stream); } else { - throw std::invalid_argument( - "fp32_router_gemm: num_tokens must be in [1, 32]"); + throw std::invalid_argument("fp32_router_gemm: unsupported num_tokens"); } } }; @@ -77,6 +77,9 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, } else if (num_experts == 128 && hidden_dim == 6144) { Fp32LoopUnroller::unroll( num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 5120) { + Fp32LoopUnroller::unroll(num_tokens, output, + mat_a, mat_b, stream); } else { throw std::invalid_argument( "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); @@ -111,9 +114,11 @@ void fp32_router_gemm( STD_TORCH_CHECK( fp32_router_gemm_supported(hidden_dim, num_experts), "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are " - "(3072, 256) and (6144, 128)"); - STD_TORCH_CHECK(num_tokens <= FP32_MAX_TOKENS, - "fp32_router_gemm: num_tokens must be in [0, 32]"); + "(3072, 256), (6144, 128), and (5120, 256)"); + const int max_tokens = + hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS; + STD_TORCH_CHECK(num_tokens <= max_tokens, + "fp32_router_gemm: num_tokens exceeds supported range"); STD_TORCH_CHECK( mat_a.scalar_type() == torch::headeronly::ScalarType::Float || mat_a.scalar_type() == torch::headeronly::ScalarType::BFloat16, diff --git a/tests/kernels/moe/test_routing.py b/tests/kernels/moe/test_routing.py index b0c3e26db9..77eac550d9 100644 --- a/tests/kernels/moe/test_routing.py +++ b/tests/kernels/moe/test_routing.py @@ -582,6 +582,20 @@ def test_fused_topk_bias( assert_routing_results_close(topk_weights, topk_ids, baseline_weights, baseline_ids) +@pytest.mark.parametrize("scoring_func", [0, 1]) +def test_fused_grouped_topk_accepts_idle_rank(scoring_func): + """Zero-token routing must leave CUDA ready for the next batch.""" + if not current_platform.is_cuda(): + pytest.skip("CUDA required") + from vllm import _custom_ops as ops + + scores = torch.empty((0, 8), device="cuda", dtype=torch.float32) + bias = torch.zeros(8, device="cuda") + weights, ids = ops.grouped_topk(scores, 2, 1, 2, True, 1.0, bias, scoring_func) + torch.testing.assert_close(weights, scores.new_empty((0, 2))) + torch.testing.assert_close(ids, scores.new_empty((0, 2), dtype=torch.int32)) + + @pytest.mark.parametrize("m,k", MK_S) @pytest.mark.parametrize("top_k", TOP_KS) @pytest.mark.parametrize( diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index d330abba53..0f55d860a6 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -3,9 +3,9 @@ """Tests for fp32_router_gemm kernel: activation×weight→fp32. Supported (hidden_size, num_experts) pairs: - (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 + (3072, 256), (6144, 128), (5120, 256) -Correctness baseline: F.linear in float32. Every M in [1, 32] is covered so +Correctness baseline: F.linear in float32. Every supported M is covered so all tuned geometries (wide-block, experts-per-block, token-group; boundaries at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy 128/1 geometry everywhere else. @@ -18,8 +18,12 @@ from vllm._custom_ops import fp32_router_gemm from vllm.platforms import current_platform # (hidden_size, num_experts) -SHAPES = [(3072, 256), (6144, 128)] -ALL_M = list(range(1, 33)) +SHAPES = [(3072, 256), (6144, 128), (5120, 256)] +CASES = [ + (m, hidden, experts) + for hidden, experts in SHAPES + for m in range(1, 49 if hidden == 5120 else 33) +] # Absolute tolerance for fp32 kernel vs float64 reference ATOL_FP32 = 2e-4 ATOL_BF16 = 2e-2 # bf16 activation has lower precision @@ -40,8 +44,7 @@ def _ref(mat_a: torch.Tensor, mat_b: torch.Tensor) -> torch.Tensor: return torch.nn.functional.linear(mat_a.float(), mat_b.float()) -@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) -@pytest.mark.parametrize("num_tokens", ALL_M) +@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES) def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int): """fp32 activation → fp32 output should match reference closely.""" _requires_sm90() @@ -58,8 +61,7 @@ def test_fp32_activation(num_tokens: int, hidden_dim: int, num_experts: int): torch.testing.assert_close(out, ref, atol=ATOL_FP32, rtol=0) -@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) -@pytest.mark.parametrize("num_tokens", ALL_M) +@pytest.mark.parametrize("num_tokens,hidden_dim,num_experts", CASES) def test_bf16_activation(num_tokens: int, hidden_dim: int, num_experts: int): """bf16 activation → fp32 output should match reference within bf16 error.""" _requires_sm90() @@ -136,17 +138,19 @@ def test_zero_tokens_returns_empty(): assert out.dtype == torch.float32 -def test_rejects_invalid_inputs(): +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +def test_rejects_invalid_inputs(hidden_dim, num_experts): """The entry must fail loudly, never compute silently wrong results.""" _requires_sm90() device = torch.device("cuda") - hidden_dim, num_experts = SHAPES[0] mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device) - # num_tokens > 32 (beyond the instantiated range) + # Beyond the instantiated range for this shape. + invalid_tokens = 49 if hidden_dim == 5120 else 33 with pytest.raises(Exception, match="num_tokens"): fp32_router_gemm( - torch.randn(33, hidden_dim, dtype=torch.float32, device=device), mat_b + torch.randn(invalid_tokens, hidden_dim, dtype=torch.float32, device=device), + mat_b, ) # unsupported (hidden_dim, num_experts) pair @@ -173,3 +177,36 @@ def test_rejects_invalid_inputs(): fp32_router_gemm( torch.randn(4, hidden_dim, dtype=torch.float16, device=device), mat_b ) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA required") +@pytest.mark.parametrize( + "tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128] +) +@torch.inference_mode() +def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens): + """Low-latency routing must retain FP32 weight precision and logits.""" + from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear + from vllm.platforms import current_platform + + if not current_platform.is_device_capability_family(100): + pytest.skip("SM100 required") + torch.manual_seed(42) + router = GateLinear( + 5120, 256, params_dtype=torch.float32, out_dtype=torch.float32 + ).cuda() + router.weight.normal_(std=0.01) + x = torch.randn(tokens, 5120, device="cuda", dtype=torch.bfloat16) + expected = (x.double() @ router.weight.double().T).float() + actual, bias = router(x) + assert bias is None and actual.dtype == torch.float32 + torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5) + torch.testing.assert_close(actual.topk(8).indices, expected.topk(8).indices) + if tokens in (1, 16, 48, 56): + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured, _ = router(x) + x.normal_() + graph.replay() + expected = (x.double() @ router.weight.double().T).float() + torch.testing.assert_close(captured, expected, rtol=2e-5, atol=2e-5) diff --git a/tests/kernels/test_ll_bf16_gemm.py b/tests/kernels/test_ll_bf16_gemm.py index 1c7a15d514..76ea03f60e 100644 --- a/tests/kernels/test_ll_bf16_gemm.py +++ b/tests/kernels/test_ll_bf16_gemm.py @@ -566,7 +566,7 @@ def test_invalid_output_dtype(): with pytest.raises(ValueError, match="output_dtype=torch.float32"): from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm - ll_bf16_gemm(a, b, output_dtype=torch.bfloat16) + ll_bf16_gemm(a, b, output_dtype=torch.float16) def test_cache_miss_compiles_dotprod(): @@ -581,5 +581,29 @@ def test_cache_miss_compiles_dotprod(): _assert_close(out, _ref(a, b), context="cache miss dotprod") +@pytest.mark.parametrize("M", [1, 2, 8, 16]) +@pytest.mark.parametrize("N", [1792, 3584]) +@torch.inference_mode() +def test_bf16_output_and_fp32_cache_coexist(M, N): + """BF16 epilogues retain FP32 accumulation and can replay with new inputs.""" + from vllm.model_executor.kernels.linear.cute_dsl.ll_bf16 import ll_bf16_gemm + + torch.manual_seed(42) + x = torch.randn(M, 5120, device="cuda", dtype=torch.bfloat16) + w = torch.randn(N, 5120, device="cuda", dtype=torch.bfloat16) * 0.01 + for dtype in (torch.float32, torch.bfloat16): + actual = ll_bf16_gemm(x, w, output_dtype=dtype) + assert actual.dtype == dtype + reference = (x.double() @ w.double().T).to(dtype) + torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = ll_bf16_gemm(x, w, output_dtype=torch.bfloat16) + x.normal_() + graph.replay() + reference = (x.double() @ w.double().T).bfloat16() + torch.testing.assert_close(actual, reference, rtol=0.008, atol=1e-5) + + if __name__ == "__main__": pytest.main([__file__, "-v"]) diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 70fed78573..d74bb67816 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -2424,6 +2424,11 @@ def grouped_topk( raise NotImplementedError( "The fused grouped_topk kernel is only available on CUDA and XPU platforms" ) + if scores.shape[0] == 0: + return ( + scores.new_empty((0, topk), dtype=torch.float32), + scores.new_empty((0, topk), dtype=torch.int32), + ) return torch.ops._moe_C.grouped_topk( scores, num_expert_group, diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py index f06de6bd0c..707e2cdf45 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_dotprod.py @@ -356,10 +356,12 @@ class LLBf16Dotprod: if tidx == 0: for m in cutlass.range_constexpr(M): partials = sm[m, None].load() - gC[m, n_idx] = partials.reduce( - cute.ReductionOp.ADD, - init_val=cutlass.Float32(0.0), - reduction_profile=0, - ) + gC[m, n_idx] = cutlass.Float32( + partials.reduce( + cute.ReductionOp.ADD, + init_val=cutlass.Float32(0.0), + reduction_profile=0, + ) + ).to(gC.element_type) if const_expr(self.use_pdl): cute.arch.griddepcontrol_launch_dependents() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py index b44a5f6e2c..2ae3d59922 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_bf16_splitk.py @@ -47,7 +47,7 @@ class LLBf16SplitK: :param ab_dtype: Element type for A and B operands. :param acc_dtype: Accumulator type used by MMA and reductions. - :param out_dtype: Output element type. The public wrapper uses fp32. + :param out_dtype: Output element type (fp32 or bf16). :param tile_n: CTA tile size in N. M is fixed to 16 for router batches. :type tile_n: int :param tile_k: K tile size staged through shared memory. @@ -580,6 +580,6 @@ class LLBf16SplitK: reduction_profile=0, ) ) - gC[local_coord] = acc + gC[local_coord] = cutlass.Float32(acc).to(gC.element_type) cute.arch.sync_threads() diff --git a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py index 23fcefdac9..7770c0f8de 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/_ll_router_common.py @@ -76,8 +76,8 @@ def make_fake_gemm_tensors( a_dtype: torch.dtype, b_dtype: torch.dtype, divisibility: int, + output_dtype: torch.dtype = torch.float32, ): - from cutlass import Float32 from quack.compile_utils import make_fake_tensor hidden_states = make_fake_tensor( @@ -86,7 +86,7 @@ def make_fake_gemm_tensors( router_weight = make_fake_tensor( cutlass_dtype(b_dtype), (N, K), divisibility=divisibility ) - output = make_fake_tensor(Float32, (M, N), divisibility=1) + output = make_fake_tensor(cutlass_dtype(output_dtype), (M, N), divisibility=1) return hidden_states, router_weight, output @@ -97,6 +97,7 @@ def validate_common_gemm_inputs( *, op_name: str, k_multiple: int | None = None, + output_dtypes: tuple[torch.dtype, ...] = (torch.float32,), ) -> None: if hidden_states.dim() != 2 or router_weight.dim() != 2: raise ValueError("hidden_states and router_weight must be 2D tensors") @@ -106,8 +107,11 @@ def validate_common_gemm_inputs( raise ValueError( "hidden_states and router_weight must be on the same CUDA device" ) - if output_dtype != torch.float32: - raise ValueError(f"{op_name} only supports output_dtype=torch.float32") + if output_dtype not in output_dtypes: + raise ValueError( + f"{op_name} only supports output_dtype=" + + " or ".join(str(dtype) for dtype in output_dtypes) + ) if hidden_states.shape[1] != router_weight.shape[1]: raise ValueError( "hidden_states and router_weight must have matching K dimensions" diff --git a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py index 95cb78c6a1..9d1573bb80 100644 --- a/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py +++ b/vllm/model_executor/kernels/linear/cute_dsl/ll_bf16.py @@ -47,6 +47,8 @@ _SM100F_TUNED_SPLITK_CONFIGS: dict[tuple[int, int], dict[int, tuple[int, int]]] **{M: (8, 5) for M in (5, 8)}, 9: (8, 2), }, + (5120, 1792): {M: (2, 4) for M in range(5, 17)}, + (5120, 3584): {M: (2, 4) for M in range(5, 17)}, (7168, 256): {14: (8, 2)}, (6144, 256): {M: (8, 2) for M in (9, 12, 16)}, (7168, 384): {M: (7, 5) for M in (13, 16)}, @@ -101,39 +103,52 @@ class LLBf16Gemm: bs: int = 0 split_k: int = 0 num_stages: int = 0 + output_dtype: torch.dtype = torch.float32 def __init__(self, *, prefetch_pdl_weights: bool = False) -> None: self._prefetch_pdl_weights = prefetch_pdl_weights - # Dot-prod: keyed on (M, K, bs), because M and K are Constexpr. - self._compiled_cache: dict[tuple[int, int, int], Any] = {} - # Split-K: keyed on (split_k, num_stages), fully shape-dynamic. - self._splitk_cache: dict[tuple[int, int], Any] = {} + # M and K are compile-time constants for dot-product. + self._compiled_cache: dict[tuple[int, int, int, torch.dtype], Any] = {} + # Split-K is shape-dynamic; both caches distinguish output dtype. + self._splitk_cache: dict[tuple[int, int, torch.dtype], Any] = {} - def dispatch(self, *, M: int, K: int, N: int) -> CompileKey: + def dispatch( + self, *, M: int, K: int, N: int, output_dtype: torch.dtype = torch.float32 + ) -> CompileKey: tuned_bs, tuned_splitk = _arch_tuned_configs() if M <= _DEFAULT_DOTPROD_MAX_M or K < 2048: bs = tuned_bs.get((K, N), {}).get(M, _DEFAULT_DOTPROD_BS) - return self.CompileKey(backend="dotprod", M=M, K=K, bs=bs) + return self.CompileKey( + backend="dotprod", M=M, K=K, bs=bs, output_dtype=output_dtype + ) split_k, num_stages = tuned_splitk.get((K, N), {}).get( M, _DEFAULT_SPLITK_CONFIG ) - return self.CompileKey(backend="splitk", split_k=split_k, num_stages=num_stages) + return self.CompileKey( + backend="splitk", + split_k=split_k, + num_stages=num_stages, + output_dtype=output_dtype, + ) def get_warmup_keys( self, *, shapes: Iterable[tuple[int, int]], m_values: Iterable[int], + output_dtype: torch.dtype = torch.float32, ) -> list[CompileKey]: return list( dict.fromkeys( - self.dispatch(M=M, K=K, N=N) for K, N in shapes for M in m_values + self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype) + for K, N in shapes + for M in m_values ) ) @staticmethod - def _fake_gemm_tensors(*, M, K, N, divisibility: int): + def _fake_gemm_tensors(*, M, K, N, divisibility: int, output_dtype: torch.dtype): return make_fake_gemm_tensors( M=M, K=K, @@ -141,10 +156,13 @@ class LLBf16Gemm: a_dtype=torch.bfloat16, b_dtype=torch.bfloat16, divisibility=divisibility, + output_dtype=output_dtype, ) def _compile_splitk(self, compile_key: CompileKey) -> None: cute, _ = cute_context() + from cutlass import BFloat16, Float32 + from ._ll_bf16_splitk import LLBf16SplitK hidden_states, router_weight, output = self._fake_gemm_tensors( @@ -152,8 +170,12 @@ class LLBf16Gemm: K=cute.sym_int(), N=cute.sym_int(), divisibility=8, + output_dtype=compile_key.output_dtype, ) gemm = LLBf16SplitK( + out_dtype=BFloat16 + if compile_key.output_dtype == torch.bfloat16 + else Float32, tile_n=16, tile_k=256, num_stages=compile_key.num_stages, @@ -169,7 +191,9 @@ class LLBf16Gemm: current_cuda_stream(), options="--enable-tvm-ffi", ) - self._splitk_cache[(compile_key.split_k, compile_key.num_stages)] = compiled + self._splitk_cache[ + (compile_key.split_k, compile_key.num_stages, compile_key.output_dtype) + ] = compiled logger.debug( "Compiled ll_bf16_splitk: sk=%d ns=%d", compile_key.split_k, @@ -187,6 +211,7 @@ class LLBf16Gemm: K=compile_key.K, N=N, divisibility=stride_divisibility, + output_dtype=compile_key.output_dtype, ) gemm = LLBf16Dotprod( k=compile_key.K, @@ -205,7 +230,9 @@ class LLBf16Gemm: current_cuda_stream(), options="--enable-tvm-ffi --ptxas-options -maxrregcount=64", ) - self._compiled_cache[(compile_key.M, compile_key.K, compile_key.bs)] = compiled + self._compiled_cache[ + (compile_key.M, compile_key.K, compile_key.bs, compile_key.output_dtype) + ] = compiled logger.debug( "Compiled ll_bf16_dotprod: M=%d, K=%d, bs=%d", compile_key.M, @@ -215,12 +242,21 @@ class LLBf16Gemm: def compile(self, compile_key: CompileKey) -> None: if compile_key.backend == "splitk": - splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + splitk_cache_key = ( + compile_key.split_k, + compile_key.num_stages, + compile_key.output_dtype, + ) if splitk_cache_key not in self._splitk_cache: self._compile_splitk(compile_key) return - dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + dotprod_cache_key = ( + compile_key.M, + compile_key.K, + compile_key.bs, + compile_key.output_dtype, + ) if dotprod_cache_key not in self._compiled_cache: self._compile_dotprod(compile_key) @@ -229,8 +265,11 @@ class LLBf16Gemm: *, shapes: Iterable[tuple[int, int]], m_values: Iterable[int], + output_dtype: torch.dtype = torch.float32, ) -> None: - for compile_key in self.get_warmup_keys(shapes=shapes, m_values=m_values): + for compile_key in self.get_warmup_keys( + shapes=shapes, m_values=m_values, output_dtype=output_dtype + ): self.compile(compile_key) @staticmethod @@ -251,6 +290,7 @@ class LLBf16Gemm: output_dtype, op_name="ll_bf16_gemm", k_multiple=8, + output_dtypes=(torch.float32, torch.bfloat16), ) def __call__( @@ -258,19 +298,28 @@ class LLBf16Gemm: hidden_states: torch.Tensor, # [M, K] bf16 router_weight: torch.Tensor, # [N, K] bf16 output_dtype: torch.dtype = torch.float32, - ) -> torch.Tensor: # [M, N] fp32 + ) -> torch.Tensor: # [M, N] self._validate_inputs(hidden_states, router_weight, output_dtype) M, K = hidden_states.shape N = router_weight.shape[0] - compile_key = self.dispatch(M=M, K=K, N=N) + compile_key = self.dispatch(M=M, K=K, N=N, output_dtype=output_dtype) if compile_key.backend == "splitk": - splitk_cache_key = (compile_key.split_k, compile_key.num_stages) + splitk_cache_key = ( + compile_key.split_k, + compile_key.num_stages, + compile_key.output_dtype, + ) if splitk_cache_key not in self._splitk_cache: self.compile(compile_key) kernel = self._splitk_cache[splitk_cache_key] else: - dotprod_cache_key = (compile_key.M, compile_key.K, compile_key.bs) + dotprod_cache_key = ( + compile_key.M, + compile_key.K, + compile_key.bs, + compile_key.output_dtype, + ) if dotprod_cache_key not in self._compiled_cache: self.compile(compile_key) kernel = self._compiled_cache[dotprod_cache_key] diff --git a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py index 2190eea157..395d7d68ea 100644 --- a/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py +++ b/vllm/model_executor/layers/fused_moe/router/bf16x3_router_gemm_cutedsl.py @@ -517,6 +517,7 @@ _TILE_CONFIG_OVERRIDES = { 256: (128, 32), 512: (128, 19), }, + (5120, 256): {64: (64, 64)}, # GB200 (6144, 128): { # MiniMax-M3 64: (32, 64), 128: (64, 64), @@ -533,7 +534,7 @@ def _pick_tile_config(N: int, K: int, M: int, num_sms: int) -> tuple[int, int]: if 32 < N <= 512: token_bucket = triton.next_power_of_2(N) overrides = _TILE_CONFIG_OVERRIDES.get((K, M)) - if overrides is not None: + if overrides is not None and token_bucket in overrides: BN, split_k = overrides[token_bucket] return BN, min(split_k, k_tiles) @@ -609,6 +610,8 @@ def _bf16x3_router_gemm(X: torch.Tensor, W: torch.Tensor) -> torch.Tensor: """Return ``X @ W.T`` using the SM100 BF16x3 router GEMM kernel.""" N, K = X.shape M, _ = W.shape + if N == 0: + return X.new_empty((N, M), dtype=torch.float32) num_sms = torch.cuda.get_device_properties(X.device).multi_processor_count BN, split_k = _pick_tile_config(N, K, M, num_sms) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index c6b1da890b..7d14f105c9 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -24,8 +24,8 @@ class GateLinear(ReplicatedLinear): 1. cuteDSL ll_bf16_gemm (SM90+, M<=16, bf16 in, fp32 out, any shape w/ K divisible by 8) - 2. fp32 specialized kernel (bf16/fp32 in, fp32 out, M<=32; SM90+ - model-specific shapes, gfx950 any shape) + 2. fp32 specialized kernel (bf16/fp32 in, fp32 out; SM90+ + model-specific shapes and batch limits, gfx950 any shape with M<=32) 3. cuteDSL ll_fp32w_gemm (SM90+, M<=32, bf16/fp32 in, fp32 out, any shape) 4. bf16x3 CuteDSL kernel (SM100, bf16 in, fp32 weight) @@ -46,9 +46,8 @@ class GateLinear(ReplicatedLinear): split-K from the token count and skip the batch-invariant linear path. """ - # (hidden_size, num_experts) pairs with an instantiated fp32 kernel: - # (3072, 256) -> MiniMax-M2/M2.5, (6144, 128) -> MiniMax-M3 - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128)} + # (hidden_size, num_experts) pairs with an instantiated fp32 kernel. + FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 256)} FP32_MAX_TOKENS = 32 # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses # whose router shape stays ahead further up the range may raise it. @@ -454,9 +453,9 @@ def fp32_router_gemm_dispatch_impl( weight: torch.Tensor, allow_bf16x3_router_gemm: bool, ) -> torch.Tensor: - """Dynamically run fp32 specialized gemm if num_tokens <= FP32_MAX_TOKENS, - otherwise optionally run the experimental BF16x3 kernel for medium/large - SM100 router batches, then fall back to F.linear. + """Run the FP32 kernel for small batches, with shape-specific limits. + + Larger SM100 BF16 batches use BF16x3; other inputs fall back to F.linear. This must be wrapped in a custom op because our torch.compile integration does not support runtime dispatching on num_tokens. """ @@ -473,7 +472,15 @@ def fp32_router_gemm_dispatch_impl( if weight.dtype == torch.bfloat16: return torch.mm(x, weight.T, out_dtype=torch.float32) return torch.nn.functional.linear(x.float(), weight) - if x.shape[0] <= _FP32_ROUTER_GEMM_MAX_TOKENS: + max_tokens = _FP32_ROUTER_GEMM_MAX_TOKENS + if ( + allow_bf16x3_router_gemm + and x.dtype == torch.bfloat16 + and weight.shape == (256, 5120) + ): + # GB200: four token groups win through 48 for multiples of eight. + max_tokens = 48 if x.shape[0] % 8 == 0 else 16 + if x.shape[0] <= max_tokens: return ops.fp32_router_gemm(x, weight) if allow_bf16x3_router_gemm and x.dtype == torch.bfloat16: From b088eba94431b711ffaa2f356e33361cf38b8fc5 Mon Sep 17 00:00:00 2001 From: Afonso Ribeiro Date: Fri, 16 Oct 2026 23:46:02 +0100 Subject: [PATCH 24/49] [Perf][Parser] Stream MiniMax M2 tool argument values incrementally (16 KB call: 5.9 s to 26 ms) (#62035) Signed-off-by: Afonso Ribeiro Co-authored-by: Flora Feng <4florafeng@gmail.com> --- tests/parser/engine/test_minimax_m2.py | 121 +++++++++++++++++++++++++ vllm/parser/minimax_m2.py | 78 +++++++++++++--- 2 files changed, 185 insertions(+), 14 deletions(-) diff --git a/tests/parser/engine/test_minimax_m2.py b/tests/parser/engine/test_minimax_m2.py index 9b2574dec4..4a3f5cbdf4 100644 --- a/tests/parser/engine/test_minimax_m2.py +++ b/tests/parser/engine/test_minimax_m2.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import dataclasses import json import pytest @@ -15,11 +16,14 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionToolsParam, FunctionDefinition, ) +from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.minimax_m2 import ( THINK_END, TOOL_CALL_END, TOOL_CALL_START, MinimaxM2Parser, + _minimax_m2_arg_value_extender, + minimax_m2_config, ) @@ -238,6 +242,123 @@ class TestStreaming: assert collect_tool_arguments(results) == "" +WRITE_CALL_HEAD = ( + '\n\n' + '/repo/app.py\n' + '644\n' + '' +) +WRITE_CALL_TAIL = "\n\n" +# Escapes, non-ASCII, blank lines and tag-like text that must stay inside +# the value, including prefixes of the closing and next-parameter tags. +CONTENT_LINE = 'if a < b > c: print("x\\y") #

int: + calls = 0 + + def counting_converter(raw_args: str, partial: bool) -> str: + nonlocal calls + calls += partial + return converter(raw_args, partial) + + parser = ParserEngine( + mock_tokenizer, + parser_engine_config=dataclasses.replace( + base, arg_converter=counting_converter + ), + ) + text = WRITE_CALL_HEAD + CONTENT_LINE * lines + WRITE_CALL_TAIL + simulate_tool_streaming( + parser, + mock_request, + [text[i : i + 4] for i in range(0, len(text), 4)], + ) + return calls + + assert streaming_conversions(lines=20) == streaming_conversions(lines=200) + + def test_extender_stops_at_split_end_tag(self): + extend = _minimax_m2_arg_value_extender('') + assert extend is not None + deltas = ['a"\n', ""] + assert [extend(d) for d in deltas] == ['a\\"\\n', "x") + assert extend is not None + assert [extend(d) for d in ["< parameter", " name ="]] == ["< parameter", None] + + @pytest.mark.parametrize( + "raw_args", + [ + 'x', + '12' + '3' + ), + 'x', + ], + ids=[ + "closed_value", + "open_name", + "repeated_key_not_last", + "earlier_unmatched_opener", + "bare_name_opens_quote", + ], + ) + def test_no_extender_unless_open_value_is_stable(self, raw_args): + assert _minimax_m2_arg_value_extender(raw_args) is None + + class TestReasoning: def test_extract_reasoning_without_start_token(self, parser, mock_request): reasoning, content = parser.extract_reasoning( diff --git a/vllm/parser/minimax_m2.py b/vllm/parser/minimax_m2.py index c96f852b9d..d804148ec7 100644 --- a/vllm/parser/minimax_m2.py +++ b/vllm/parser/minimax_m2.py @@ -22,6 +22,7 @@ import regex as re from vllm.parser.engine.events import EventType from vllm.parser.engine.parser_engine import ParserEngine from vllm.parser.engine.parser_engine_config import ( + ArgValueExtender, ParserEngineConfig, ParserState, Transition, @@ -57,37 +58,85 @@ _PARTIAL_PARAM_RE = re.compile( ) -def _minimax_m2_arg_converter(raw_args: str, partial: bool) -> str: +# What ends an open value in ``_PARAM_RE``: a closing tag or the next opener. +_PARAM_VALUE_END_RE = re.compile(r"<\s*/\s*parameter\s*>|<\s*parameter\s+name\s*=") + + +def _param_name(match: re.Match) -> str: + return ( + match.group("dq_name") + or match.group("sq_name") + or match.group("bare_name") + or "" + ).strip() + + +def _minimax_m2_parse_args( + raw_args: str, partial: bool +) -> tuple[dict[str, object], re.Match | None]: + """Return the arguments and the match of the open trailing value, if any.""" params: dict[str, object] = {} for match in _PARAM_RE.finditer(raw_args): - name = ( - match.group("dq_name") - or match.group("sq_name") - or match.group("bare_name") - or "" - ).strip() + name = _param_name(match) if not name: continue # Keep the value verbatim (like the glm47 converter): it is inline # between `>` and ``, so surrounding whitespace is data. params[name] = match.group("value") + open_match = None if partial: remaining = _PARAM_RE.sub("", raw_args) match = _PARTIAL_PARAM_RE.search(remaining) if match: - name = ( - match.group("dq_name") - or match.group("sq_name") - or match.group("bare_name") - or "" - ).strip() + name = _param_name(match) if name: # Verbatim, same as the complete-match loop above. params[name] = match.group("value") + open_match = match - return json.dumps(params, ensure_ascii=False) + return params, open_match + + +def _minimax_m2_arg_converter(raw_args: str, partial: bool) -> str: + return json.dumps(_minimax_m2_parse_args(raw_args, partial)[0], ensure_ascii=False) + + +def _value_end_window(text: str) -> str: + """Return the suffix of *text* that may still grow into a value end.""" + start = text.rfind("<") + if start == -1: + return "" + tail = text[start:] + return tail if _PARAM_VALUE_END_RE.fullmatch(tail, partial=True) else "" + + +def _minimax_m2_arg_value_extender(raw_args: str) -> ArgValueExtender | None: + # The open value runs to the end of raw_args and is the last JSON value, + # so a delta appends its JSON-escaped text unless it closes the value. + params, match = _minimax_m2_parse_args(raw_args, partial=True) + if match is None or next(reversed(params)) != _param_name(match): + return None + # Later text could still re-parse the opener: an earlier unmatched one, or + # a bare name that starts a quote, may close into a different name. + bare_name = match.group("bare_name") + if _PARAM_VALUE_END_RE.search(match.string, 0, match.start()) or ( + bare_name is not None and bare_name[0] in "\"'" + ): + return None + + window = _value_end_window(match.group("value")) + + def extend(delta: str) -> str | None: + nonlocal window + text = window + delta + if _PARAM_VALUE_END_RE.search(text): + return None + window = _value_end_window(text) + return json.dumps(delta, ensure_ascii=False)[1:-1] + + return extend @functools.cache @@ -191,6 +240,7 @@ def minimax_m2_config() -> ParserEngineConfig: }, }, arg_converter=_minimax_m2_arg_converter, + arg_value_extender=_minimax_m2_arg_value_extender, stream_arg_deltas=True, tool_args_json=False, validate_tool_names=True, From 28e92eeff25f96c1b78a2f7552f6747c9073798e Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Sat, 17 Oct 2026 01:58:56 +0300 Subject: [PATCH 25/49] [ROCm][Perf][Kimi-K3] Use AITER Gluon chunk_kda for KDA prefill (#59552) Signed-off-by: Omar Muhammad Signed-off-by: Nicholas Susanto Co-authored-by: Nicholas Susanto Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: Kevin H. Luu --- .../models/kimi_k3/test_amd_kda_checkpoint.py | 5 +- tests/models/kimi_k3/test_amd_kda_chunk.py | 25 ++++++- vllm/_aiter_ops.py | 65 +++++++++++++++++++ .../warmup/kimi_k3_triton_warmup.py | 56 +++++++++++++++- vllm/models/kimi_k3/amd/kda.py | 22 ++++--- vllm/models/kimi_k3/amd/ops/kda_prefill.py | 44 +++++++++++-- 6 files changed, 197 insertions(+), 20 deletions(-) diff --git a/tests/models/kimi_k3/test_amd_kda_checkpoint.py b/tests/models/kimi_k3/test_amd_kda_checkpoint.py index 6f3129a224..2f8bd3f0e4 100644 --- a/tests/models/kimi_k3/test_amd_kda_checkpoint.py +++ b/tests/models/kimi_k3/test_amd_kda_checkpoint.py @@ -482,7 +482,7 @@ def _spec_for(monkeypatch, **attrs) -> MambaSpec: layer = KimiK3DeltaAttention.__new__(KimiK3DeltaAttention) for name, value in dict( use_prefill_checkpoint=True, - use_fused_chunk=True, + kda_prefill_backend="fused", num_spec=0, use_safe_gate=True, ).items(): @@ -506,7 +506,8 @@ def test_opt_in_publishes_chunk_walk_alignment(monkeypatch) -> None: "attr,value", [ ("use_prefill_checkpoint", False), - ("use_fused_chunk", False), + ("kda_prefill_backend", "triton"), + ("kda_prefill_backend", "gluon"), ("use_safe_gate", False), ], ) diff --git a/tests/models/kimi_k3/test_amd_kda_chunk.py b/tests/models/kimi_k3/test_amd_kda_chunk.py index c53e9bacdb..2dcf886c91 100644 --- a/tests/models/kimi_k3/test_amd_kda_chunk.py +++ b/tests/models/kimi_k3/test_amd_kda_chunk.py @@ -77,7 +77,7 @@ def _run(inp: dict, use_fused: bool, out: torch.Tensor | None = None, **kw): output_final_state=True, use_qk_l2norm_in_kernel=True, cu_seqlens=inp["cu"], - use_fused_chunk=use_fused, + backend=kw.pop("backend", "fused" if use_fused else "triton"), out=out, **kw, ) @@ -112,6 +112,27 @@ def test_fused_chunk_matches_triton(seqlens: list[int]) -> None: torch.testing.assert_close(state_fused, state_ref, rtol=1e-3, atol=1e-3) +@pytest.mark.parametrize( + "seqlens", + [[1], [64], [1024], [4096], [513, 64, 1, 1200], [512] * 6], +) +def test_gluon_chunk_matches_triton(seqlens: list[int]) -> None: + """AITER's Gluon chunk_kda must match the Triton chunk path too.""" + from vllm._aiter_ops import rocm_aiter_ops + + if not rocm_aiter_ops.is_kda_gluon_prefill_available(): + pytest.skip("AITER Gluon chunk_kda is not available here") + + inp = _inputs(seqlens, seed=len(seqlens)) + o_ref, state_ref = _run(inp, use_fused=False) + o_gluon, state_gluon = _run(inp, use_fused=False, backend="gluon") + + torch.testing.assert_close(o_gluon.float(), o_ref.float(), rtol=2e-2, atol=2e-2) + torch.testing.assert_close(state_gluon, state_ref, rtol=5e-3, atol=5e-3) + state_rel = (state_gluon - state_ref).norm() / state_ref.norm() + assert state_rel < 5e-3, f"final state relative L2 error {state_rel:.2e}" + + @pytest.mark.parametrize("use_fused", [False, True]) def test_out_argument_receives_the_result(use_fused: bool) -> None: """``out`` is a promise about where the result lands. @@ -648,7 +669,7 @@ def _run_paged(inp: dict, cache, idx, warm, use_fused: bool, **kw): has_initial_state=warm, use_qk_l2norm_in_kernel=True, cu_seqlens=inp["cu"], - use_fused_chunk=use_fused, + backend="fused" if use_fused else "triton", **kw, ) diff --git a/vllm/_aiter_ops.py b/vllm/_aiter_ops.py index 6a6c99840a..29364f204a 100644 --- a/vllm/_aiter_ops.py +++ b/vllm/_aiter_ops.py @@ -2606,6 +2606,27 @@ class rocm_aiter_ops: return cls.is_linear_enabled() and on_gfx950() + @staticmethod + @functools.cache + def _kda_gluon_prefill_importable() -> bool: + try: + from aiter.ops.triton.attention.chunk_kda import chunk_kda # noqa: F401 + + return True + except ImportError: + return False + + @classmethod + @if_aiter_supported + def is_kda_gluon_prefill_available(cls) -> bool: + from vllm.platforms.rocm import get_cdna_version + + return ( + cls._AITER_ENABLED + and get_cdna_version() >= 4 + and cls._kda_gluon_prefill_importable() + ) + @classmethod def get_aiter_allreduce(cls): """Return the TP device communicator's AITER custom-allreduce if it has @@ -5111,6 +5132,50 @@ class rocm_aiter_ops: layer_input.view(*outer_shape, hidden_size), ) + @staticmethod + def kda_chunk_prefill( + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + g: torch.Tensor, + beta: torch.Tensor, + A_log: torch.Tensor, + dt_bias: torch.Tensor | None, + lower_bound: float | None, + cu_seqlens: torch.Tensor | None, + chunk_indices: torch.Tensor | None, + chunk_offsets: torch.Tensor | None, + scale: float, + out: torch.Tensor | None, + initial_state: torch.Tensor | None, + output_final_state: bool, + state_cache: torch.Tensor | None, + state_indices: torch.Tensor | None, + has_initial_state: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + from aiter.ops.triton.attention.chunk_kda import chunk_kda + + return chunk_kda( + q=q, + k=k, + v=v, + g=g, + beta=beta, + A_log=A_log, + dt_bias=dt_bias, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + out=out, + initial_state=initial_state, + output_final_state=output_final_state, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) + rocm_aiter_ops.register_ops_once() _sync_aiter_situv2_moe_env() diff --git a/vllm/model_executor/warmup/kimi_k3_triton_warmup.py b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py index e7883fff44..68186ac085 100644 --- a/vllm/model_executor/warmup/kimi_k3_triton_warmup.py +++ b/vllm/model_executor/warmup/kimi_k3_triton_warmup.py @@ -19,9 +19,11 @@ if TYPE_CHECKING: logger = init_logger(__name__) -def _get_kda_layer(worker: Worker) -> KimiK3DeltaAttention | None: +def _get_kda_layer( + worker: Worker, module: str = "vllm.models.kimi_k3.nvidia.kda" +) -> KimiK3DeltaAttention | None: # Kimi model construction already imports kda. Avoid importing it here. - kda = sys.modules.get("vllm.models.kimi_k3.nvidia.kda") + kda = sys.modules.get(module) if kda is None: return None @@ -172,9 +174,59 @@ def _warm_recurrent_kda( ) +def _warm_gluon_chunk_kda(layer, input_dtype: torch.dtype) -> None: + from vllm.models.kimi_k3.amd.ops.kda_prefill import ( + chunk_kda_prefill, + get_gluon_chunk_kda_warmup_profiles, + ) + + if layer.kda_prefill_backend != "gluon": + return + kv_cache = layer.kv_cache + if not isinstance(kv_cache, (list, tuple)) or len(kv_cache) < 2: + return + state = kv_cache[1] + if not isinstance(state, torch.Tensor) or not state.numel(): + return + + logger.info("Warming up Kimi-K3 Gluon chunk KDA prefill kernels.") + h, d = int(layer.local_num_heads), int(layer.head_dim) + for num_seqs, seq_len in get_gluon_chunk_kda_warmup_profiles(): + num_tokens = num_seqs * seq_len + shape = (1, num_tokens, h, d) + chunk_kda_prefill( + q=torch.zeros(shape, dtype=input_dtype, device=state.device), + k=torch.zeros(shape, dtype=input_dtype, device=state.device), + v=torch.zeros(shape, dtype=input_dtype, device=state.device), + raw_g=torch.zeros(shape, dtype=input_dtype, device=state.device), + raw_beta=torch.zeros( + (1, num_tokens, h), dtype=input_dtype, device=state.device + ), + A_log=layer.A_log, + g_bias=layer.dt_bias, + lower_bound=layer.gate_lower_bound, + state_cache=state, + state_indices=torch.zeros(num_seqs, dtype=torch.int32, device=state.device), + has_initial_state=torch.zeros( + num_seqs, dtype=torch.bool, device=state.device + ), + use_qk_l2norm_in_kernel=True, + cu_seqlens=torch.arange( + 0, num_tokens + 1, seq_len, dtype=torch.int32, device=state.device + ), + backend="gluon", + out=torch.empty(shape, dtype=input_dtype, device=state.device), + ) + + @torch.inference_mode() def kimi_k3_triton_warmup(worker: Worker) -> None: """Warm Kimi-K3 Triton kernels reachable by this server.""" + if current_platform.is_rocm(): + layer = _get_kda_layer(worker, "vllm.models.kimi_k3.amd.kda") + if layer is not None: + _warm_gluon_chunk_kda(layer, worker.model_config.dtype) + return if not current_platform.is_cuda(): return diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 611a0c5d96..08ad0b650d 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -95,7 +95,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): assert isinstance(spec, MambaSpec) # Only the fused path exports the checkpoint snapshots can_checkpoint = ( - self.use_prefill_checkpoint and self.use_fused_chunk and self.use_safe_gate + self.use_prefill_checkpoint + and self.kda_prefill_backend == "fused" + and self.use_safe_gate ) if can_checkpoint and vllm_config.cache_config.mamba_cache_mode == "align": logger.info_once("Kimi-K3 KDA prefill checkpoint enabled.") @@ -242,13 +244,15 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): "The fused KDA chunk kernel requires gfx950 and a build that " "includes it." ) - self.use_fused_chunk = backend == "fused" or ( - backend == "auto" and is_fused_kda_chunk_supported() - ) - logger.info_once( - "Kimi-K3 KDA prefill backend: %s", - "fused" if self.use_fused_chunk else "triton", - ) + self.kda_prefill_backend = backend + if backend == "auto": + if self.use_safe_gate and rocm_aiter_ops.is_kda_gluon_prefill_available(): + self.kda_prefill_backend = "gluon" + elif is_fused_kda_chunk_supported(): + self.kda_prefill_backend = "fused" + else: + self.kda_prefill_backend = "triton" + logger.info_once("Kimi-K3 KDA prefill backend: %s", self.kda_prefill_backend) # TODO: Remove the extra config arg once tested self.use_prefill_checkpoint = ( additional_config.get("kda_prefill_checkpoint", True) @@ -694,7 +698,7 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): cu_seqlens=prefill_query_start_loc, chunk_indices=m.chunk_indices, chunk_offsets=m.chunk_offsets, - use_fused_chunk=self.use_fused_chunk, + backend=self.kda_prefill_backend, out=core_attn_out[:, nd_tok:num_actual_tokens] if direct else None, checkpoint_state=checkpoint_state, checkpoint_offsets=checkpoint_offsets, diff --git a/vllm/models/kimi_k3/amd/ops/kda_prefill.py b/vllm/models/kimi_k3/amd/ops/kda_prefill.py index 5ba2fd6379..040f3afe66 100644 --- a/vllm/models/kimi_k3/amd/ops/kda_prefill.py +++ b/vllm/models/kimi_k3/amd/ops/kda_prefill.py @@ -9,6 +9,7 @@ path. import torch +from vllm._aiter_ops import rocm_aiter_ops from vllm.logger import init_logger from vllm.model_executor.layers.mamba.ops.gather_initial_states import ( gather_initial_states, @@ -24,6 +25,11 @@ from vllm.third_party.flash_linear_attention.ops.utils import FLA_CHUNK_SIZE logger = init_logger(__name__) +# Consumed by kimi_k3_triton_warmup.py during kernel_warmup(). +def get_gluon_chunk_kda_warmup_profiles() -> tuple[tuple[int, int], ...]: + return ((1, 64), (2, 64), (4, 64), (8, 64), (16, 64), (1, 4096), (2, 2048)) + + def chunk_kda_prefill( q: torch.Tensor, k: torch.Tensor, @@ -40,7 +46,7 @@ def chunk_kda_prefill( cu_seqlens: torch.Tensor | None = None, chunk_indices: torch.Tensor | None = None, chunk_offsets: torch.Tensor | None = None, - use_fused_chunk: bool = False, + backend: str = "triton", out: torch.Tensor | None = None, checkpoint_state: torch.Tensor | None = None, checkpoint_offsets: torch.Tensor | None = None, @@ -68,9 +74,10 @@ def chunk_kda_prefill( cu_seqlens: int32 cumulative sequence lengths. chunk_indices: precomputed chunk indices for ``cu_seqlens``. chunk_offsets: int32 per-sequence first chunk index. - use_fused_chunk: request the two-kernel ROCm path. It is used only when - every one of its preconditions holds; otherwise the Triton path - runs unchanged. + backend: ``"triton"`` (the vendored chunk path), ``"fused"`` (the + two-kernel ROCm path; used only when every one of its + preconditions holds, otherwise the Triton path runs unchanged) or + ``"gluon"`` (AITER's ``chunk_kda``). out: buffer the result must land in. Honoured by both backends, so the caller can hand in a slice of its own output and skip a copy. checkpoint_state: destination for mid-prefill recurrent state @@ -89,6 +96,11 @@ def chunk_kda_prefill( The output and, when requested, the final recurrent state. """ + if backend not in ("triton", "fused", "gluon"): + raise ValueError( + "chunk_kda_prefill backend must be 'triton', 'fused' or 'gluon', " + f"got {backend!r}." + ) if scale is None: scale = k.shape[-1] ** -0.5 @@ -96,7 +108,7 @@ def chunk_kda_prefill( # needs the raw projections and a bounded gate rather than the pre-normalized # tensors the Triton path takes. fused = ( - use_fused_chunk + backend == "fused" and use_qk_l2norm_in_kernel and cu_seqlens is not None and lower_bound is not None @@ -116,6 +128,28 @@ def chunk_kda_prefill( if state_indices is None or has_initial_state is None: raise ValueError("state_cache needs state_indices and has_initial_state") + if backend == "gluon": + return rocm_aiter_ops.kda_chunk_prefill( + q=q, + k=k, + v=v, + g=raw_g, + beta=raw_beta, + A_log=A_log, + dt_bias=g_bias, + lower_bound=lower_bound, + cu_seqlens=cu_seqlens, + chunk_indices=chunk_indices, + chunk_offsets=chunk_offsets, + scale=scale, + out=out, + initial_state=initial_state, + output_final_state=output_final_state, + state_cache=state_cache, + state_indices=state_indices, + has_initial_state=has_initial_state, + ) + # Only the fused walk addresses the paged rows directly. The Triton path # gathers the rows it needs and scatters the results back here, so callers # pass the cache the same way for either backend. From 2268d83bd9eb9a60138b3591ad9ec8c170b093df Mon Sep 17 00:00:00 2001 From: alp Date: Fri, 16 Oct 2026 16:09:12 -0700 Subject: [PATCH 26/49] [Bugfix][Multimodal] Honor RGB PNG transparency when loading images (#61986) Signed-off-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> Co-authored-by: mottopanikeiku <176798723+mottopanikeiku@users.noreply.github.com> --- docs/features/multimodal_inputs.md | 6 +-- tests/multimodal/media/test_image.py | 18 ++++++++ tests/multimodal/test_hasher.py | 12 ++++- tests/multimodal/test_image.py | 65 ++++++++++++++++++++++++++++ vllm/multimodal/image.py | 6 +-- vllm/multimodal/media/image.py | 11 ++--- 6 files changed, 102 insertions(+), 16 deletions(-) diff --git a/docs/features/multimodal_inputs.md b/docs/features/multimodal_inputs.md index cb091c1c79..e4cd5f56a2 100644 --- a/docs/features/multimodal_inputs.md +++ b/docs/features/multimodal_inputs.md @@ -187,7 +187,7 @@ Multi-image input can be extended to perform video captioning. We show this with #### Custom RGBA Background Color -When loading RGBA images (images with transparency), vLLM converts them to RGB format. By default, transparent pixels are replaced with white background. You can customize this background color using the `rgba_background_color` parameter in `media_io_kwargs`. +When loading images with transparency (including RGBA images and PNGs with a tRNS chunk), vLLM converts them to RGB format. By default, transparent pixels are replaced with white background. You can customize this background color using the `rgba_background_color` parameter in `media_io_kwargs`. ??? code @@ -212,7 +212,7 @@ When loading RGBA images (images with transparency), vLLM converts them to RGB f !!! note - The `rgba_background_color` accepts RGB values as a list `[R, G, B]` or tuple `(R, G, B)` where each value is 0-255 - - This setting only affects RGBA images with transparency; RGB images are unchanged + - This setting affects images with transparency, including RGB PNGs with a transparent color; opaque RGB images are unchanged - If not specified, the default white background `(255, 255, 255)` is used for backward compatibility #### Moondream3 Prompt Recipes { #moondream3-prompt-recipes } @@ -1089,7 +1089,7 @@ By passing this metadata, the model can better understand the temporal distribut #### Custom RGBA Background Color -To use a custom background color for RGBA images, pass the `rgba_background_color` parameter via `--media-io-kwargs`: +To use a custom background color for images with transparency, pass the `rgba_background_color` parameter via `--media-io-kwargs`: ```bash # Example: Black background for dark theme diff --git a/tests/multimodal/media/test_image.py b/tests/multimodal/media/test_image.py index 60cfb86cc0..1980adf197 100644 --- a/tests/multimodal/media/test_image.py +++ b/tests/multimodal/media/test_image.py @@ -85,6 +85,24 @@ def test_image_media_io_rgba_custom_background(tmp_path): assert green_numpy[0][0][2] == 0 # B +@pytest.mark.parametrize("background_color", [(255, 255, 255), (12, 34, 56)]) +def test_image_media_io_rgb_png_transparency(tmp_path, background_color): + image = Image.new("RGB", (2, 1), (255, 0, 0)) + image.putpixel((1, 0), (0, 0, 255)) + path = tmp_path / "rgb_transparency.png" + image.save(path, transparency=(255, 0, 0)) + + with Image.open(path) as source: + reference = Image.alpha_composite( + Image.new("RGBA", source.size, (*background_color, 255)), + source.convert("RGBA"), + ).convert("RGB") + + image_io = ImageMediaIO(rgba_background_color=background_color) + converted = image_io.load_file(path) + np.testing.assert_array_equal(np.asarray(converted), np.asarray(reference)) + + def test_image_media_io_no_mode_conversion(tmp_path): """image_mode=None skips conversion and preserves the original mode.""" # RGBA image: opaque black pixel on a fully transparent background diff --git a/tests/multimodal/test_hasher.py b/tests/multimodal/test_hasher.py index b8d9dc7c9d..bdc1e98421 100644 --- a/tests/multimodal/test_hasher.py +++ b/tests/multimodal/test_hasher.py @@ -237,8 +237,16 @@ def _rgba_png_bytes() -> bytes: return buf.getvalue() -def test_hash_collision_media_io_config(): - data = _rgba_png_bytes() +def _rgb_trns_png_bytes() -> bytes: + image = Image.new("RGB", (8, 8), (255, 0, 0)) + buf = BytesIO() + image.save(buf, format="PNG", transparency=(255, 0, 0)) + return buf.getvalue() + + +@pytest.mark.parametrize("png_bytes", [_rgba_png_bytes, _rgb_trns_png_bytes]) +def test_hash_collision_media_io_config(png_bytes): + data = png_bytes() white = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data) black = ImageMediaIO(rgba_background_color=(0, 0, 0)).load_bytes(data) white2 = ImageMediaIO(rgba_background_color=(255, 255, 255)).load_bytes(data) diff --git a/tests/multimodal/test_image.py b/tests/multimodal/test_image.py index ea7f18611c..4365321cf4 100644 --- a/tests/multimodal/test_image.py +++ b/tests/multimodal/test_image.py @@ -5,6 +5,7 @@ from io import BytesIO from pathlib import Path import numpy as np +import pybase64 import pytest from PIL import Image, ImageChops @@ -13,6 +14,8 @@ from vllm.multimodal.image import ( convert_image_mode, normalize_image, ) +from vllm.multimodal.media import MediaConnector +from vllm.multimodal.parse import MultiModalDataParser pytestmark = pytest.mark.cpu_test @@ -69,6 +72,68 @@ def test_palette_with_trns_to_rgb(): assert (r, g, b) == (255, 0, 0) +def _reopen(image: Image.Image, image_format: str, **save_kwargs) -> Image.Image: + buffer = BytesIO() + image.save(buffer, image_format, **save_kwargs) + return Image.open(buffer) + + +def _rgb_png() -> Image.Image: + image = Image.new("RGB", (2, 1), (255, 0, 0)) + image.putpixel((1, 0), (0, 0, 255)) + return _reopen(image, "PNG", transparency=(255, 0, 0)) + + +def _gif() -> Image.Image: + # The transparent index 0 and the opaque index 1 are both black. + image = Image.new("P", (2, 1)) + image.putpalette([0] * 768) + image.putpixel((1, 0), 1) + return _reopen(image, "GIF", transparency=0) + + +def _rgba_png() -> Image.Image: + image = Image.new("RGBA", (2, 1), (255, 0, 0, 0)) + image.putpixel((1, 0), (0, 0, 255, 255)) + return _reopen(image, "PNG") + + +@pytest.mark.parametrize( + "make_image", + [ + _rgb_png, + lambda: _rgb_png().convert("RGB"), + _gif, + lambda: _gif().convert("RGB"), + _rgba_png, + ], + ids=["rgb_png", "rgb_png_to_rgb", "gif", "gif_to_rgb", "rgba_png"], +) +def test_transparency_offline_matches_online(make_image): + """A PIL image passed offline and the same image sent online as a PNG + end up with the same pixels, matching Pillow's own RGBA conversion. + + After ``convert("RGB")``, Pillow keeps the transparent color in ``info`` + and saves it as a tRNS chunk, so every pixel of that color is + transparent in both pipelines, including opaque GIF pixels that share it. + """ + image = make_image() + buffer = BytesIO() + image.save(buffer, "PNG") + url = f"data:image/png;base64,{pybase64.b64encode(buffer.getvalue()).decode()}" + + parser = MultiModalDataParser() + offline = parser.parse_mm_data({"image": image})["image"].get(0) + online = parser.parse_mm_data({"image": MediaConnector().fetch_image(url)}) + online = online["image"].get(0) + np.testing.assert_array_equal(np.asarray(offline), np.asarray(online)) + + expected = Image.alpha_composite( + Image.new("RGBA", image.size, (255, 255, 255, 255)), image.convert("RGBA") + ).convert("RGB") + np.testing.assert_array_equal(np.asarray(online), np.asarray(expected)) + + def test_l_mode_no_trns_to_rgb(): """L-mode without transparency should convert directly.""" img = Image.new("L", (4, 4), 128) diff --git a/vllm/multimodal/image.py b/vllm/multimodal/image.py index 13df7f3c9a..96eb606e60 100644 --- a/vllm/multimodal/image.py +++ b/vllm/multimodal/image.py @@ -104,12 +104,12 @@ def convert_image_mode( to_mode: str, background_color: tuple[int, int, int] | list[int] = (255, 255, 255), ) -> Image.Image: - if image.mode == to_mode: - return image - if to_mode == "RGB" and _has_transparency(image): if image.mode != "RGBA": image = image.convert("RGBA") return rgba_to_rgb(image, background_color) + if image.mode == to_mode: + return image + return image.convert(to_mode) diff --git a/vllm/multimodal/media/image.py b/vllm/multimodal/media/image.py index 1e2e106189..5dbd6e9381 100644 --- a/vllm/multimodal/media/image.py +++ b/vllm/multimodal/media/image.py @@ -16,7 +16,7 @@ from vllm.utils.sparse_utils import ( safe_to_dense, ) -from ..image import convert_image_mode, normalize_image, open_image, rgba_to_rgb +from ..image import convert_image_mode, normalize_image, open_image from .base import MediaDecodedBudget, MediaIO, MediaWithBytes MAGIC_NUMPY_PREFIX = b"\x93NUMPY" # https://numpy.org/devdocs/reference/generated/numpy.lib.format.html#format-version-1-0 @@ -74,14 +74,9 @@ class ImageMediaIO(MediaIO[Image.Image]): """Convert image mode with custom background color.""" if isinstance(image, MediaWithBytes): image = image.media - if self.image_mode is None or image.mode == self.image_mode: + if self.image_mode is None: return image - elif image.mode == "RGBA" and self.image_mode == "RGB": - return rgba_to_rgb(image, self.rgba_background_color) - else: - return convert_image_mode( - image, self.image_mode, self.rgba_background_color - ) + return convert_image_mode(image, self.image_mode, self.rgba_background_color) def load_bytes(self, data: bytes) -> MediaWithBytes[Image.Image]: try: From b5445874080ce85c575f2bcb847852235f1af7bc Mon Sep 17 00:00:00 2001 From: wangxiyuan Date: Sat, 17 Oct 2026 07:16:47 +0800 Subject: [PATCH 27/49] [Core][Feat] Pluggable KVCacheConfigBuilder for platform/model-specific KV cache planning (#53558) Signed-off-by: wangxiyuan Signed-off-by: Yifan Qiao Co-authored-by: Claude Opus 5.5 Co-authored-by: Yifan Qiao --- tests/v1/core/test_kv_cache_config_builder.py | 294 ++++++++++++++++++ vllm/config/model.py | 4 + vllm/model_executor/models/interfaces.py | 42 +++ vllm/model_executor/models/registry.py | 3 + vllm/platforms/interface.py | 22 ++ vllm/v1/core/kv_cache_config_builder.py | 103 ++++++ vllm/v1/core/kv_cache_planning.py | 214 +++++++++++-- vllm/v1/worker/gpu/cudagraph_utils.py | 22 +- vllm/v1/worker/gpu_model_runner.py | 20 +- 9 files changed, 675 insertions(+), 49 deletions(-) create mode 100644 tests/v1/core/test_kv_cache_config_builder.py create mode 100644 vllm/v1/core/kv_cache_config_builder.py diff --git a/tests/v1/core/test_kv_cache_config_builder.py b/tests/v1/core/test_kv_cache_config_builder.py new file mode 100644 index 0000000000..4f4c8d0262 --- /dev/null +++ b/tests/v1/core/test_kv_cache_config_builder.py @@ -0,0 +1,294 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for KVCacheConfigBuilder resolution.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from vllm.platforms import Platform +from vllm.v1.core.kv_cache_config_builder import ( + KVCacheConfigBuilder, + get_kv_cache_config_builder, +) +from vllm.v1.core.kv_cache_planning import ( + DefaultKVCacheConfigBuilder, + get_kv_cache_groups_from_builder, +) +from vllm.v1.kv_cache_interface import FullAttentionSpec, KVCacheGroupSpec + + +def _make_vllm_config(supports_grouping: bool = False) -> MagicMock: + """Create a minimal mock VllmConfig for builder resolution tests. + + ``supports_grouping`` selects whether the resolved model class declares + ``SupportsKVCacheGrouping`` (exposing a ``get_kv_cache_groups`` classmethod). + """ + cfg = MagicMock() + cfg.model_config.supports_kv_cache_grouping = supports_grouping + cfg.attention_config.hisparse_config = None + if supports_grouping: + cfg.model_config.registry.resolve_model_cls.return_value = ( + _GroupedModel, + "arch", + ) + return cfg + + +CUSTOM_PATH = "tests.v1.core.test_kv_cache_config_builder.CustomBuilder" +DEFAULT_PATH = "vllm.v1.core.kv_cache_planning.DefaultKVCacheConfigBuilder" + + +class CustomBuilder(DefaultKVCacheConfigBuilder): + """A test builder subclass.""" + + pass + + +class _GroupedModel: + """A test model class with a ``get_kv_cache_groups`` classmethod.""" + + @classmethod + def get_kv_cache_groups(cls, vllm_config, kv_cache_spec): + return [] + + +class _PartialGroupedModel: + """A test model class whose grouping omits layers.""" + + @classmethod + def get_kv_cache_groups(cls, vllm_config, kv_cache_spec): + from vllm.v1.kv_cache_interface import KVCacheGroupSpec + + return [KVCacheGroupSpec(["a"], kv_cache_spec["a"])] + + +class TestPlatformHookResolution: + """The platform hook selects only the placement builder.""" + + def test_default_hook_returns_default_builder(self): + cfg = _make_vllm_config() + assert Platform.get_kv_cache_config_builder_cls(cfg) == DEFAULT_PATH + + +class TestBuilderResolution: + @patch("vllm.platforms.current_platform") + @pytest.mark.parametrize( + ("builder_cls_path", "expected_cls"), + [ + (None, DefaultKVCacheConfigBuilder), + (CUSTOM_PATH, CustomBuilder), + ], + ) + def test_resolves_hook_selected_builder( + self, mock_platform, builder_cls_path, expected_cls + ): + mock_platform.get_kv_cache_config_builder_cls.return_value = ( + builder_cls_path or DEFAULT_PATH + ) + cfg = _make_vllm_config() + assert type(get_kv_cache_config_builder(cfg)) is expected_cls + + @patch("vllm.platforms.current_platform") + def test_resolves_fresh_instance_per_call(self, mock_platform): + mock_platform.get_kv_cache_config_builder_cls.return_value = CUSTOM_PATH + cfg = _make_vllm_config() + assert get_kv_cache_config_builder(cfg) is not get_kv_cache_config_builder(cfg) + + def test_public_interface_declares_two_placement_hooks(self): + # get_kv_cache_configs lives in kv_cache_planning and is Core-owned; + # the builder only declares the placement hooks. + assert KVCacheConfigBuilder.__abstractmethods__ == { + "get_kv_cache_groups", + "get_kv_cache_config_from_groups", + } + assert issubclass(DefaultKVCacheConfigBuilder, KVCacheConfigBuilder) + assert ( + "pool_bytes_per_block" + not in DefaultKVCacheConfigBuilder.__abstractmethods__ + ) + assert ( + "max_memory_usage_bytes" + not in DefaultKVCacheConfigBuilder.__abstractmethods__ + ) + + +class TestModelGroupingPriority: + """A model class with a get_kv_cache_groups classmethod owns grouping.""" + + def test_model_classmethod_owns_grouping(self): + cfg = _make_vllm_config(supports_grouping=True) + assert DefaultKVCacheConfigBuilder().get_kv_cache_groups(cfg, {}) == [] + + def test_model_grouping_rejects_hisparse(self): + cfg = _make_vllm_config(supports_grouping=True) + cfg.attention_config.hisparse_config = object() + with pytest.raises(ValueError, match="HiSparse"): + DefaultKVCacheConfigBuilder().get_kv_cache_groups(cfg, {}) + + +class TestCoreChecksGrouping: + """Core checks the grouping whether a model or a platform produced it.""" + + _SPEC = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=128, dtype="float16" + ) + + def test_model_grouping_is_checked(self): + cfg = _make_vllm_config(supports_grouping=True) + cfg.model_config.registry.resolve_model_cls.return_value = ( + _PartialGroupedModel, + "arch", + ) + with pytest.raises(AssertionError, match="miss layers"): + get_kv_cache_groups_from_builder( + DefaultKVCacheConfigBuilder(), cfg, {"a": self._SPEC, "b": self._SPEC} + ) + + @pytest.mark.parametrize( + ("groups", "match"), + [ + ([KVCacheGroupSpec(["a"], _SPEC)], "miss layers"), + ( + [KVCacheGroupSpec(["a", "b"], _SPEC), KVCacheGroupSpec(["b"], _SPEC)], + "overlaps", + ), + ], + ) + def test_platform_grouping_is_checked(self, groups, match): + class FixedGroupsBuilder(DefaultKVCacheConfigBuilder): + def get_kv_cache_groups(self, vllm_config, kv_cache_spec): + return groups + + with pytest.raises(AssertionError, match=match): + get_kv_cache_groups_from_builder( + FixedGroupsBuilder(), + _make_vllm_config(), + {"a": self._SPEC, "b": self._SPEC}, + ) + + +class TestPlatformCustomPriority: + """A vendor platform can override the hook to customize priority.""" + + def test_platform_builder_wins_over_model_declaration(self): + class PlatformFirstPlatform(Platform): + @classmethod + def get_kv_cache_config_builder_cls(cls, vllm_config): + return CUSTOM_PATH + + cfg = _make_vllm_config() + assert PlatformFirstPlatform.get_kv_cache_config_builder_cls(cfg) == CUSTOM_PATH + with patch("vllm.platforms.current_platform", PlatformFirstPlatform): + assert isinstance(get_kv_cache_config_builder(cfg), CustomBuilder) + + +class TestProjectionPreservesGroupMetadata: + """PP projection must carry host_resident/role/enable_kv_transfer.""" + + def test_partial_projection_preserves_metadata(self): + from vllm.v1.core.kv_cache_planning import _project_kv_cache_groups_to_worker + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupRole, + ) + + spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=128, dtype="float16" + ) + group = MagicMock() + group.layer_names = ["layer1", "layer2", "layer3"] + group.kv_cache_spec = spec + group.is_eagle_group = False + group.host_resident = True + group.enable_kv_transfer = False + group.role = KVCacheGroupRole.HISPARSE_SOURCE + + projected = _project_kv_cache_groups_to_worker( + [group], {"layer1": spec, "layer2": spec} + ) + assert len(projected) == 1 + p = projected[0] + assert p.host_resident is True + assert p.enable_kv_transfer is False + assert p.role is KVCacheGroupRole.HISPARSE_SOURCE + assert p.layer_names == ["layer1", "layer2"] + + +class TestBytesPerBlockProbe: + """The single-block probe must not be polluted by num_gpu_blocks_override.""" + + def test_probe_suppresses_override(self): + import torch + + from vllm.config import CacheConfig, ModelConfig, VllmConfig + from vllm.v1.attention.backends.utils import resolve_kv_cache_layout + from vllm.v1.core.kv_cache_planning import ( + DefaultKVCacheConfigBuilder, + _pool_bytes_per_block, + _pool_bytes_per_block_from_builder, + ) + from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheGroupSpec, + ) + + spec = FullAttentionSpec( + block_size=16, num_kv_heads=1, head_size=128, dtype=torch.float16 + ) + vc = VllmConfig( + model_config=ModelConfig(max_model_len=100), + cache_config=CacheConfig(block_size=16), + ) + resolve_kv_cache_layout(vc, [["BLHNC"]], [spec]) + groups = [ + KVCacheGroupSpec(["a"], spec), + KVCacheGroupSpec(["b"], spec), + ] + + pure_c = _pool_bytes_per_block(vc, groups) + vc.cache_config.num_gpu_blocks_override = 500 + probed_c = _pool_bytes_per_block_from_builder( + DefaultKVCacheConfigBuilder(), vc, groups + ) + assert probed_c == pure_c + assert vc.cache_config.num_gpu_blocks_override == 500 + + @pytest.mark.parametrize("declares_pool_bytes", [False, True]) + def test_several_device_allocations_need_declared_pool_bytes( + self, declares_pool_bytes + ): + """A placement with one allocation per group cannot be sized from one + backing; the builder must declare its per-block bytes.""" + from vllm.v1.core.kv_cache_planning import _pool_bytes_per_block_from_builder + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheTensor, + ) + + class PerGroupAllocationBuilder(DefaultKVCacheConfigBuilder): + def get_kv_cache_config_from_groups(self, vllm_config, groups, num_blocks): + tensors = [ + KVCacheTensor( + size=page * num_blocks, + layers=group.layer_names, + layer_stride=page * num_blocks, + block_stride=page, + ) + for group, page in zip(groups, (64, 32)) + ] + return KVCacheConfig(num_blocks, tensors, groups) + + def _get_pool_bytes_per_block(self, vllm_config, groups): + return 64 + 32 if declares_pool_bytes else None + + vc = _make_vllm_config() + vc.cache_config.num_gpu_blocks_override = None + groups = [KVCacheGroupSpec([name], MagicMock()) for name in "ab"] + builder = PerGroupAllocationBuilder() + if declares_pool_bytes: + assert _pool_bytes_per_block_from_builder(builder, vc, groups) == 96 + else: + with pytest.raises(ValueError, match="_get_pool_bytes_per_block"): + _pool_bytes_per_block_from_builder(builder, vc, groups) diff --git a/vllm/config/model.py b/vllm/config/model.py index 7a883eac54..6dc2d16d52 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -2034,6 +2034,10 @@ class ModelConfig: def has_noops(self) -> bool: return self._model_info.has_noops + @property + def supports_kv_cache_grouping(self) -> bool: + return self._model_info.supports_kv_cache_grouping + @property def has_inner_state(self): return self._model_info.has_inner_state diff --git a/vllm/model_executor/models/interfaces.py b/vllm/model_executor/models/interfaces.py index 5f51ece0b3..818c161050 100644 --- a/vllm/model_executor/models/interfaces.py +++ b/vllm/model_executor/models/interfaces.py @@ -64,6 +64,7 @@ if TYPE_CHECKING: from vllm.tasks import ScoreType from vllm.tokenizers import TokenizerLike from vllm.v1.attention.backends.registry import MambaAttentionBackendEnum + from vllm.v1.kv_cache_interface import KVCacheGroupSpec, KVCacheSpec from vllm.v1.worker.encoder_cudagraph_defs import ( EncoderCudaGraphCaptureInputs, EncoderCudaGraphConfig, @@ -669,6 +670,47 @@ def supports_score_template( return getattr(model, "supports_score_template", False) +@runtime_checkable +class SupportsKVCacheGrouping(Protocol): + """The interface required for models that own their KV cache grouping. + + .. note:: + Experimental. The classmethod contract may change as model-owned + grouping stabilizes. + + Implementing it lets a model replace the default KV cache grouping (which + layers share a block table) while Core keeps owning placement, byte + accounting, and admission. + """ + + @classmethod + def get_kv_cache_groups( + cls, + vllm_config: "VllmConfig", + kv_cache_spec: dict[str, "KVCacheSpec"], + ) -> list["KVCacheGroupSpec"]: + """Split the model's layers into logical KV cache groups.""" + ... + + +@overload +def supports_kv_cache_grouping( + model: type[object], +) -> TypeIs[type[SupportsKVCacheGrouping]]: ... + + +@overload +def supports_kv_cache_grouping( + model: object, +) -> TypeIs[SupportsKVCacheGrouping]: ... + + +def supports_kv_cache_grouping( + model: type[object] | object, +) -> TypeIs[type[SupportsKVCacheGrouping]] | TypeIs[SupportsKVCacheGrouping]: + return isinstance(model, SupportsKVCacheGrouping) + + @runtime_checkable class SupportsLoRA(Protocol): """The interface required for all models that support LoRA.""" diff --git a/vllm/model_executor/models/registry.py b/vllm/model_executor/models/registry.py index 23e67c9f50..50280061e2 100644 --- a/vllm/model_executor/models/registry.py +++ b/vllm/model_executor/models/registry.py @@ -50,6 +50,7 @@ from .interfaces import ( is_hybrid, requires_raw_input_tokens, supports_hc_sp, + supports_kv_cache_grouping, supports_multimodal, supports_multimodal_encoder_tp_data, supports_multimodal_raw_input_only, @@ -882,6 +883,7 @@ class _ModelInfo: supports_replayssm: bool supports_transcription: bool supports_transcription_only: bool + supports_kv_cache_grouping: bool supported_video_pruning_methods: tuple[str, ...] supports_mm_device_do_normalize: bool @@ -913,6 +915,7 @@ class _ModelInfo: supports_transcription_only=( supports_transcription(model) and model.supports_transcription_only ), + supports_kv_cache_grouping=supports_kv_cache_grouping(model), has_noops=has_noops(model), supported_video_pruning_methods=getattr( model, "supported_video_pruning_methods", () diff --git a/vllm/platforms/interface.py b/vllm/platforms/interface.py index 3034d7bc1a..65f165ad3a 100644 --- a/vllm/platforms/interface.py +++ b/vllm/platforms/interface.py @@ -1467,6 +1467,28 @@ class Platform: "Unknown vLLM environment variable detected: %s", env ) + @classmethod + def get_kv_cache_config_builder_cls(cls, vllm_config: "VllmConfig") -> str: + """Resolve the fully-qualified class path of the platform's KV cache + config builder. + + .. note:: + Experimental, like the ``KVCacheConfigBuilder`` interface it selects. + + The builder supplies grouping (``get_kv_cache_groups``) and tensor + placement (``get_kv_cache_config_from_groups``). Vendors may override + this to return a platform-specific builder. The default builder honors a + model implementing + :class:`~vllm.model_executor.models.interfaces.SupportsKVCacheGrouping`; + a builder that overrides ``get_kv_cache_groups`` owns grouping and + bypasses that model hook unless it calls ``super()``. + + Returns: + Fully-qualified class path of a ``KVCacheConfigBuilder`` subclass. + + """ + return "vllm.v1.core.kv_cache_planning.DefaultKVCacheConfigBuilder" + class UnspecifiedPlatform(Platform): _enum = PlatformEnum.UNSPECIFIED diff --git a/vllm/v1/core/kv_cache_config_builder.py b/vllm/v1/core/kv_cache_config_builder.py new file mode 100644 index 0000000000..09332946cc --- /dev/null +++ b/vllm/v1/core/kv_cache_config_builder.py @@ -0,0 +1,103 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Interface and resolution for pluggable KV cache config builders.""" + +from abc import ABC, abstractmethod +from typing import TYPE_CHECKING + +from vllm.utils.import_utils import resolve_obj_by_qualname + +if TYPE_CHECKING: + from vllm.config import VllmConfig + from vllm.v1.kv_cache_interface import ( + KVCacheConfig, + KVCacheGroupSpec, + KVCacheSpec, + ) + + +class KVCacheConfigBuilder(ABC): + """Interface for platform KV cache placement. + + .. note:: + This interface is experimental and subject to change. The exact hooks + and the ``num_blocks`` contract may evolve before it stabilizes. + + Subclasses normally inherit from ``DefaultKVCacheConfigBuilder`` and + override only the hooks they need. The cross-worker planning flow + (:func:`vllm.v1.core.kv_cache_planning.get_kv_cache_configs`) is Core-owned; + this builder supplies the placement decisions: how layers are split into + groups and how a group list is materialized into tensors for a given block + count. Grouping coordination lives in ``DefaultKVCacheConfigBuilder`` (a + model implementing :class:`~vllm.model_executor.models.interfaces. + SupportsKVCacheGrouping` owns its grouping there); a platform that replaces + the builder fully owns its own grouping. + """ + + @abstractmethod + def get_kv_cache_groups( + self, + vllm_config: "VllmConfig", + kv_cache_spec: dict[str, "KVCacheSpec"], + ) -> list["KVCacheGroupSpec"]: + """Split the given layers into KV cache groups. + + Planning passes all workers' layers; profiling passes one worker's. + Core rejects a result that does not put each layer in exactly one group. + + The default builder honors a model's + :class:`~vllm.model_executor.models.interfaces.SupportsKVCacheGrouping` + declaration here before falling back to the default grouping; a + platform-specific builder decides its own grouping. + """ + raise NotImplementedError + + @abstractmethod + def get_kv_cache_config_from_groups( + self, + vllm_config: "VllmConfig", + kv_cache_groups: list["KVCacheGroupSpec"], + num_blocks: int, + ) -> "KVCacheConfig": + """Materialize the groups for exactly ``num_blocks`` blocks. + + Each device backing's size is ``bytes_per_block * num_blocks``. All + device tensors share one backing allocation unless the builder + overrides ``_get_pool_bytes_per_block``. Host-resident tensors, such as + HiSparse source storage, are excluded from the device byte accounting + and sized independently. + """ + raise NotImplementedError + + def _get_pool_bytes_per_block( + self, + vllm_config: "VllmConfig", + kv_cache_groups: list["KVCacheGroupSpec"], + ) -> int | None: + """Return the device bytes one block ID consumes, or ``None``. + + Private and transitional; expected to go away with worker-driven + allocation (RFC #42449). With ``None`` (the default), Core derives the + value from a one-block placement with one device backing. Override it + only for placements with several device allocations. + """ + return None + + +def get_kv_cache_config_builder( + vllm_config: "VllmConfig", +) -> KVCacheConfigBuilder: + """Resolve the builder selected by the current platform. + + Resolution priority is owned by the platform hook + (:meth:`vllm.platforms.interface.Platform.get_kv_cache_config_builder_cls`). + Builders are stateless (every method takes ``vllm_config``), so a fresh + instance is returned per call and engines with different configs in the + same process always get the correct builder. + """ + from vllm.platforms import current_platform + + builder_cls = resolve_obj_by_qualname( + current_platform.get_kv_cache_config_builder_cls(vllm_config) + ) + return builder_cls() diff --git a/vllm/v1/core/kv_cache_planning.py b/vllm/v1/core/kv_cache_planning.py index 6e258a68a9..10016409af 100644 --- a/vllm/v1/core/kv_cache_planning.py +++ b/vllm/v1/core/kv_cache_planning.py @@ -14,6 +14,7 @@ from vllm.logger import init_logger from vllm.utils.math_utils import cdiv, round_up from vllm.utils.mem_utils import format_gib from vllm.utils.torch_utils import get_dtype_size +from vllm.v1.core.kv_cache_config_builder import KVCacheConfigBuilder from vllm.v1.core.kv_cache_utils import _kv_cache_device_name from vllm.v1.hisparse.layout import ( get_hisparse_host_pool_bytes, @@ -261,6 +262,48 @@ def _pool_bytes_per_block( ) +def _pool_bytes_per_block_from_builder( + builder: KVCacheConfigBuilder, + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], +) -> int: + """Bytes consumed by one block in the worker's device KV cache pool. + + Core places one block with the builder and reads the size of its single + device backing. A placement with several device allocations declares the + value through ``KVCacheConfigBuilder._get_pool_bytes_per_block``. HiSparse + keeps the pure computation, since probing its placement would check the host + pool against ``max_model_len`` before auto-fit shortens it. + ``num_gpu_blocks_override`` is suppressed so it cannot inflate the probe. + """ + if not kv_cache_groups: + return 0 + if vllm_config.attention_config.hisparse_config is not None: + return _pool_bytes_per_block(vllm_config, kv_cache_groups) + pool_bytes = builder._get_pool_bytes_per_block(vllm_config, kv_cache_groups) + if pool_bytes is not None: + return pool_bytes + saved_override = vllm_config.cache_config.num_gpu_blocks_override + vllm_config.cache_config.num_gpu_blocks_override = None + try: + probe = builder.get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, num_blocks=1 + ) + finally: + vllm_config.cache_config.num_gpu_blocks_override = saved_override + backing_sizes = { + tensor.size for tensor in probe.kv_cache_tensors if not tensor.host_resident + } + if len(backing_sizes) != 1: + raise ValueError( + "KV cache tensors must share one device backing allocation, but got " + f"sizes {sorted(backing_sizes)}. Override " + "KVCacheConfigBuilder._get_pool_bytes_per_block() for placements " + "with several device allocations." + ) + return backing_sizes.pop() + + def get_uniform_page_size(kv_cache_specs: Iterable[KVCacheSpec]) -> int: """Get the page size of the KV cache.""" page_sizes = {layer.page_size_bytes for layer in kv_cache_specs} @@ -1264,6 +1307,36 @@ def get_kv_cache_groups( return groups +def get_kv_cache_groups_from_builder( + builder: KVCacheConfigBuilder, + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], +) -> list[KVCacheGroupSpec]: + """Group the layers with the builder and check the result. + + Core checks every grouping, whether it comes from the default grouping, a + model's ``SupportsKVCacheGrouping`` or a platform builder. + """ + groups = builder.get_kv_cache_groups(vllm_config, kv_cache_spec) + _check_exact_partition(groups, kv_cache_spec) + return groups + + +def _check_exact_partition( + groups: list[KVCacheGroupSpec], + kv_cache_spec: dict[str, KVCacheSpec], +) -> None: + """Ensure the groups partition every input layer exactly once.""" + grouped_layers: set[str] = set() + for group in groups: + assert not set(group.layer_names) & grouped_layers, ( + f"KV cache group overlaps: {set(group.layer_names) & grouped_layers}." + ) + grouped_layers |= set(group.layer_names) + missing = set(kv_cache_spec) - grouped_layers + assert not missing, f"KV cache groups miss layers: {sorted(missing)}." + + def _layer_tp_replicas(spec: KVCacheSpec, tp_size: int, dcp_size: int) -> int: if not isinstance(spec, AttentionSpec) or spec.max_tp_shards is None: return 1 @@ -1285,6 +1358,7 @@ def kv_cache_groups_tp_replicas( def _max_memory_usage_bytes_from_groups( + builder: KVCacheConfigBuilder, vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], ) -> int: @@ -1301,7 +1375,9 @@ def _max_memory_usage_bytes_from_groups( if not kv_cache_groups: return 0 - bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) + bytes_per_block = _pool_bytes_per_block_from_builder( + builder, vllm_config, kv_cache_groups + ) total_blocks = 0 for group in kv_cache_groups: spec = group.kv_cache_spec @@ -1317,6 +1393,7 @@ def _max_memory_usage_bytes_from_groups( def _estimate_max_model_len_from_groups( + builder: KVCacheConfigBuilder, vllm_config: VllmConfig, kv_cache_groups: list[KVCacheGroupSpec], available_memory: int, @@ -1334,13 +1411,16 @@ def _estimate_max_model_len_from_groups( vllm_config.model_config.max_model_len = model_len if hisparse_enabled: try: - get_kv_cache_config_from_groups( - vllm_config, kv_cache_groups, available_memory + num_blocks = available_memory // _pool_bytes_per_block_from_builder( + builder, vllm_config, kv_cache_groups + ) + builder.get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, num_blocks ) except ValueError: return False return ( - _max_memory_usage_bytes_from_groups(vllm_config, kv_cache_groups) + _max_memory_usage_bytes_from_groups(builder, vllm_config, kv_cache_groups) <= available_memory ) @@ -1362,6 +1442,7 @@ def _estimate_max_model_len_from_groups( def _auto_fit_max_model_len( + builder: KVCacheConfigBuilder, vllm_config: VllmConfig, projected_groups_per_worker: list[list[KVCacheGroupSpec]], available_memory: list[int], @@ -1372,6 +1453,7 @@ def _auto_fit_max_model_len( workers. Args: + builder: The builder whose hooks drive block accounting. vllm_config: The global VllmConfig (will be modified in-place) projected_groups_per_worker: KV cache groups projected to each worker. available_memory: Memory available for KV cache in bytes for each @@ -1395,7 +1477,9 @@ def _auto_fit_max_model_len( for groups, avail_mem in zip(projected_groups_per_worker, available_memory): if not groups: continue - worker_max = _estimate_max_model_len_from_groups(vllm_config, groups, avail_mem) + worker_max = _estimate_max_model_len_from_groups( + builder, vllm_config, groups, avail_mem + ) if worker_max < auto_fit_max: auto_fit_max = worker_max limiting_worker_mem = avail_mem @@ -1473,6 +1557,9 @@ def _project_kv_cache_groups_to_worker( worker_layer_names, group_spec, is_eagle_group=group.is_eagle_group and bool(worker_layer_names), + host_resident=group.host_resident, + enable_kv_transfer=group.enable_kv_transfer, + role=group.role, ) ) return projected_groups @@ -1512,6 +1599,10 @@ def get_kv_cache_configs( The generated KVCacheConfigs for each worker. """ + from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder + + builder = get_kv_cache_config_builder(vllm_config) + # Merge the KV cache specs of all workers. Different PP stages may have # different layer names, and different TP ranks of the same PP stage should # have the same KV cache spec. @@ -1542,7 +1633,9 @@ def get_kv_cache_configs( # Get global KV cache groups. This also handles spec unification for # hybrid models when disable_hybrid_kv_cache_manager is enabled. # After this call, merged_kv_cache_specs may be modified in-place. - global_kv_cache_groups = get_kv_cache_groups(vllm_config, merged_kv_cache_specs) + global_kv_cache_groups = get_kv_cache_groups_from_builder( + builder, vllm_config, merged_kv_cache_specs + ) # If original_max_model_len was -1, automatically # determine the maximum model length that fits in available GPU memory. @@ -1553,11 +1646,10 @@ def get_kv_cache_configs( ] # If `num_gpu_blocks_override` is set, the cache size that will actually - # be allocated is decoupled from the profiled `available_memory`: - # `may_override_num_blocks` in `get_kv_cache_config_from_groups` clamps - # `num_blocks` to the override. Reflect that in `available_memory` here so - # auto-fit, the admission check, and the per-worker config builder all - # plan against the same effective capacity. + # be allocated is decoupled from the profiled `available_memory`. Reflect + # that in `available_memory` here so auto-fit, the admission check, and + # the per-worker config builder all plan against the same effective + # capacity. override = vllm_config.cache_config.num_gpu_blocks_override if override is not None: adjusted_memory: list[int] = [] @@ -1565,7 +1657,9 @@ def get_kv_cache_configs( if not groups: adjusted_memory.append(avail_mem) continue - bytes_per_block = _pool_bytes_per_block(vllm_config, groups) + bytes_per_block = _pool_bytes_per_block_from_builder( + builder, vllm_config, groups + ) logger.info( "Overriding num_gpu_blocks=%d with num_gpu_blocks_override=%d", avail_mem // bytes_per_block, @@ -1578,12 +1672,18 @@ def get_kv_cache_configs( # the capacity check both plan against usable blocks. Allocation below # still uses the full memory. check_memory = [ - avail_mem - _pool_bytes_per_block(vllm_config, groups) if groups else avail_mem + ( + avail_mem - _pool_bytes_per_block_from_builder(builder, vllm_config, groups) + if groups + else avail_mem + ) for groups, avail_mem in zip(projected_groups_per_worker, available_memory) ] if vllm_config.model_config.original_max_model_len == -1: - _auto_fit_max_model_len(vllm_config, projected_groups_per_worker, check_memory) + _auto_fit_max_model_len( + builder, vllm_config, projected_groups_per_worker, check_memory + ) # Check if the available memory is enough per worker. for groups, avail_mem in zip(projected_groups_per_worker, check_memory): @@ -1591,9 +1691,9 @@ def get_kv_cache_configs( continue _check_enough_kv_cache_memory( avail_mem, - partial(_max_memory_usage_bytes_from_groups, vllm_config, groups), + partial(_max_memory_usage_bytes_from_groups, builder, vllm_config, groups), vllm_config.model_config.max_model_len, - partial(_estimate_max_model_len_from_groups, vllm_config, groups), + partial(_estimate_max_model_len_from_groups, builder, vllm_config, groups), ) kv_cache_configs: list[KVCacheConfig] = [] @@ -1603,9 +1703,15 @@ def get_kv_cache_configs( assert sum(len(group.layer_names) for group in projected_groups) == len( kv_cache_spec_one_worker ), "Some layers are not assigned to any group." + bytes_per_block = _pool_bytes_per_block_from_builder( + builder, vllm_config, projected_groups + ) + num_blocks = ( + available_memory_one_worker // bytes_per_block if projected_groups else 1 + ) kv_cache_configs.append( - get_kv_cache_config_from_groups( - vllm_config, projected_groups, available_memory_one_worker + builder.get_kv_cache_config_from_groups( + vllm_config, projected_groups, num_blocks ) ) @@ -1618,13 +1724,12 @@ def get_kv_cache_configs( for i, kv_cache_config in enumerate(kv_cache_configs): if kv_cache_config.num_blocks == min_num_blocks: continue - # Re-plan with exactly the memory the smallest rank can afford, so - # strides and offsets stay consistent with the shrunken allocation. - groups = kv_cache_config.kv_cache_groups - kv_cache_configs[i] = get_kv_cache_config_from_groups( + # Re-plan with exactly the smallest rank's block count, so strides and + # offsets stay consistent with the shrunken allocation. + kv_cache_configs[i] = builder.get_kv_cache_config_from_groups( vllm_config, - groups, - min_num_blocks * _pool_bytes_per_block(vllm_config, groups), + kv_cache_config.kv_cache_groups, + min_num_blocks, ) for kv_cache_config in kv_cache_configs: @@ -1635,3 +1740,64 @@ def get_kv_cache_configs( ) return kv_cache_configs + + +class DefaultKVCacheConfigBuilder(KVCacheConfigBuilder): + """Default grouping and placement, backed by the module-level functions here. + + .. note:: + This class is part of the experimental ``KVCacheConfigBuilder`` + interface and is subject to change. + + Subclasses override only the hooks they need. The planning flow and its + invariants (``num_gpu_blocks_override``, null-block reservation, auto-fit, + admission checks, cross-rank convergence) stay in the Core-owned + :func:`get_kv_cache_configs`. + """ + + def get_kv_cache_groups( + self, + vllm_config: VllmConfig, + kv_cache_spec: dict[str, KVCacheSpec], + ) -> list[KVCacheGroupSpec]: + """See ``KVCacheConfigBuilder.get_kv_cache_groups``. + + A model implementing ``SupportsKVCacheGrouping`` owns the grouping + decision; otherwise the default grouping is used. + """ + if not vllm_config.model_config.supports_kv_cache_grouping: + return get_kv_cache_groups(vllm_config, kv_cache_spec) + if vllm_config.attention_config.hisparse_config is not None: + raise ValueError( + "A model-declared get_kv_cache_groups cannot be combined " + "with HiSparse, which owns its own KV cache grouping. " + "Remove one of them." + ) + model_cls, _ = vllm_config.model_config.registry.resolve_model_cls( + vllm_config.model_config.architectures, + vllm_config.model_config, + ) + return model_cls.get_kv_cache_groups(vllm_config, kv_cache_spec) + + def get_kv_cache_config_from_groups( + self, + vllm_config: VllmConfig, + kv_cache_groups: list[KVCacheGroupSpec], + num_blocks: int, + ) -> KVCacheConfig: + """See ``KVCacheConfigBuilder.get_kv_cache_config_from_groups``.""" + if len(kv_cache_groups) == 0: + # Attention free models do not have KV cache. + # Return num_blocks=1 as BlockPool always needs a null_block. + return KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=kv_cache_groups, + prefix_cache_retention_interval=( + vllm_config.cache_config.prefix_cache_retention_interval + ), + ) + bytes_per_block = _pool_bytes_per_block(vllm_config, kv_cache_groups) + return get_kv_cache_config_from_groups( + vllm_config, kv_cache_groups, num_blocks * bytes_per_block + ) diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index a1955b09c7..4a1dd0fbfb 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -1036,26 +1036,22 @@ def _extrapolate_full_graph_memory(mem_samples: list[int], total_graphs: int) -> def _init_minimal_kv_cache_for_profiling(runner: "GPUModelRunner") -> None: """Allocate the smallest KV cache that still lets every graph be captured.""" - from vllm.v1.core.kv_cache_planning import ( - get_kv_cache_config_from_groups, - get_kv_cache_groups, - ) + from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder + from vllm.v1.core.kv_cache_planning import get_kv_cache_groups_from_builder kv_cache_spec = runner.get_kv_cache_spec() - kv_cache_groups = get_kv_cache_groups(runner.vllm_config, kv_cache_spec) + builder = get_kv_cache_config_builder(runner.vllm_config) + kv_cache_groups = get_kv_cache_groups_from_builder( + builder, runner.vllm_config, kv_cache_spec + ) # At least one block per sequence is required to capture the graphs. min_blocks = ( min(runner.max_num_reqs, runner.compilation_config.max_cudagraph_capture_size) or 1 ) - saved_override = runner.cache_config.num_gpu_blocks_override - runner.cache_config.num_gpu_blocks_override = min_blocks - try: - minimal_config = get_kv_cache_config_from_groups( - runner.vllm_config, kv_cache_groups, available_memory=0 - ) - finally: - runner.cache_config.num_gpu_blocks_override = saved_override + minimal_config = builder.get_kv_cache_config_from_groups( + runner.vllm_config, kv_cache_groups, num_blocks=min_blocks + ) runner.initialize_kv_cache(minimal_config, is_profiling=True) runner.cache_config.num_gpu_blocks = minimal_config.num_blocks diff --git a/vllm/v1/worker/gpu_model_runner.py b/vllm/v1/worker/gpu_model_runner.py index f18e219d8e..3f1c3fe4c0 100644 --- a/vllm/v1/worker/gpu_model_runner.py +++ b/vllm/v1/worker/gpu_model_runner.py @@ -6174,27 +6174,23 @@ class GPUModelRunner( gc.collect() def _init_minimal_kv_cache_for_profiling(self) -> None: - from vllm.v1.core.kv_cache_planning import ( - get_kv_cache_config_from_groups, - get_kv_cache_groups, - ) + from vllm.v1.core.kv_cache_config_builder import get_kv_cache_config_builder + from vllm.v1.core.kv_cache_planning import get_kv_cache_groups_from_builder kv_cache_spec = self.get_kv_cache_spec() KVCacheSpecRegistry.check_kv_cache_spec_registry(kv_cache_spec) - kv_cache_groups = get_kv_cache_groups(self.vllm_config, kv_cache_spec) + builder = get_kv_cache_config_builder(self.vllm_config) + kv_cache_groups = get_kv_cache_groups_from_builder( + builder, self.vllm_config, kv_cache_spec + ) # the minimum number of blocks required is 1 block *per sequence* min_blocks = ( min(self.max_num_reqs, self.compilation_config.max_cudagraph_capture_size) or 1 ) - - # Temporarily change num_gpu_blocks_override to allocate a minimal KV cache - saved_override = self.cache_config.num_gpu_blocks_override - self.cache_config.num_gpu_blocks_override = min_blocks - minimal_config = get_kv_cache_config_from_groups( - self.vllm_config, kv_cache_groups, available_memory=0 + minimal_config = builder.get_kv_cache_config_from_groups( + self.vllm_config, kv_cache_groups, num_blocks=min_blocks ) - self.cache_config.num_gpu_blocks_override = saved_override self.initialize_kv_cache(minimal_config, is_profiling=True) self.cache_config.num_gpu_blocks = minimal_config.num_blocks From a94c3157119e52adacb4d580d4e950983d14e353 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Sat, 17 Oct 2026 07:41:54 +0800 Subject: [PATCH 28/49] [DP] Engines own wave wake-ups; drop the front-end FIRST_REQ (#59559) Signed-off-by: fh3 Co-authored-by: fh3 Co-authored-by: Claude Opus 5.5 --- rust/src/engine-core-client/src/client.rs | 7 +- .../src/coordinator/external.rs | 97 ++----------- .../src/coordinator/handle.rs | 91 ++---------- .../src/coordinator/inproc.rs | 108 ++------------- rust/src/engine-core-client/src/error.rs | 4 - .../engine-core-client/src/tests/client.rs | 131 ++++-------------- tests/v1/distributed/test_async_llm_dp.py | 86 +++++++----- tests/v1/engine/test_engine_core.py | 50 ++++++- tests/v1/engine/test_engine_core_client.py | 6 +- .../elastic_ep/external_elastic_ep.py | 2 +- vllm/v1/engine/coordinator.py | 42 ++---- vllm/v1/engine/core.py | 7 +- vllm/v1/engine/core_client.py | 118 ++++++---------- 13 files changed, 226 insertions(+), 523 deletions(-) diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index 49910d80c7..a9bcc01959 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -657,12 +657,9 @@ impl EngineCoreClient { ); let result: Result<()> = async { + // The engine that receives it wakes the other ranks if they are idle. if let Some(coordinator) = self.coordinator.as_ref() { - let snapshot = coordinator.snapshot(); - req.current_wave = snapshot.current_wave; - if !snapshot.engines_running { - coordinator.notify_first_request(engine_id.clone())?; - } + req.current_wave = coordinator.snapshot().current_wave; } debug!( diff --git a/rust/src/engine-core-client/src/coordinator/external.rs b/rust/src/engine-core-client/src/coordinator/external.rs index 606224ea9e..d0418352a6 100644 --- a/rust/src/engine-core-client/src/coordinator/external.rs +++ b/rust/src/engine-core-client/src/coordinator/external.rs @@ -3,34 +3,16 @@ use std::sync::Arc; -use serde_tuple::{Deserialize_tuple, Serialize_tuple}; +use serde_tuple::Deserialize_tuple; use thiserror_ext::AsReport; -use tokio::sync::mpsc; use tracing::{debug, warn}; -use zeromq::prelude::{SocketRecv, SocketSend}; +use zeromq::prelude::SocketRecv; use zeromq::{XSubSocket, ZmqMessage}; use crate::client::imp::ClientInner; -use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; +use crate::coordinator::handle::CoordinatorState; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; -use crate::protocol::{OpaqueValue, decode_msgpack, encode_msgpack}; - -/// Frontend-to-coordinator wakeup message sent when the first request arrives -/// while all engines are paused. -/// -/// This matches the frontend-side msgpack tuple sent by Python -/// `DPAsyncMPClient._ensure_stats_update_task` to the coordinator front socket. -/// -/// Original Python definition: -/// -#[derive(Debug, Clone, PartialEq, Eq, Serialize_tuple, Deserialize_tuple)] -struct CoordinatorWakeupMessage { - /// Engine index that already has the triggering request and should be - /// excluded from the coordinator's `START_DP_WAVE` rebroadcast. - exclude_engine_index: u32, - /// DP wave number observed by the frontend when the request was admitted. - wave: u32, -} +use crate::protocol::{OpaqueValue, decode_msgpack}; /// Coordinator-to-frontend state publish received on the front-side coordinator /// socket. @@ -56,56 +38,17 @@ struct CoordinatorStateUpdate { /// Background half of an external Python-owned coordinator connection. /// -/// This owns the command receiver and one frontend-facing XSUB socket. It -/// mirrors the subset of Python's coordinator protocol needed by the Rust -/// bootstrapped frontend: receive `(counts, wave, running)` publishes, ignore -/// `counts`, and send `(exclude_engine_index, wave)` wakeup messages when the -/// first request arrives while engines are paused. +/// This owns one frontend-facing XSUB socket. It mirrors the subset of Python's +/// coordinator protocol needed by the Rust bootstrapped frontend: receive +/// `(counts, wave, running)` publishes and ignore `counts`. pub(crate) struct ExternalCoordinatorService { state: Arc, - command_rx: mpsc::UnboundedReceiver, socket: XSubSocket, } impl ExternalCoordinatorService { - pub(super) fn new( - state: Arc, - command_rx: mpsc::UnboundedReceiver, - socket: XSubSocket, - ) -> Self { - Self { - state, - command_rx, - socket, - } - } - - /// Apply one frontend-originated command to the external coordinator state - /// machine. - async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> { - match command { - CoordinatorCommand::FirstRequest { - target_engine_id, - wave, - } => { - let target_engine_index = target_engine_id.engine_index().ok_or_else(|| { - Error::UnsupportedCoordinatorEngineId { - engine_id: target_engine_id.to_vec(), - } - })?; - debug!( - wave, - exclude_engine_index = target_engine_index, - "notifying external coordinator about first request while engines were paused" - ); - let payload = encode_msgpack(&CoordinatorWakeupMessage { - exclude_engine_index: target_engine_index, - wave, - })?; - self.socket.send(ZmqMessage::from(payload)).await?; - } - } - Ok(()) + pub(super) fn new(state: Arc, socket: XSubSocket) -> Self { + Self { state, socket } } /// Apply one publish received from the xsub socket containing a coordinator @@ -136,26 +79,14 @@ impl ExternalCoordinatorService { Ok(()) } - /// Drive the coordinator event loop until either side of the control plane - /// is closed or a fatal error is observed. + /// Drive the coordinator event loop until the task is aborted or a fatal + /// error is observed. pub(crate) async fn run(mut self, inner: Arc) { let result: Result<()> = async { loop { - tokio::select! { - // Received frontend-originated command from the handle. - command = self.command_rx.recv() => { - let Some(command) = command else { - warn!("external coordinator command channel closed, shutting down service"); - return Ok(()); - }; - self.handle_command(command).await?; - } - // Received publish from the external coordinator socket. - publish = self.socket.recv() => { - let publish = publish.map_err(Error::from)?; - self.handle_publish(publish).await?; - } - } + // Received publish from the external coordinator socket. + let publish = self.socket.recv().await.map_err(Error::from)?; + self.handle_publish(publish).await?; } } .await; diff --git a/rust/src/engine-core-client/src/coordinator/handle.rs b/rust/src/engine-core-client/src/coordinator/handle.rs index b68c76ddec..9ebc1717e3 100644 --- a/rust/src/engine-core-client/src/coordinator/handle.rs +++ b/rust/src/engine-core-client/src/coordinator/handle.rs @@ -4,101 +4,56 @@ use std::sync::Arc; use parking_lot::Mutex; -use tokio::sync::mpsc; use zeromq::prelude::Socket; use zeromq::{XPubSocket, XSubSocket}; use crate::coordinator::external::ExternalCoordinatorService; use crate::coordinator::inproc::InProcCoordinatorRunner; -use crate::error::{Error, Result, bail_control_closed}; -use crate::transport::EngineId; +use crate::error::Result; /// Snapshot to the coordinator state for request routing and stamping. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub(crate) struct CoordinatorStateSnapshot { /// The current DP wave, which will be stamped on outgoing requests. pub current_wave: u32, - /// Whether the engines are currently running or paused, which determines if - /// the frontend must trigger a new wave on the next request. + /// Whether the engines are currently running or paused, as last reported + /// by the engines themselves. pub engines_running: bool, } -impl CoordinatorStateSnapshot { - /// Resume the engines for a `FirstRequest` and return the wave to broadcast - /// and the engine to exclude from the wakeup. - /// - /// The request may have been stamped with a `request_wave` older than - /// `current_wave` if a `WaveComplete` advanced it after the command was - /// enqueued. Such a request still needs serving, so the current wave is - /// broadcast to every engine (`exclude = None`); the wave is never rewound. - /// A non-stale request excludes the engine that already received it. Mirrors - /// the Python coordinator's front-end path. - pub(crate) fn start_wave_for_first_request( - &mut self, - request_wave: u32, - target_engine_index: u32, - ) -> (u32, Option) { - self.engines_running = true; - let exclude = (request_wave >= self.current_wave).then_some(target_engine_index); - (self.current_wave, exclude) - } -} - /// Shared in-process coordinator state. pub(crate) type CoordinatorState = Mutex; -/// Commands sent from the frontend request path into the background runner. -#[derive(Debug)] -pub(crate) enum CoordinatorCommand { - /// The first request arrived while all engines were paused. - /// - /// The coordinator should broadcast `START_DP_WAVE` with the current wave - /// and the target engine index as the excluded engine. - FirstRequest { - target_engine_id: EngineId, - wave: u32, - }, -} - /// Frontend-facing coordinator handle used by `EngineCoreClient::call()`. /// -/// This side stays intentionally small: it can read the latest wave snapshot -/// and enqueue a `FirstRequest` transition when the request path observes the -/// system in the paused state. +/// The frontend only reads the latest wave snapshot: waves are started by the +/// engines, which notify the coordinator when an idle engine receives a request. #[derive(Clone)] pub(crate) struct CoordinatorHandle { state: Arc, - command_tx: mpsc::UnboundedSender, } impl CoordinatorHandle { - fn new_parts() -> ( - Self, - Arc, - mpsc::UnboundedReceiver, - ) { + fn new_parts() -> (Self, Arc) { let state = Arc::new(Mutex::new(CoordinatorStateSnapshot { current_wave: 0, engines_running: false, })); - let (command_tx, command_rx) = mpsc::unbounded_channel(); ( Self { state: state.clone(), - command_tx, }, state, - command_rx, ) } /// Build the paired frontend handle and background runner around one /// engine-facing coordinator broadcast socket. pub(crate) fn new_inproc(coordinator_input: XPubSocket) -> (Self, InProcCoordinatorRunner) { - let (handle, state, command_rx) = Self::new_parts(); + let (handle, state) = Self::new_parts(); ( handle, - InProcCoordinatorRunner::new(state, command_rx, coordinator_input), + InProcCoordinatorRunner::new(state, coordinator_input), ) } @@ -107,41 +62,15 @@ impl CoordinatorHandle { pub(crate) async fn connect_external( coordinator_address: &str, ) -> Result<(Self, ExternalCoordinatorService)> { - let (handle, state, command_rx) = Self::new_parts(); + let (handle, state) = Self::new_parts(); let mut socket = XSubSocket::new(); socket.connect(coordinator_address).await?; socket.subscribe("").await?; - Ok(( - handle, - ExternalCoordinatorService::new(state, command_rx, socket), - )) + Ok((handle, ExternalCoordinatorService::new(state, socket))) } /// Snapshot the coordinator state for request routing and stamping. pub(crate) fn snapshot(&self) -> CoordinatorStateSnapshot { *self.state.lock() } - - /// Notify the runner that a new request arrived while engines were paused. - /// - /// The handle flips `engines_running` optimistically so concurrent request - /// submissions coalesce behind one `START_DP_WAVE` broadcast instead of all - /// trying to trigger the wave independently. - pub(crate) fn notify_first_request(&self, target_engine_id: EngineId) -> Result<()> { - let mut state = self.state.lock(); - if state.engines_running { - return Ok(()); - } - - let command = CoordinatorCommand::FirstRequest { - target_engine_id, - wave: state.current_wave, - }; - if self.command_tx.send(command).is_err() { - bail_control_closed!("in-process coordinator command channel already shut down"); - } - - state.engines_running = true; - Ok(()) - } } diff --git a/rust/src/engine-core-client/src/coordinator/inproc.rs b/rust/src/engine-core-client/src/coordinator/inproc.rs index cbd69095f8..d780b64e99 100644 --- a/rust/src/engine-core-client/src/coordinator/inproc.rs +++ b/rust/src/engine-core-client/src/coordinator/inproc.rs @@ -12,7 +12,7 @@ use zeromq::{XPubSocket, ZmqError, ZmqMessage}; use crate::client::imp::ClientInner; use crate::coordinator::bootstrap::is_engine_subscription; -use crate::coordinator::handle::{CoordinatorCommand, CoordinatorState}; +use crate::coordinator::handle::CoordinatorState; use crate::error::{Error, Result, bail_unexpected_coordinator_output}; use crate::protocol::encode_msgpack; use crate::protocol::output::{DpControlMessage, DpControlOutput, EngineCoreOutputs}; @@ -30,32 +30,25 @@ use crate::protocol::request::EngineCoreRequestType; struct StartDpWaveMessage { /// DP wave number that all engines should start processing. wave: u32, - /// Engine index that already received the triggering request and so does not - /// need an extra wakeup. `None` wakes every engine (used when the triggering - /// request was for a stale wave). + /// Engine index that started the wave and so does not need an extra wakeup. + /// `None` wakes every engine. exclude_engine_index: Option, } /// Background half of the in-process coordinator. /// -/// This owns the command receiver and the engine-facing coordinator input -/// socket. It is the single place where wave transitions are serialized and -/// where `START_DP_WAVE` broadcasts are emitted. +/// This owns the engine-facing coordinator input socket. It is the single place +/// where wave transitions are serialized and where `START_DP_WAVE` broadcasts +/// are emitted. pub(crate) struct InProcCoordinatorRunner { state: Arc, - command_rx: mpsc::UnboundedReceiver, coordinator_input: XPubSocket, } impl InProcCoordinatorRunner { - pub(super) fn new( - state: Arc, - command_rx: mpsc::UnboundedReceiver, - coordinator_input: XPubSocket, - ) -> Self { + pub(super) fn new(state: Arc, coordinator_input: XPubSocket) -> Self { Self { state, - command_rx, coordinator_input, } } @@ -82,34 +75,6 @@ impl InProcCoordinatorRunner { Ok(()) } - /// Apply one frontend-originated command to the coordinator state machine. - async fn handle_command(&mut self, command: CoordinatorCommand) -> Result<()> { - match command { - CoordinatorCommand::FirstRequest { - target_engine_id, - wave, - } => { - let target_engine_index = target_engine_id.engine_index().ok_or_else(|| { - Error::UnsupportedCoordinatorEngineId { - engine_id: target_engine_id.to_vec(), - } - })?; - let (current_wave, exclude) = { - let mut state = self.state.lock(); - state.start_wave_for_first_request(wave, target_engine_index) - }; - debug!( - current_wave, - request_wave = wave, - ?exclude, - "starting DP wave after first request while engines were paused" - ); - self.broadcast_start_wave(current_wave, exclude).await?; - } - } - Ok(()) - } - /// Apply one engine-originated control output to the coordinator state /// machine. async fn handle_outputs(&mut self, outputs: EngineCoreOutputs) -> Result<()> { @@ -160,7 +125,7 @@ impl InProcCoordinatorRunner { debug!( wave, exclude_engine_index = engine_index, - "starting DP wave after stale-wave notification from engine" + "starting DP wave on notification from engine" ); self.broadcast_start_wave(wave, Some(engine_index)).await?; } @@ -175,8 +140,8 @@ impl InProcCoordinatorRunner { Ok(()) } - /// Drive the coordinator event loop until either side of the control plane - /// is closed or a fatal error is observed. + /// Drive the coordinator event loop until the engine control output is + /// closed or a fatal error is observed. /// /// Any fatal error closes the main client registries so request streams and /// future calls observe a stable shutdown cause. @@ -216,14 +181,6 @@ impl InProcCoordinatorRunner { Err(error) => return Err(error.into()), } } - // Received frontend-originated command from the handle. - command = self.command_rx.recv() => { - let Some(command) = command else { - warn!("coordinator command channel closed, shutting down coordinator runner"); - return Ok(()); - }; - self.handle_command(command).await?; - } // Received engine-originated control output from the coordinator socket. outputs = output_rx.recv() => { let Some(outputs) = outputs else { @@ -242,48 +199,3 @@ impl InProcCoordinatorRunner { inner.close_registries(Arc::new(error)); } } - -#[cfg(test)] -mod tests { - use crate::coordinator::handle::CoordinatorStateSnapshot; - - /// A `FirstRequest` for the current wave starts that wave and excludes the - /// engine that already received the triggering request. - #[test] - fn first_request_for_current_wave_excludes_target() { - let mut state = CoordinatorStateSnapshot { - current_wave: 3, - engines_running: false, - }; - - let (wave, exclude) = state.start_wave_for_first_request(3, 2); - - assert_eq!(wave, 3); - assert_eq!(exclude, Some(2)); - assert!(state.engines_running); - assert_eq!(state.current_wave, 3); - } - - /// A `FirstRequest` whose wave was superseded by a racing `WaveComplete` - /// (`request_wave < current_wave`) must still start the request's wave: it - /// broadcasts the current wave and wakes every engine (`exclude = None`) - /// rather than rewinding the wave or dropping the request. - #[test] - fn stale_first_request_starts_current_wave_for_all_engines() { - let mut state = CoordinatorStateSnapshot { - current_wave: 4, - engines_running: false, - }; - - // Request stamped with wave 3 while the coordinator already advanced to 4. - let (wave, exclude) = state.start_wave_for_first_request(3, 2); - - assert_eq!( - wave, 4, - "must broadcast the current wave, not the stale one" - ); - assert_eq!(exclude, None, "a stale request must wake every engine"); - assert!(state.engines_running); - assert_eq!(state.current_wave, 4, "wave must not be rewound"); - } -} diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index d55f2cbbcb..6217208e3b 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -55,8 +55,6 @@ pub enum Error { UnexpectedCoordinatorOutput { message: String }, #[error("unexpected output on main dispatcher path: {message}")] UnexpectedDispatcherOutput { message: String }, - #[error("coordinator requires a Python-compatible two-byte engine id, got {engine_id:?}")] - UnsupportedCoordinatorEngineId { engine_id: Vec }, #[error("unsupported auxiliary frame(s): expected 1 frame, got {frame_count}")] UnsupportedAuxFrames { frame_count: usize }, #[error("external coordinator mode is not implemented yet")] @@ -66,8 +64,6 @@ pub enum Error { context: &'static str, field: &'static str, }, - #[error("engine control channel closed unexpectedly: {message}")] - ControlClosed { message: String }, #[error("request `{request_id}` is already in flight")] DuplicateRequestId { request_id: String }, #[error( diff --git a/rust/src/engine-core-client/src/tests/client.rs b/rust/src/engine-core-client/src/tests/client.rs index 7b7a880bde..c2216168b8 100644 --- a/rust/src/engine-core-client/src/tests/client.rs +++ b/rust/src/engine-core-client/src/tests/client.rs @@ -267,6 +267,19 @@ async fn recv_start_dp_wave(sub: &mut SubSocket) -> (u32, Option) { rmp_serde::from_slice(&frames[1]).expect("decode START_DP_WAVE payload") } +async fn send_start_wave(push: &mut PushSocket, engine_index: u32, wave: u32) { + send_outputs( + push, + DpControlOutput { + engine_index, + timestamp: 0.0, + control: DpControlMessage::StartWave(wave), + } + .into(), + ) + .await; +} + async fn connect_client_with_ipc( config: EngineCoreClientConfig, ipc: &IpcNamespace, @@ -413,12 +426,6 @@ async fn recv_xpub_subscription(xpub: &mut XPubSocket) { assert_eq!(frames[0].as_ref(), b"\x01"); } -async fn recv_external_coordinator_wakeup(xpub: &mut XPubSocket) -> (u32, u32) { - let frames = recv_xpub_message(xpub).await; - assert_eq!(frames.len(), 1); - rmp_serde::from_slice(&frames[0]).expect("decode external coordinator wakeup") -} - async fn send_external_coordinator_publish( xpub: &mut XPubSocket, payload: &T, @@ -617,6 +624,14 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); + // An idle engine handed a request starts the wave itself. + let add = recv_engine_message(&mut data_socket.dealer).await; + assert_eq!(add[0].as_ref(), &[0x00]); + let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); + assert_eq!(request.request_id, "req-1"); + assert_eq!(request.current_wave, 0); + send_start_wave(&mut coordinator.output_push, 0, 0).await; + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!((wave, exclude_engine), (0, Some(0))); @@ -629,12 +644,6 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { let replay = recv_start_dp_wave(&mut coordinator.input_sub).await; assert_eq!(replay, (0, None)); - let add = recv_engine_message(&mut data_socket.dealer).await; - assert_eq!(add[0].as_ref(), &[0x00]); - let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); - assert_eq!(request.request_id, "req-1"); - assert_eq!(request.current_wave, 0); - assert!( timeout( Duration::from_millis(200), @@ -670,14 +679,15 @@ async fn coordinator_wave_control_tracks_pause_running_and_rebroadcasts() { ) .await; - let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (1, Some(0))); - let add = recv_engine_message(&mut data_socket.dealer).await; assert_eq!(add[0].as_ref(), &[0x00]); let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); assert_eq!(request.request_id, "req-3"); assert_eq!(request.current_wave, 1); + send_start_wave(&mut coordinator.output_push, 0, 1).await; + + let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; + assert_eq!((wave, exclude_engine), (1, Some(0))); send_outputs( &mut data_socket.push, @@ -896,9 +906,6 @@ async fn coordinator_accepts_stats_only_outputs() { engine.coordinator.take().expect("coordinator sockets should be present"); let data_socket = engine.data_sockets.first_mut().expect("data socket"); - let (wave, exclude_engine) = recv_start_dp_wave(&mut coordinator.input_sub).await; - assert_eq!((wave, exclude_engine), (0, Some(0))); - send_outputs( &mut coordinator.output_push, RequestBatchOutputs { @@ -3248,7 +3255,7 @@ async fn bootstrapped_external_coordinator_connects_and_subscribes() { } #[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_one_wakeup() { +async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends_no_wakeup() { init_tracing(); let ipc = IpcNamespace::new().unwrap(); let input_address = ipc.input_endpoint(); @@ -3289,18 +3296,11 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); - let wakeup = timeout( - Duration::from_secs(1), - recv_external_coordinator_wakeup(&mut stats_socket), - ) - .await - .unwrap(); - assert_eq!(wakeup, (0, 7)); - + // The engine that receives the request starts the wave, not the frontend. assert!( timeout( Duration::from_millis(200), - recv_external_coordinator_wakeup(&mut stats_socket) + recv_xpub_message(&mut stats_socket) ) .await .is_err() @@ -3333,78 +3333,3 @@ async fn bootstrapped_external_coordinator_updates_wave_ignores_counts_and_sends client.shutdown().await.unwrap(); } - -#[tokio::test(flavor = "multi_thread", worker_threads = 2)] -async fn bootstrapped_external_coordinator_running_state_suppresses_wakeup() { - init_tracing(); - let ipc = IpcNamespace::new().unwrap(); - let input_address = ipc.input_endpoint(); - let output_address = ipc.output_endpoint(); - let coordinator_address = ipc.endpoint("stats.sock"); - - let mut stats_socket = XPubSocket::new(); - stats_socket.bind(&coordinator_address).await.unwrap(); - - let client_task = tokio::spawn({ - let input_address = input_address.clone(); - let output_address = output_address.clone(); - let coordinator_address = coordinator_address.clone(); - async move { - EngineCoreClient::connect(bootstrapped_test_config( - input_address, - output_address, - 1, - Duration::from_secs(2), - 0, - Some(CoordinatorMode::External { - address: coordinator_address, - }), - )) - .await - .unwrap() - } - }); - - let (mut dealer, mut push) = - setup_bootstrapped_mock_engine(input_address, output_address, &[0x00, 0x00]).await; - let client = client_task.await.unwrap(); - recv_xpub_subscription(&mut stats_socket).await; - - send_external_coordinator_publish(&mut stats_socket, &(Value::Nil, 5_u32, true)).await; - tokio::time::sleep(Duration::from_millis(50)).await; - - let mut stream = client.call(sample_request()).await.unwrap().into_outputs(); - - assert!( - timeout( - Duration::from_millis(200), - recv_external_coordinator_wakeup(&mut stats_socket) - ) - .await - .is_err() - ); - - let add = recv_engine_message(&mut dealer).await; - let request: EngineCoreRequest = rmp_serde::from_slice(&add[1]).unwrap(); - assert_eq!(request.current_wave, 5); - - send_outputs( - &mut push, - RequestBatchOutputs { - outputs: vec![request_output( - "req-1", - vec![], - Some(EngineCoreFinishReason::Length), - )], - finished_requests: Some(BTreeSet::from(["req-1".to_string()])), - ..Default::default() - } - .into(), - ) - .await; - - let final_output = timeout(Duration::from_secs(1), stream.next()).await.unwrap(); - assert!(final_output.is_some()); - - client.shutdown().await.unwrap(); -} diff --git a/tests/v1/distributed/test_async_llm_dp.py b/tests/v1/distributed/test_async_llm_dp.py index 7b350fc55e..d5beea2f63 100644 --- a/tests/v1/distributed/test_async_llm_dp.py +++ b/tests/v1/distributed/test_async_llm_dp.py @@ -18,6 +18,7 @@ from vllm.inputs import PromptType from vllm.outputs import RequestOutput from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind +from vllm.v1.engine import EngineCoreRequestType from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.core_client import DPLBAsyncMPClient from vllm.v1.metrics.loggers import StatLoggerBase @@ -336,12 +337,12 @@ async def _poll_flag(engine: AsyncLLM, want: bool, timeout: float) -> bool: @pytest.mark.asyncio async def test_dp_pause_late_request_does_not_block_drain(): - """A request arriving after pause must not leave the coordinator believing + """A request arriving after pause must not make the coordinator believe the engines are running. - Paused engines discard START_DP_WAVE, so nothing can report wave - completion afterwards; if forwarding the wake also marks the engines as - running, drain has no way back and can only end in a timeout. + Paused engines start no wave for it, so nothing could report wave + completion afterwards; a running flag set on its behalf would leave drain + no way back but a timeout. MoE only: wave coordination is enabled iff the model is MoE, so a dense model never reaches the coordinator state this exercises. @@ -365,11 +366,9 @@ async def test_dp_pause_late_request_does_not_block_drain(): # A design that never reports the engines as running would satisfy # every drain assertion below while destroying the signal, so pin it - # down first. The front-end sets its own copy optimistically when it - # forwards the wake, so sample only after several coordinator - # publishes (every 100ms while stats change) have overwritten it. - # Rank 0 reports the wave only after its first step, which can outlast - # the 2s on a cold device, so wait for the report rather than sample it. + # down first. Rank 0 reports the wave only after its first step, which + # can outlast the 2s on a cold device, so wait for the report rather + # than sample it. await asyncio.sleep(2) assert not long_request.done(), "the warmup request was too short to sample" assert await _poll_flag(engine, True, timeout=30), ( @@ -382,24 +381,12 @@ async def test_dp_pause_late_request_does_not_block_drain(): await engine.pause_generation(mode="abort") await engine.wait_for_requests_to_drain(drain_timeout=30) - # Awaiting add_request guarantees the new-request notification has been - # sent to the coordinator - the message that used to latch the flag. collector = await engine.add_request( request_id="late", prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=5), ) - - # The front-end marks the engines running off the back of that - # notification. This is what makes the test non-vacuous: it is the - # path that used to leave the coordinator stuck. - assert await _poll_flag(engine, True, timeout=5), ( - "the late request did not notify the coordinator" - ) - - # It must settle back by itself. Unfixed it never does, because the - # paused engines discard the wake and so never report wave completion. - assert await _poll_flag(engine, False, timeout=60) + assert not await _poll_flag(engine, True, timeout=5) # The late request was held rather than dropped: it completes on resume. await engine.resume_generation() @@ -409,6 +396,46 @@ async def test_dp_pause_late_request_does_not_block_drain(): break +@pytest.mark.asyncio +async def test_dp_request_reaching_rank_after_resume_wakes_peers(): + """A request can reach its rank only after resume (e.g. from another API + server), with no other work anywhere. That rank must wake its peers itself, + or it steps alone. Modelled here by holding the request's ADD on the client + side across pause and resume.""" + with ExitStack() as after: + engine = AsyncLLM.from_engine_args( + _get_dp_pause_engine_args(expert_parallel=True) + ) + after.callback(engine.shutdown) + client = engine.engine_core + params = SamplingParams(max_tokens=5) + + await _consume(engine.generate(DP_PAUSE_PROMPT, params, "warmup")) + await engine.pause_generation(mode="abort") + assert await _poll_flag(engine, False, timeout=30) + + held: list[tuple] = [] + send_input = client._send_input + + def hold_add(request_type, request, engine_id=None): + if request_type == EngineCoreRequestType.ADD: + held.append((request_type, request, engine_id)) + return asyncio.sleep(0) + return send_input(request_type, request, engine_id) + + client._send_input = hold_add + task = asyncio.create_task( + _consume(engine.generate(DP_PAUSE_PROMPT, params, "late")) + ) + await asyncio.sleep(2) + assert held + + await engine.resume_generation() + client._send_input = send_input + await send_input(*held[0]) + await asyncio.wait_for(task, timeout=60) + + @pytest.mark.asyncio async def test_dp_sleep_late_request_does_not_block_drain(): """The same latch, reached through sleep rather than pause. @@ -441,14 +468,7 @@ async def test_dp_sleep_late_request_does_not_block_drain(): prompt=DP_PAUSE_PROMPT, params=SamplingParams(max_tokens=5), ) - - assert await _poll_flag(engine, True, timeout=5), ( - "the request did not notify the coordinator" - ) - - # Sleeping engines cannot report wave completion, so a coordinator that - # marked them running when it forwarded the wake never hears otherwise. - assert await _poll_flag(engine, False, timeout=60) + assert not await _poll_flag(engine, True, timeout=5) await engine.wake_up() assert not await engine.is_sleeping() @@ -628,7 +648,7 @@ async def test_dp_pause_keep_race_staggered_engines(): @pytest.mark.asyncio async def test_dp_pause_barrier_request_deadlock(): - """Test that start_dp_wave is ignored while paused. + """Test that a request sent while paused does not wake any engine. Sequence: 1. Pause all engines (PAUSED_ALL). @@ -697,7 +717,7 @@ async def test_dp_pause_barrier_request_deadlock(): # Yield so generate() preprocessing completes and # add_request_async is called (which, in buggy code, - # would send FIRST_REQ and wake engine 1). + # would wake engine 1). for _ in range(200): await asyncio.sleep(0) @@ -725,7 +745,7 @@ async def test_dp_pause_barrier_request_deadlock(): for t in mid_barrier_tasks: t.cancel() pytest.fail( - "Staggered barrier deadlocked — FIRST_REQ sent while " + "Staggered barrier deadlocked — a request sent while " "paused caused collective-op mismatch between engines" ) diff --git a/tests/v1/engine/test_engine_core.py b/tests/v1/engine/test_engine_core.py index 2d9f7e8632..df65b9f481 100644 --- a/tests/v1/engine/test_engine_core.py +++ b/tests/v1/engine/test_engine_core.py @@ -25,7 +25,7 @@ from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform from vllm.utils.torch_utils import set_default_torch_num_threads from vllm.v1.core.sched.interface import PauseState -from vllm.v1.engine import EngineCoreRequest +from vllm.v1.engine import EngineCoreRequest, EngineCoreRequestType from vllm.v1.engine.core import DPEngineCoreProc, EngineCore, EngineCoreProc from vllm.v1.executor.abstract import Executor from vllm.v1.executor.uniproc_executor import UniProcExecutor @@ -689,6 +689,54 @@ def test_dp_sync_interval_idle_pause_consensus_on_first_step(monkeypatch): assert not core.pending_pause +@pytest.mark.parametrize( + "request_wave,engines_running,pause_state,announces", + [ + pytest.param(3, False, PauseState.UNPAUSED, True, id="current-wave"), + pytest.param(2, False, PauseState.UNPAUSED, True, id="stale-wave"), + pytest.param(4, False, PauseState.UNPAUSED, True, id="newer-wave"), + pytest.param(3, True, PauseState.UNPAUSED, False, id="running"), + pytest.param(3, False, PauseState.PAUSED_ALL, False, id="paused"), + ], +) +def test_idle_dp_rank_announces_wave_for_new_work( + request_wave, engines_running, pause_state, announces +): + """An idle rank handed work announces the wave itself, since only engines + start waves; a paused rank waits for resume.""" + core = object.__new__(DPEngineCoreProc) + core.has_coordinator = True + core.current_wave = 3 + core.engines_running = engines_running + core.scheduler = MagicMock(pause_state=pause_state) + core.output_queue = MagicMock() + + with patch.object(EngineCore, "add_request"): + core.add_request(MagicMock(), request_wave) + + if announces: + assert core.engines_running + _, outputs = core.output_queue.put_nowait.call_args.args[0] + assert outputs.start_wave == max(request_wave, 3) + else: + core.output_queue.put_nowait.assert_not_called() + + +@pytest.mark.parametrize("paused", [False, True]) +def test_paused_dp_rank_ignores_start_wave(paused): + """A START_DP_WAVE still in flight when the pause completes must not wake + the paused rank.""" + core = object.__new__(DPEngineCoreProc) + core.ignore_start_dp_wave = paused + core.engine_index = 0 + core.current_wave = 3 + core.engines_running = False + + core._handle_client_request(EngineCoreRequestType.START_DP_WAVE, (3, 1)) + + assert core.engines_running != paused + + def _pausable_engine_core_proc() -> EngineCoreProc: """A bare EngineCoreProc holding just the state pause_scheduler touches.""" core = object.__new__(EngineCoreProc) diff --git a/tests/v1/engine/test_engine_core_client.py b/tests/v1/engine/test_engine_core_client.py index b58aed8547..e2804e9aa5 100644 --- a/tests/v1/engine/test_engine_core_client.py +++ b/tests/v1/engine/test_engine_core_client.py @@ -314,7 +314,7 @@ async def test_dplb_scale_down_routes_after_stale_stats_snapshot(): client.current_wave = 0 client.resources = SimpleNamespace(stats_update_task=None) client.stats_update_address = "inproc://scale-down-stats" - client.first_req_sock_addr = "inproc://scale-down-first-request" + client.scale_sock_addr = "inproc://scale-down-notify" pause_started = asyncio.Event() async def pause_scheduler(*args, **kwargs): @@ -326,11 +326,11 @@ async def test_dplb_scale_down_routes_after_stale_stats_snapshot(): with ( zmq.asyncio.Context() as ctx, ctx.socket(zmq.XPUB) as coordinator, - ctx.socket(zmq.PAIR) as first_req_socket, + ctx.socket(zmq.PAIR) as scale_socket, ): client.ctx = ctx coordinator.bind(client.stats_update_address) - first_req_socket.bind(client.first_req_sock_addr) + scale_socket.bind(client.scale_sock_addr) client._ensure_stats_update_task() scale_task = None try: diff --git a/vllm/distributed/elastic_ep/external_elastic_ep.py b/vllm/distributed/elastic_ep/external_elastic_ep.py index 6db6e5f4b1..66d94b4371 100644 --- a/vllm/distributed/elastic_ep/external_elastic_ep.py +++ b/vllm/distributed/elastic_ep/external_elastic_ep.py @@ -838,7 +838,7 @@ class ExternalElasticEPScaleCoordinator: if prepared.dp_rank == 0: self.client._coord_store = self.reconfig_store_ref self.client._ensure_stats_update_task() - await self.client.first_req_send_socket.send( + await self.client.scale_send_socket.send( msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", bootstrap.new_data_parallel_size) ) diff --git a/vllm/v1/engine/coordinator.py b/vllm/v1/engine/coordinator.py index 28a4c7af6d..524794dac0 100644 --- a/vllm/v1/engine/coordinator.py +++ b/vllm/v1/engine/coordinator.py @@ -41,12 +41,9 @@ class DPCoordinator: DPEngineCoreProc._has_global_unfinished_reqs method. * Broadcasts the START_DP_WAVE message to engines to move them from paused - to running state when one engine receives a new request. This can happen - in two cases: - 1) A front-end sending a new request while the engines are paused will - concurrently notify the coordinator. - 2) An engine receiving a request for a stale request wave while in paused - state will notify the coordinator. + to running state when an idle engine receives a new request and notifies + the coordinator. Only the engines start waves: a front-end's view of the + running state can be stale, and a paused engine discards START_DP_WAVE. Engines will move into running state when receiving a new request or START_DP_WAVE message. @@ -353,27 +350,9 @@ class DPCoordinatorProc: ) continue # Skip normal engine notification processing - # Wave coordination: handle new-request messages from front-end. - # Only process these when wave coordination is enabled - if self.enable_wave_coordination: - # We received a message on the front-end XPUB socket, - # from an API server sending a new request while the - # engines are paused, so that we can wake the other - # engines. - engine_to_exclude, wave = decoded - if not engines_running: - if wave < current_wave: - # If the wave number is stale, ensure the message - # is handled by all the engines. - engine_to_exclude = None - - # engines_running is only set from the engines' - # own notifications; a paused engine discards - # START_DP_WAVE, so sending it is not evidence - # that the engines are running. - self._send_start_wave( - publish_back, current_wave, engine_to_exclude - ) + logger.error( + "DP Coordinator received unexpected message from front-end" + ) if output_back in events: # We received a message from one of the engines. @@ -448,13 +427,12 @@ class DPCoordinatorProc: wave > current_wave or (wave == current_wave and not engines_running) ): - # 3. The engine received request for a non-current wave - # so we must ensure that other engines progress to the - # next wave (race condition handling). + # 3. An engine started a wave (e.g. an idle engine + # received a request), so the others must join it. logger.debug( - "Starting wave %d after notification of " - "stale wave request from engine.", + "Starting wave %d on notification from engine %d.", wave, + eng_index, ) current_wave = wave engines_running = True diff --git a/vllm/v1/engine/core.py b/vllm/v1/engine/core.py index c9843dce96..2b275b65c5 100644 --- a/vllm/v1/engine/core.py +++ b/vllm/v1/engine/core.py @@ -2269,15 +2269,14 @@ class DPEngineCoreProc(EngineCoreProc): def add_request(self, request: Request, request_wave: int = 0): super().add_request(request, request_wave) - if self.has_coordinator and request_wave != self.current_wave: + if self.has_coordinator: if request_wave > self.current_wave: self.current_wave = request_wave - elif ( + if ( not self.engines_running and self.scheduler.pause_state == PauseState.UNPAUSED ): - # Request received for an already-completed wave, notify - # front-end that we need to start the next one. + # Engines alone start waves: have the coordinator wake the others. self.engines_running = True self.output_queue.put_nowait( (-1, EngineCoreOutputs(start_wave=self.current_wave)) diff --git a/vllm/v1/engine/core_client.py b/vllm/v1/engine/core_client.py index cfe8976668..0c87c9a052 100644 --- a/vllm/v1/engine/core_client.py +++ b/vllm/v1/engine/core_client.py @@ -506,8 +506,8 @@ class BackgroundResources: coordinator: DPCoordinator | None = None output_socket: zmq.Socket | zmq.asyncio.Socket | None = None input_socket: zmq.Socket | zmq.asyncio.Socket | None = None - first_req_send_socket: zmq.asyncio.Socket | None = None - first_req_rcv_socket: zmq.asyncio.Socket | None = None + scale_send_socket: zmq.asyncio.Socket | None = None + scale_rcv_socket: zmq.asyncio.Socket | None = None stats_update_socket: zmq.asyncio.Socket | None = None output_queue_task: asyncio.Task | None = None stats_update_task: asyncio.Task | None = None @@ -535,8 +535,8 @@ class BackgroundResources: sockets = ( self.output_socket, self.input_socket, - self.first_req_send_socket, - self.first_req_rcv_socket, + self.scale_send_socket, + self.scale_rcv_socket, self.stats_update_socket, ) @@ -1493,9 +1493,9 @@ class DPAsyncMPClient(AsyncMPClient): self.eep_scaling_cache: ElasticScalingCache | None = None - self.first_req_sock_addr = get_open_zmq_inproc_path() - self.first_req_send_socket = self.resources.first_req_send_socket = ( - make_zmq_socket(self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=True) + self.scale_sock_addr = get_open_zmq_inproc_path() + self.scale_send_socket = self.resources.scale_send_socket = make_zmq_socket( + self.ctx, self.scale_sock_addr, zmq.PAIR, bind=True ) try: # If we are running in an asyncio event loop, start the stats task. @@ -1518,75 +1518,48 @@ class DPAsyncMPClient(AsyncMPClient): with ( make_zmq_socket(self.ctx, stats_addr, zmq.XSUB, linger=0) as socket, make_zmq_socket( - self.ctx, self.first_req_sock_addr, zmq.PAIR, bind=False, linger=0 - ) as first_req_rcv_socket, + self.ctx, self.scale_sock_addr, zmq.PAIR, bind=False, linger=0 + ) as scale_rcv_socket, ): assert isinstance(socket, zmq.asyncio.Socket) - assert isinstance(first_req_rcv_socket, zmq.asyncio.Socket) + assert isinstance(scale_rcv_socket, zmq.asyncio.Socket) self.resources.stats_update_socket = socket - self.resources.first_req_rcv_socket = first_req_rcv_socket + self.resources.scale_rcv_socket = scale_rcv_socket # Send subscription message. await socket.send(b"\x01") poller = zmq.asyncio.Poller() poller.register(socket, zmq.POLLIN) - poller.register(first_req_rcv_socket, zmq.POLLIN) + poller.register(scale_rcv_socket, zmq.POLLIN) while True: - events = await poller.poll() - if ( - not self.engines_running - and len(events) == 2 - or (events[0][0] == first_req_rcv_socket) - ): - # Check if this is a regular request notification or - # scale up notification - buf = first_req_rcv_socket.recv(flags=zmq.NOBLOCK).result() - - decoded = msgspec.msgpack.decode(buf) - if ( - isinstance(decoded, (list, tuple)) - and len(decoded) == 2 - and decoded[0] == "SCALE_ELASTIC_EP" - ): - # Extract new engine count from the decoded message - new_engine_count = decoded[1] - # Update engine_ranks_managed and count_slice - parallel_config = self.vllm_config.parallel_config - dp_size = parallel_config.data_parallel_size - dp_rank = parallel_config.data_parallel_rank - assert dp_rank == 0 - assert dp_size == new_engine_count - assert not parallel_config.data_parallel_hybrid_lb - # External LB clients continue to manage only local ranks. - if not parallel_config.data_parallel_external_lb: - self.engine_ranks_managed = list(range(dp_size)) - if len(self.lb_engines) < new_engine_count: - self.lb_engines += [ - [0, 0, 0.0] - for _ in range( - new_engine_count - len(self.lb_engines) - ) - ] - else: - self.lb_engines = self.lb_engines[:new_engine_count] - # Send scale up notification to coordinator - scale_msg = msgspec.msgpack.encode( - ("SCALE_ELASTIC_EP", new_engine_count) - ) - await socket.send(scale_msg) - continue - - # we're sending a request while the engines are - # paused, so that it can wake the others up - # (to run dummy EP loop). - assert decoded[0] == "FIRST_REQ" - target_eng_index = decoded[1] - self.engines_running = True - msg = msgspec.msgpack.encode( - (target_eng_index, self.current_wave) - ) - await socket.send(msg) + events = dict(await poller.poll()) + if scale_rcv_socket in events: + buf = scale_rcv_socket.recv(flags=zmq.NOBLOCK).result() + msg_type, new_engine_count = msgspec.msgpack.decode(buf) + assert msg_type == "SCALE_ELASTIC_EP" + # Update engine_ranks_managed and count_slice + parallel_config = self.vllm_config.parallel_config + dp_size = parallel_config.data_parallel_size + dp_rank = parallel_config.data_parallel_rank + assert dp_rank == 0 + assert dp_size == new_engine_count + assert not parallel_config.data_parallel_hybrid_lb + # External LB clients continue to manage only local ranks. + if not parallel_config.data_parallel_external_lb: + self.engine_ranks_managed = list(range(dp_size)) + if len(self.lb_engines) < new_engine_count: + self.lb_engines += [ + [0, 0, 0.0] + for _ in range( + new_engine_count - len(self.lb_engines) + ) + ] + else: + self.lb_engines = self.lb_engines[:new_engine_count] + # Forward the scale notification to the coordinator. + await socket.send(buf) + continue buf = None while True: @@ -1624,14 +1597,9 @@ class DPAsyncMPClient(AsyncMPClient): request.current_wave = self.current_wave request.client_index = self.client_index + # The engine that receives it wakes the other ranks if they are idle. chosen_engine = self.get_core_engine_for_request(request) - to_await = self._send_input(EngineCoreRequestType.ADD, request, chosen_engine) - if not self.engines_running: - # Notify coordinator that we're sending a request - req_msg = msgspec.msgpack.encode(("FIRST_REQ", chosen_engine)) - await self.first_req_send_socket.send(req_msg) - - await to_await + await self._send_input(EngineCoreRequestType.ADD, request, chosen_engine) self._ensure_output_queue_task() @@ -2100,7 +2068,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): scale_up_marker = msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", new_data_parallel_size) ) - await self.first_req_send_socket.send(scale_up_marker) + await self.scale_send_socket.send(scale_up_marker) logger.info( "[Elastic EP] Scale up completed, new data parallel size: %s", @@ -2176,7 +2144,7 @@ class DPLBAsyncMPClient(DPAsyncMPClient): scale_down_marker = msgspec.msgpack.encode( ("SCALE_ELASTIC_EP", new_data_parallel_size) ) - await self.first_req_send_socket.send(scale_down_marker) + await self.scale_send_socket.send(scale_down_marker) await wait_future await self.resume_scheduler_async() except Exception: From 523d74e2ac78867aa35d535835e14f78fb7433af Mon Sep 17 00:00:00 2001 From: Beichen Huang <71347503+BeichenHuang@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:42:30 -0500 Subject: [PATCH 29/49] =?UTF-8?q?[Bugfix][LoRA]=20Fix=20garbage/NaN=20lora?= =?UTF-8?q?=5Fexpand=20output=20on=20SM90=20with=20fp32=20i=E2=80=A6=20(#6?= =?UTF-8?q?1906)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: BeichenHuang Co-authored-by: Claude Opus 5.5 Co-authored-by: aoshen02 --- vllm/lora/ops/triton_ops/lora_expand_op.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/vllm/lora/ops/triton_ops/lora_expand_op.py b/vllm/lora/ops/triton_ops/lora_expand_op.py index 184e58247b..05fea52a71 100644 --- a/vllm/lora/ops/triton_ops/lora_expand_op.py +++ b/vllm/lora/ops/triton_ops/lora_expand_op.py @@ -15,6 +15,7 @@ from vllm.lora.ops.triton_ops.utils import ( get_lora_op_configs, supports_pdl, ) +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton from vllm.utils.torch_utils import direct_register_custom_op @@ -231,6 +232,10 @@ def _lora_expand( torch.bfloat16, ]: CAST_TYPE = True + # Pipelined wgmma with a register A operand is miscompiled by + # ptxas < 13.3 on SM90 (triton-lang/triton#9433). + if current_platform.is_cuda() and current_platform.is_device_capability(90): + NUM_STAGES = 1 # TODO (varun): This grid formulation maximizes parallelization at the # cost of wasteful thread block launch when only a few input tokens require From be25977532f94db389f8e43776c946615b73f4d7 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Fri, 16 Oct 2026 20:10:08 -0400 Subject: [PATCH 30/49] [CI][Tests] Account for FP8 KV-cache rounding in MLA reference (#62067) Signed-off-by: wzhao18 Co-authored-by: Codex --- tests/v1/attention/test_mla_backends.py | 29 +++++++++++++++++++++---- 1 file changed, 25 insertions(+), 4 deletions(-) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 891aa8880d..9d4be102d9 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -1618,6 +1618,14 @@ def _with_even_cpu_query_split( ) +def _fp8_cache_roundtrip(tensor: torch.Tensor, scale: float) -> torch.Tensor: + """Reproduce cache storage rounding for the SDPA reference.""" + fp8_dtype = current_platform.fp8_dtype() + limits = torch.finfo(fp8_dtype) + quantized = (tensor.float() / scale).clamp(limits.min, limits.max).to(fp8_dtype) + return (quantized.float() * scale).to(tensor.dtype) + + def _run_backend_correctness( default_vllm_config, dist_init, @@ -1801,6 +1809,19 @@ def _run_backend_correctness( # K_PE (rope component): [s_len, 1, qk_rope_head_dim] k_pe_full = torch.randn(s_len, 1, qk_rope_head_dim, dtype=dtype, device=device) + kv_c_cached, k_pe_cached = kv_c_full, k_pe_full + kv_c_prefill, k_pe_prefill = kv_c_full, k_pe_full + if kv_cache_dtype.startswith("fp8"): + kv_c_cached = _fp8_cache_roundtrip(kv_c_full, k_scale) + k_pe_cached = _fp8_cache_roundtrip(k_pe_full, k_scale) + # Prefill reads cached context and unquantized new tokens. + kv_c_prefill = torch.cat( + [kv_c_cached[:context_len], kv_c_full[context_len:]] + ) + k_pe_prefill = torch.cat( + [k_pe_cached[:context_len], k_pe_full[context_len:]] + ) + # Determine if this sequence uses the decode pipeline or prefill # pipeline for each backend # NOTE: For spec decode tests with uniform query_len > 1, backends that @@ -1850,10 +1871,10 @@ def _run_backend_correctness( q_mqa = torch.cat([ql_nope, q_pe], dim=-1) # K: [s_len, kv_lora_rank + qk_rope_head_dim] # (broadcasted to all heads) - k_mqa = torch.cat([kv_c_full, k_pe_full.squeeze(1)], dim=-1) + k_mqa = torch.cat([kv_c_cached, k_pe_cached.squeeze(1)], dim=-1) k_mqa = k_mqa.unsqueeze(1).expand(-1, num_q_heads, -1) # V: [s_len, kv_lora_rank] (broadcasted to all heads) - v_mqa = kv_c_full.unsqueeze(1).expand(-1, num_q_heads, -1) + v_mqa = kv_c_cached.unsqueeze(1).expand(-1, num_q_heads, -1) # Create custom attention mask for decode path: # - Query tokens can attend to all context tokens @@ -1882,12 +1903,12 @@ def _run_backend_correctness( ####################################################### # Prefill path: MHA-style attention with full sequence # Apply kv_b_proj to the full kv_c tensor - kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_full, kv_b_proj_weight) + kv_nope_full = torch.einsum("sl,lnh->snh", kv_c_prefill, kv_b_proj_weight) k_nope_full, v_full = kv_nope_full.split([qk_nope_head_dim, v_head_dim], dim=-1) # Build attention inputs for full sequence q_mha = torch.cat([q_nope, q_pe], dim=-1) # [q_len, num_heads, total_dim] - k_pe_full_expanded = k_pe_full.expand(-1, num_q_heads, -1) + k_pe_full_expanded = k_pe_prefill.expand(-1, num_q_heads, -1) k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1) # Create custom attention mask: From 783380af397ceeae91a40f4f759439eb8c1ab690 Mon Sep 17 00:00:00 2001 From: Vadim Gimpelson <156319763+vadiklyutiy@users.noreply.github.com> Date: Sat, 17 Oct 2026 04:24:51 +0400 Subject: [PATCH 31/49] [Bugfix] Fix NaN output of the Mamba2 chunk scan kernel on sm90 (#61718) Signed-off-by: Vadim Gimpelson Co-authored-by: Claude --- tests/kernels/mamba/test_mamba_ssm_ssd.py | 47 +++ .../layers/mamba/ops/ssd_chunk_scan.py | 277 ++++++++++-------- 2 files changed, 197 insertions(+), 127 deletions(-) diff --git a/tests/kernels/mamba/test_mamba_ssm_ssd.py b/tests/kernels/mamba/test_mamba_ssm_ssd.py index 22abeab1ed..8522477045 100644 --- a/tests/kernels/mamba/test_mamba_ssm_ssd.py +++ b/tests/kernels/mamba/test_mamba_ssm_ssd.py @@ -6,6 +6,10 @@ import torch import torch.nn.functional as F from einops import rearrange, repeat +from vllm.model_executor.layers.mamba.ops.ssd_chunk_scan import ( + _chunk_scan_fwd, + _chunk_scan_fwd_kernel, +) from vllm.model_executor.layers.mamba.ops.ssd_combined import ( mamba_chunk_scan_combined_varlen, ) @@ -578,3 +582,46 @@ def test_mamba_chunk_scan_cont_batch_prefill_chunking(chunk_size, seqlens): rtol=rtol, msg=lambda x, i=i: f"seq{i} state " + x, ) + + +def test_chunk_scan_fwd_every_config_matches(monkeypatch): + """The autotuner picks a config by speed only, so every config in the list + must give the same finite output. On sm90, configs that pipeline a wgmma + with `cb` in registers wrote NaN into whole rows (#61603).""" + from triton.runtime.errors import OutOfResources + + set_random_seed(0) + nheads, headdim, ngroups, dstate, chunk_size = 64, 64, 8, 128, 128 + cu_chunk_seqlens = torch.tensor([0, 128, 181], dtype=torch.int32, device=DEVICE) + seqlen, nchunks = 181, 2 + x = torch.randn(seqlen, nheads, headdim, dtype=torch.bfloat16, device=DEVICE) + C = torch.randn(seqlen, ngroups, dstate, dtype=torch.bfloat16, device=DEVICE) + cb = torch.randn(nchunks, ngroups, chunk_size, chunk_size, device=DEVICE) + dt = torch.rand(nheads, nchunks, chunk_size, device=DEVICE) * 0.1 + dA_cumsum = -torch.rand(nheads, nchunks, chunk_size, device=DEVICE).cumsum(-1) + states = torch.randn(nchunks, nheads, headdim, dstate, device=DEVICE) * 0.1 + seq_idx = torch.zeros(nchunks, dtype=torch.int32, device=DEVICE) + D = torch.randn(nheads, headdim, device=DEVICE) * 0.1 + + ref = None + for config in list(_chunk_scan_fwd_kernel.configs): + monkeypatch.setattr(_chunk_scan_fwd_kernel, "configs", [config]) + for _ in range(10): + out = torch.empty_like(x) + try: + _chunk_scan_fwd( + cb, x, dt, dA_cumsum, C, states, cu_chunk_seqlens, out, seq_idx, D + ) + except OutOfResources: + # The autotuner skips configs that do not fit this device. + break + assert torch.isfinite(out).all(), f"non-finite output with {config}" + if ref is None: + ref = out + torch.testing.assert_close( + out, + ref, + atol=2e-2, + rtol=2e-2, + msg=lambda m, config=config: f"{config}: {m}", + ) diff --git a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py index bd554ede5a..237dd33b52 100644 --- a/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py +++ b/vllm/model_executor/layers/mamba/ops/ssd_chunk_scan.py @@ -9,139 +9,162 @@ from packaging import version from vllm.model_executor.layers.mamba.ops.triton_helpers import fast_exp +from vllm.platforms import current_platform from vllm.triton_utils import tl, triton TRITON_22 = version.parse(triton.__version__) >= version.parse("2.2.0") +def _drop_racy_wgmma_configs(configs: list[triton.Config]) -> list[triton.Config]: + """Drop the configs that Triton miscompiles on sm90. + + With whole warpgroups (num_warps % 4 == 0) and BLOCK_SIZE_M % 64 == 0, + `tl.dot(cb, x)` lowers to wgmma with `cb` in registers. With + num_stages > 1, Triton leaves that wgmma running across the loop + back-edge while the next iteration overwrites the registers of `cb`, + which writes NaN/Inf into whole output rows. + """ + if not (current_platform.is_cuda() and current_platform.is_device_capability(90)): + return configs + return [ + c + for c in configs + if c.num_stages == 1 + or c.num_warps % 4 != 0 + or c.kwargs["BLOCK_SIZE_M"] % 64 != 0 + ] + + @triton.autotune( - configs=[ - # ================================================================= - # Higher warp count configs for better latency hiding - # More warps = more instructions in flight = better memory latency hiding - # ================================================================= - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=8, - ), - # Smaller tiles with more stages for software pipelining - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=3, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64}, - num_stages=2, - num_warps=4, - ), - # ================================================================= - # Low register pressure configs (num_stages=1) for large dstate - # ================================================================= - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, - num_stages=1, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=1, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=1, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=1, - num_warps=4, - ), - # num_stages=2 configs - moderate register pressure - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, - num_stages=2, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=2, - num_warps=4, - ), - # Original configs for larger dstate values - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64}, - num_stages=3, - num_warps=8, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=4, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, - num_stages=5, - num_warps=2, - ), - triton.Config( - {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=5, - num_warps=2, - ), - triton.Config( - {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, - num_stages=4, - num_warps=2, - ), - ], + configs=_drop_racy_wgmma_configs( + [ + # ================================================================= + # Higher warp count configs for better latency hiding + # More warps = more instructions in flight = better memory latency hiding + # ================================================================= + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=8, + ), + # Smaller tiles with more stages for software pipelining + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=3, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 64}, + num_stages=2, + num_warps=4, + ), + # ================================================================= + # Low register pressure configs (num_stages=1) for large dstate + # ================================================================= + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, + num_stages=1, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=1, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=1, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=1, + num_warps=4, + ), + # num_stages=2 configs - moderate register pressure + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, + num_stages=2, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=2, + num_warps=4, + ), + # Original configs for larger dstate values + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64}, + num_stages=3, + num_warps=8, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 64}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 64}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=4, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32}, + num_stages=5, + num_warps=2, + ), + triton.Config( + {"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=5, + num_warps=2, + ), + triton.Config( + {"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32}, + num_stages=4, + num_warps=2, + ), + ] + ), key=["chunk_size", "hdim", "dstate", "IS_CAUSAL"], ) @triton.jit From f019e3e1055424075b440dfe40198d5d45e1dfe1 Mon Sep 17 00:00:00 2001 From: Omar Muhammad Date: Sat, 17 Oct 2026 03:54:54 +0300 Subject: [PATCH 32/49] [ROCm] attn res + sigmoid_mul + conv fusions (#52968) Signed-off-by: Omar Muhammad Signed-off-by: Nicholas Susanto Co-authored-by: Nicholas Susanto Co-authored-by: Douglas Lehr <91553416+dllehr-amd@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../kimi_k3/test_amd_kda_fused_conv1d.py | 163 ++++++++++++++++++ vllm/models/kimi_k3/amd/kda.py | 79 +++++---- vllm/models/kimi_k3/amd/mla.py | 12 +- 3 files changed, 219 insertions(+), 35 deletions(-) create mode 100644 tests/models/kimi_k3/test_amd_kda_fused_conv1d.py diff --git a/tests/models/kimi_k3/test_amd_kda_fused_conv1d.py b/tests/models/kimi_k3/test_amd_kda_fused_conv1d.py new file mode 100644 index 0000000000..803655b893 --- /dev/null +++ b/tests/models/kimi_k3/test_amd_kda_fused_conv1d.py @@ -0,0 +1,163 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch +import torch.nn.functional as F + +from vllm.model_executor.layers.mamba.ops.causal_conv1d import causal_conv1d_fn +from vllm.platforms import current_platform +from vllm.v1.attention.backends.utils import NULL_BLOCK_ID + +pytestmark = pytest.mark.skipif( + not current_platform.is_rocm(), reason="AITER fused conv1d is ROCm-only" +) + +WIDTH = 4 # Kimi-K3 short_conv_kernel_size +STATE_LEN = WIDTH - 1 +HEAD_DIM = 128 +LOCAL_HEADS = 4 # per-TP-rank heads; projection = 512, packed dim = 1536 + + +def _fused_fn(): + gdn = pytest.importorskip("aiter.ops.triton.gated_delta_net") + return gdn.causal_conv1d_split_qkv_triton_tile_fn + + +def _make_inputs(seqlens, has_initial_state, seed=0, device="cuda"): + torch.manual_seed(seed) + proj = HEAD_DIM * LOCAL_HEADS + dim = 3 * proj + total = sum(seqlens) + n_seqs = len(seqlens) + # Token-major activations, as produced by the in_proj GEMM. + mixed_qkv = torch.randn(total, dim, dtype=torch.bfloat16, device=device) + weight = torch.randn(dim, WIDTH, dtype=torch.bfloat16, device=device) * 0.3 + # More cache lines than sequences, assigned in a shuffled order. Line 0 is + # vLLM's null block (NULL_BLOCK_ID), which causal_conv1d_fn skips, so real + # sequences never map to it. + n_lines = n_seqs + 2 + conv_state = torch.randn( + n_lines, dim, STATE_LEN, dtype=torch.bfloat16, device=device + ) + cache_indices = (torch.randperm(n_lines - 1, device=device)[:n_seqs] + 1).to( + torch.int32 + ) + query_start_loc = torch.tensor( + [0] + list(torch.tensor(seqlens).cumsum(0).tolist()), + dtype=torch.int32, + device=device, + ) + has_initial = torch.full( + (n_seqs,), has_initial_state, dtype=torch.bool, device=device + ) + return mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial + + +def _reference( + mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial +): + """fp32 per-sequence depthwise causal conv + SiLU. + + Returns (out[T, dim], new_conv_state). + """ + x = mixed_qkv.float().t() # (dim, T) + w = weight.float() + dim = x.shape[0] + out = torch.empty_like(x) + new_state = conv_state.float().clone() + qsl = query_start_loc.tolist() + for s in range(len(qsl) - 1): + start, end = qsl[s], qsl[s + 1] + ci = int(cache_indices[s]) + xs = x[:, start:end] + if bool(has_initial[s]): + xin = torch.cat([conv_state[ci].float(), xs], dim=1) + else: + xin = F.pad(xs, (STATE_LEN, 0)) + y = F.conv1d(xin.unsqueeze(0), w.unsqueeze(1), None, groups=dim)[0] + out[:, start:end] = F.silu(y) + new_state[ci] = xin[:, -STATE_LEN:] + return out.t(), new_state + + +def _three_launch( + mixed_qkv, weight, conv_state, cache_indices, query_start_loc, has_initial +): + """The pre-fusion KDA prefill path: one causal_conv1d_fn per projection.""" + proj = mixed_qkv.shape[1] // 3 + outs = [] + for x, state, w in zip( + mixed_qkv.split(proj, dim=-1), + conv_state.split(proj, dim=-2), + weight.split(proj, dim=0), + ): + outs.append( + causal_conv1d_fn( + x.transpose(0, 1), + w, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial, + cache_indices=cache_indices, + query_start_loc=query_start_loc, + ).transpose(0, 1) + ) + return outs + + +@pytest.mark.parametrize( + "seqlens", [[1024], [8, 1, 300], [64, 64, 64], [513, 64, 1, 1200]] +) +@pytest.mark.parametrize("has_initial_state", [False, True]) +@pytest.mark.parametrize("block_m", [8, 128]) +def test_fused_split_qkv_conv1d_matches_reference(seqlens, has_initial_state, block_m): + fused_fn = _fused_fn() + mixed_qkv, weight, conv_state, cache_indices, qsl, has_initial = _make_inputs( + seqlens, has_initial_state + ) + proj = mixed_qkv.shape[1] // 3 + + ref_out, ref_state = _reference( + mixed_qkv, weight, conv_state, cache_indices, qsl, has_initial + ) + + state_3 = conv_state.clone() + q3, k3, v3 = _three_launch( + mixed_qkv, weight, state_3, cache_indices, qsl, has_initial + ) + + state_f = conv_state.clone() + # Exactly the call made by KimiK3DeltaAttention._forward. + qf, kf, vf = fused_fn( + mixed_qkv.transpose(0, 1), + weight, + None, + state_f, + qsl, + proj, + proj, + cache_indices=cache_indices, + has_initial_state=has_initial, + activation="silu", + block_m=block_m, + metadata=None, + preserve_input_dtype=True, + ) + + fused = torch.cat([qf, kf, vf], dim=-1) + assert fused.shape == mixed_qkv.shape and fused.dtype == mixed_qkv.dtype + assert all(t.is_contiguous() for t in (qf, kf, vf)) + + tol = dict(atol=2e-2, rtol=2e-2) + torch.testing.assert_close(fused.float(), ref_out, **tol) + torch.testing.assert_close(fused, torch.cat([q3, k3, v3], dim=-1), **tol) + + # conv_state writeback: assigned lines updated, the others untouched. + torch.testing.assert_close(state_f.float(), ref_state, **tol) + torch.testing.assert_close(state_f, state_3, **tol) + untouched = torch.ones(conv_state.shape[0], dtype=torch.bool, device="cuda") + untouched[cache_indices.long()] = False + assert untouched[NULL_BLOCK_ID] + assert torch.equal(state_f[untouched], conv_state[untouched]) diff --git a/vllm/models/kimi_k3/amd/kda.py b/vllm/models/kimi_k3/amd/kda.py index 08ad0b650d..73e357e444 100644 --- a/vllm/models/kimi_k3/amd/kda.py +++ b/vllm/models/kimi_k3/amd/kda.py @@ -134,6 +134,14 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): self.conv_size = kda_config["short_conv_kernel_size"] self.use_full_rank_gate = True + self._fused_qkv_conv_fn = None + if rocm_aiter_ops.is_enabled(): + from aiter.ops.triton.gated_delta_net import ( + causal_conv1d_split_qkv_triton_tile_fn, + ) + + self._fused_qkv_conv_fn = causal_conv1d_split_qkv_triton_tile_fn + # Keep f_a before the narrow beta shard, then pad each TP-local row to # select the aligned BF16 GEMM path. The padding also avoids an Inductor # correctness issue seen with the row-strided G view. @@ -580,35 +588,48 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): if mixed_qkv_ns is not None: assert g1_ns is not None and beta_ns is not None if m.num_prefills > 0: - q_ns, k_ns, v_ns = mixed_qkv_ns.split( - self.local_projection_size, dim=-1 - ) - - # Packed prefill conv would require copying V solely to make - # it dense for KDA. Separate calls accept the strided inputs - # and produce dense Q/K/V without that extra traffic. - # TODO: Use packed conv once every KDA prefill backend accepts - # row-strided Q/K/V directly. - def _prefill_conv( - x: torch.Tensor, - state: torch.Tensor, - weight: torch.Tensor, - ) -> torch.Tensor: - return causal_conv1d_fn( - x.transpose(0, 1), - weight, + if self._fused_qkv_conv_fn is not None: + # One launch (fused qkv) instead of 3 + q_ns, k_ns, v_ns = self._fused_qkv_conv_fn( + mixed_qkv_ns.transpose(0, 1), + conv_weights, None, - activation="silu", - conv_states=state, - has_initial_state=has_initial_state, + conv_state, + non_spec_query_start_loc, + self.local_projection_size, + self.local_projection_size, cache_indices=non_spec_state_indices_tensor, - query_start_loc=non_spec_query_start_loc, + has_initial_state=has_initial_state, + activation="silu", + block_m=8, metadata=m, - ).transpose(0, 1) + preserve_input_dtype=True, + ) + else: + q_ns, k_ns, v_ns = mixed_qkv_ns.split( + self.local_projection_size, dim=-1 + ) - q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) - k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) - v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) + def _prefill_conv( + x: torch.Tensor, + state: torch.Tensor, + weight: torch.Tensor, + ) -> torch.Tensor: + return causal_conv1d_fn( + x.transpose(0, 1), + weight, + None, + activation="silu", + conv_states=state, + has_initial_state=has_initial_state, + cache_indices=non_spec_state_indices_tensor, + query_start_loc=non_spec_query_start_loc, + metadata=m, + ).transpose(0, 1) + + q_ns = _prefill_conv(q_ns, q_conv_state, q_conv_weight) + k_ns = _prefill_conv(k_ns, k_conv_state, k_conv_weight) + v_ns = _prefill_conv(v_ns, v_conv_state, v_conv_weight) q_ns, k_ns, v_ns = ( rearrange(x, "n (h d) -> 1 n h d", d=self.head_dim) for x in (q_ns, k_ns, v_ns) @@ -752,15 +773,9 @@ class KimiK3DeltaAttention(GatedDeltaNetAttention): # ---------- merge spec and non-spec outputs ---------- if core_attn_out_spec is not None and core_attn_out_non_spec is not None: - # Mixed batches require indexed placement in the original order. - merged = torch.empty( - (1, num_actual_tokens, *core_attn_out_spec.shape[2:]), - dtype=core_attn_out_spec.dtype, - device=core_attn_out_spec.device, - ) + merged = core_attn_out[:, :num_actual_tokens] merged.index_copy_(1, spec_token_indx, core_attn_out_spec) merged.index_copy_(1, non_spec_token_indx, core_attn_out_non_spec) - core_attn_out[0, :num_actual_tokens] = merged[0, :num_actual_tokens] elif core_attn_out_non_spec is not None: if core_attn_out_non_spec.data_ptr() != core_attn_out.data_ptr(): core_attn_out[0, :num_actual_tokens] = core_attn_out_non_spec[ diff --git a/vllm/models/kimi_k3/amd/mla.py b/vllm/models/kimi_k3/amd/mla.py index dd93ce439a..48ea1792bb 100644 --- a/vllm/models/kimi_k3/amd/mla.py +++ b/vllm/models/kimi_k3/amd/mla.py @@ -23,8 +23,8 @@ _OPT_MIN_SIZE = 2048 class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): - """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm fusion and a fused - decode Q-prep path.""" + """Kimi-K3 MLA wrapper with eager AITER q/kv RMSNorm and output gate + fusions and a fused decode Q-prep path.""" def __init__(self, *args, **kwargs) -> None: super().__init__(*args, **kwargs) @@ -303,6 +303,12 @@ class KimiK3MultiHeadLatentAttentionWrapper(MultiHeadLatentAttentionWrapper): ) if self.g_proj is not None: - attn_out = attn_out * self.g_proj(hidden_states)[0].sigmoid() + gate = self.g_proj(hidden_states)[0] + if rocm_aiter_ops.is_enabled(): + from aiter.ops.triton.fusions.fused_sigmoid_mul import fused_sigmoid_mul + + attn_out = fused_sigmoid_mul(attn_out, gate) + else: + attn_out = attn_out * gate.sigmoid() return self.o_proj(attn_out)[0] From f8112704d54a15828c30496d49a73238be66f997 Mon Sep 17 00:00:00 2001 From: Harry Mellor <19981378+hmellor@users.noreply.github.com> Date: Sat, 17 Oct 2026 02:04:54 +0100 Subject: [PATCH 33/49] [Chore] Use upstream Transformers InklingProcessor (#61883) Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- vllm/models/inkling/common/mm_preprocess.py | 120 +++-- .../transformers_utils/processors/__init__.py | 6 - vllm/transformers_utils/processors/inkling.py | 504 ------------------ 3 files changed, 61 insertions(+), 569 deletions(-) delete mode 100644 vllm/transformers_utils/processors/inkling.py diff --git a/vllm/models/inkling/common/mm_preprocess.py b/vllm/models/inkling/common/mm_preprocess.py index a949e4ec74..60ea4a92eb 100644 --- a/vllm/models/inkling/common/mm_preprocess.py +++ b/vllm/models/inkling/common/mm_preprocess.py @@ -4,12 +4,15 @@ from __future__ import annotations +import math from collections.abc import Iterable, Mapping, Sequence from typing import Any, cast import numpy as np import regex as re import torch +from PIL import Image +from transformers import InklingProcessor from transformers.feature_extraction_utils import BatchFeature from vllm.config.multimodal import ( @@ -29,18 +32,19 @@ from vllm.multimodal.processing import ( PromptUpdate, PromptUpdateDetails, ) -from vllm.transformers_utils.processors.inkling import ( - AUDIO_MARKER_ID, - AUDIO_TOKEN_ID, - IMAGE_MARKER_ID, - IMAGE_TOKEN_ID, - InklingAudioFeatureExtractor, - InklingImageProcessor, - InklingProcessor, -) from ..configs import InklingMMConfig +# Block-start markers (<|content_image|>, <|content_audio_input|>), kept verbatim +IMAGE_MARKER_ID = 200005 +AUDIO_MARKER_ID = 200020 +# Per-patch / per-frame placeholders (<|unused_200054|>, <|unused_200053|>) +IMAGE_TOKEN_ID = 200054 +AUDIO_TOKEN_ID = 200053 + +# Long-edge upscale applied before patchifying (upstream defaults to no rescale) +DEFAULT_RESCALE_IMAGE_FRAC = 2.0 + # Maximum audio tokens accepted per clip. At the dMel rate of 20 tokens/s # (50 ms hop) this is ~10 minutes of audio. It bounds the persistent per-request # buffers and the encoder/memory budget; longer clips are rejected up front. @@ -57,6 +61,24 @@ class InklingMultiModalDataParser(MultiModalDataParser): return super()._parse_audio_data(data) +def _rescale_image( + image: Image.Image, frac: float | None, max_upscaled_long_edge: int | None +) -> Image.Image: + """Scale the long edge by `frac` with PIL LANCZOS, as the reference does.""" + image = image.convert("RGB") + if frac is None: + return image + long_edge = max(image.size) + target_long_edge = long_edge * frac + if max_upscaled_long_edge is not None: + target_long_edge = min(target_long_edge, max(max_upscaled_long_edge, long_edge)) + ratio = target_long_edge / long_edge + if ratio == 1.0: + return image + size = tuple(max(1, math.floor(d * ratio + 0.5)) for d in image.size) + return image.resize(size, resample=Image.Resampling.LANCZOS) + + def inkling_vision_enabled(config: InklingMMConfig) -> bool: return getattr(config.vision_config, "decoder_dmodel", None) is not None @@ -70,38 +92,13 @@ class InklingProcessingInfo(BaseProcessingInfo): return self.ctx.get_hf_config(InklingMMConfig) def get_hf_processor(self, **kwargs: object) -> InklingProcessor: - config = self.get_hf_config() - vision_config = config.vision_config - audio_config = config.audio_config - - image_processor = InklingImageProcessor( - patch_size=getattr(vision_config, "patch_size", None) or 40, - ) - - if inkling_audio_enabled(config): - audio_params = { - "n_mels": audio_config.n_mel_bins, - "num_dmel_bins": audio_config.mel_vocab_size, - "dmel_min_value": audio_config.dmel_min_value, - "dmel_max_value": audio_config.dmel_max_value, - } - else: - audio_params = {} - audio_extractor = InklingAudioFeatureExtractor(params=audio_params) - - return InklingProcessor( - image_processor=image_processor, - audio_feature_extractor=audio_extractor, - tokenizer=self.get_tokenizer(), - ) + return self.ctx.get_hf_processor(InklingProcessor, **kwargs) def get_data_parser(self) -> MultiModalDataParser: - # Audio inputs must be resampled to the dMel feature extractor's rate - # before `process_audios` (see InklingAudioFeatureExtractor._decode_one). - # Without a target_sr the default parser raises on any audio input. - extractor = self.get_hf_processor().audio_feature_extractor + # The feature extractor requires audio at its sampling rate + feature_extractor = self.get_hf_processor().feature_extractor return InklingMultiModalDataParser( - target_sr=extractor.params.sample_rate, + target_sr=feature_extractor.sampling_rate, target_channels=1, expected_hidden_size=self._get_expected_hidden_size(), ) @@ -163,9 +160,8 @@ class InklingDummyInputsBuilder(BaseDummyInputsBuilder[InklingProcessingInfo]): if num_audios: # Size the dummy at the maximum allowed audio so memory/encoder # budgeting reflects the largest clip we accept (MAX_AUDIO_TOKENS). - params = self.info.get_hf_processor().audio_feature_extractor.params - hop = int(round(params.audio_token_duration_s * params.sample_rate)) - audio_len = MAX_AUDIO_TOKENS * hop + feature_extractor = self.info.get_hf_processor().feature_extractor + audio_len = MAX_AUDIO_TOKENS * feature_extractor.hop_length mm_data["audio"] = self._get_dummy_audios( length=audio_len, num_audios=num_audios, @@ -203,31 +199,37 @@ class InklingMultiModalProcessor(BaseMultiModalProcessor[InklingProcessingInfo]) data: dict[str, Any] = {"input_ids": [prompt_ids]} if images: - img_feat = processor.process_images(images) - data["pixel_values"] = img_feat["vision_patches_bthwc"] - data["num_patches"] = torch.tensor( - img_feat["num_patches"], dtype=torch.int64 + image_processor = processor.image_processor + frac = image_processor.rescale_image_frac or DEFAULT_RESCALE_IMAGE_FRAC + max_long_edge = image_processor.rescale_image_max_upscaled_long_edge + # Resize with PIL: torchvision LANCZOS does not match the reference + images = [_rescale_image(img, frac, max_long_edge) for img in images] + img_feat = image_processor( + images, rescale_image_frac=None, return_tensors="pt" ) + data["pixel_values"] = img_feat["pixel_values"].to(torch.bfloat16) + data["num_patches"] = img_feat["num_patches"].to(torch.int64) if audios: - aud_feat = processor.process_audios(audios) - per_clip = aud_feat["dmel_bins"] - num_audio_tokens = aud_feat["num_audio_tokens"] - for i, n in enumerate(num_audio_tokens): - if int(n) > MAX_AUDIO_TOKENS: + feature_extractor = processor.feature_extractor + aud_feat = feature_extractor( + audios, + sampling_rate=feature_extractor.sampling_rate, + return_tensors="pt", + ) + audio_input_ids = processor._extract_dmel_bins(aud_feat["input_features"]) + num_audio_tokens = aud_feat["input_features_mask"].sum(-1) + for i, n in enumerate(num_audio_tokens.tolist()): + if n > MAX_AUDIO_TOKENS: raise ValueError( - f"Audio clip {i} produces {int(n)} tokens, exceeding the " + f"Audio clip {i} produces {n} tokens, exceeding the " f"maximum of {MAX_AUDIO_TOKENS} (~10 min at 20 tokens/s). " "Provide a shorter clip." ) - if per_clip: - input_audio_features = torch.cat( - [torch.as_tensor(c) for c in per_clip], dim=0 - ) - else: - input_audio_features = torch.empty(0) - data["input_audio_features"] = input_audio_features - data["num_audio_tokens"] = torch.tensor(num_audio_tokens, dtype=torch.int64) + data["input_audio_features"] = torch.cat( + [ids[:n] for ids, n in zip(audio_input_ids, num_audio_tokens)] + ) + data["num_audio_tokens"] = num_audio_tokens.to(torch.int64) processed_data = BatchFeature(data=data, tensor_type=None) return self._finalize_hf_mm_data( diff --git a/vllm/transformers_utils/processors/__init__.py b/vllm/transformers_utils/processors/__init__.py index db45e49d5a..c06519ef18 100644 --- a/vllm/transformers_utils/processors/__init__.py +++ b/vllm/transformers_utils/processors/__init__.py @@ -39,9 +39,6 @@ __all__ = [ "OvisProcessor", "Ovis2_5Processor", "Step3VLProcessor", - "InklingProcessor", - "InklingImageProcessor", - "InklingAudioFeatureExtractor", ] _CLASS_TO_MODULE: dict[str, str] = { @@ -74,9 +71,6 @@ _CLASS_TO_MODULE: dict[str, str] = { "OvisProcessor": "vllm.transformers_utils.processors.ovis", "Ovis2_5Processor": "vllm.transformers_utils.processors.ovis2_5", "Step3VLProcessor": "vllm.transformers_utils.processors.step3_vl", - "InklingProcessor": "vllm.transformers_utils.processors.inkling", - "InklingImageProcessor": "vllm.transformers_utils.processors.inkling", - "InklingAudioFeatureExtractor": "vllm.transformers_utils.processors.inkling", } diff --git a/vllm/transformers_utils/processors/inkling.py b/vllm/transformers_utils/processors/inkling.py deleted file mode 100644 index 09d34f7806..0000000000 --- a/vllm/transformers_utils/processors/inkling.py +++ /dev/null @@ -1,504 +0,0 @@ -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Vendored HuggingFace-convention processors for the Inkling Titan model. - -Implements the image processor (numba patchifier), the audio feature extractor -(STFT/dMel path), the composite processor, and the MM token-id constants. - -Besides raw bytes / file paths, the extractors also accept the dummy inputs -vLLM generates during profiling (PIL images / numpy audio arrays). -""" - -from __future__ import annotations - -import io -import math -from collections.abc import Sequence -from dataclasses import dataclass - -import numpy as np -import torch -import torch.nn.functional as F -from numba import njit -from transformers.feature_extraction_utils import ( - BatchFeature, - FeatureExtractionMixin, -) -from transformers.image_processing_utils import BaseImageProcessor -from transformers.image_utils import ImageInput - -# --------------------------------------------------------------------------- -# MM token-id constants -# --------------------------------------------------------------------------- - -# Block-start marker token ids. These are real tokens the model was trained on -# (``<|content_image|>`` / ``<|content_audio_input|>``) that mark the start of an -# image/audio embedding block; they are kept verbatim in ``input_ids``. -IMAGE_MARKER_ID = 200005 # <|content_image|> -AUDIO_MARKER_ID = 200020 # <|content_audio_input|> - -# Per-patch / per-frame placeholder ids marking where tower embeddings are -# scattered in. These are unused slots in the padded vocabulary and the -# corresponding positions are always overwritten by tower embeddings. -IMAGE_TOKEN_ID = 200054 # <|unused_200054|> -AUDIO_TOKEN_ID = 200053 # <|unused_200053|> - -# --------------------------------------------------------------------------- -# Image processing -# --------------------------------------------------------------------------- - -IMAGE_MEAN = np.array([0.48145466, 0.4578275, 0.40821073], dtype=np.float32) -IMAGE_STD = np.array([0.26862954, 0.2613026, 0.2757771], dtype=np.float32) -PAD_RAW_VALUE = np.float32(-1.0 / 255.0) -PAD_NORM = (np.full((3,), PAD_RAW_VALUE, dtype=np.float32) - IMAGE_MEAN) / IMAGE_STD - - -def _validate_image_rescale( - rescale_image_frac: float | None, - rescale_image_max_upscaled_long_edge: int | None, -) -> None: - if rescale_image_frac is not None and ( - not math.isfinite(rescale_image_frac) or rescale_image_frac <= 0 - ): - raise ValueError( - "rescale_image_frac must be positive and finite or None, " - f"got {rescale_image_frac}" - ) - if rescale_image_max_upscaled_long_edge is None: - return - if rescale_image_max_upscaled_long_edge <= 0: - raise ValueError( - "rescale_image_max_upscaled_long_edge must be positive or None, " - f"got {rescale_image_max_upscaled_long_edge}" - ) - if rescale_image_frac is None or rescale_image_frac <= 1.0: - raise ValueError( - "rescale_image_max_upscaled_long_edge requires rescale_image_frac > 1, " - f"got {rescale_image_frac}" - ) - - -def _scaled_image_dimensions( - width: int, - height: int, - rescale_image_frac: float | None, - rescale_image_max_upscaled_long_edge: int | None, -) -> tuple[int, int]: - """Return the long-edge-scaled ``(width, height)``.""" - if rescale_image_frac is None: - return width, height - - long_edge = max(width, height) - if long_edge == 0: - return width, height - - target_long_edge = float(long_edge) * rescale_image_frac - if rescale_image_max_upscaled_long_edge is not None: - effective_cap = max(rescale_image_max_upscaled_long_edge, long_edge) - target_long_edge = min(target_long_edge, float(effective_cap)) - - ratio = target_long_edge / float(long_edge) - if ratio == 1.0: - return width, height - - def scale(value: int) -> int: - return max(1, math.floor(float(value) * ratio + 0.5)) - - return scale(width), scale(height) - - -def _load_image_bytes(image) -> bytes: - """Encode a PIL image as raw PNG bytes for preprocessing. - - The HF processor is always handed ``PIL.Image`` instances by vLLM, so no - other input types need to be supported here. - """ - if image.mode != "RGB": - image = image.convert("RGB") - buf = io.BytesIO() - image.save(buf, format="PNG") - return buf.getvalue() - - -@njit(cache=True) -def _fill_patches_numba( - arr: np.ndarray, - patch_size: int, - patches: np.ndarray, - mean: np.ndarray, - std: np.ndarray, - pad_norm: np.ndarray, -) -> None: - h = arr.shape[0] - w = arr.shape[1] - nph = (h + patch_size - 1) // patch_size - npw = w // patch_size + 1 - inv255 = np.float32(1.0 / 255.0) - - for k in range(nph * npw): - i = k // npw - j = k - i * npw - y_base = i * patch_size - x_base = j * patch_size - - for y in range(patch_size): - iy = y_base + y - for x in range(patch_size): - ix = x_base + x - if iy < h and ix < w: - for c in range(3): - raw = np.float32(arr[iy, ix, c]) * inv255 - patches[k, y, x, c] = (raw - mean[c]) / std[c] - else: - for c in range(3): - patches[k, y, x, c] = pad_norm[c] - - -def _encode_image_bytes( - image_bytes: bytes, - *, - patch_size: int, - rescale_image_frac: float | None, - rescale_image_max_upscaled_long_edge: int | None, -) -> torch.Tensor: - if patch_size <= 0: - raise ValueError("patch_size must be greater than zero") - _validate_image_rescale( - rescale_image_frac, - rescale_image_max_upscaled_long_edge, - ) - - from PIL import Image - - image = Image.open(io.BytesIO(image_bytes)).convert("RGB") - scaled_size = _scaled_image_dimensions( - image.width, - image.height, - rescale_image_frac=rescale_image_frac, - rescale_image_max_upscaled_long_edge=rescale_image_max_upscaled_long_edge, - ) - if scaled_size != image.size: - image = image.resize(scaled_size, resample=Image.Resampling.LANCZOS) - arr = np.array(image, dtype=np.uint8, copy=True) - height, width, _ = arr.shape - - nph = (height + patch_size - 1) // patch_size - npw = width // patch_size + 1 - num_patches = nph * npw - - patches = np.empty((num_patches, patch_size, patch_size, 3), dtype=np.float32) - _fill_patches_numba(arr, patch_size, patches, IMAGE_MEAN, IMAGE_STD, PAD_NORM) - - return ( - torch.from_numpy(patches) - .to(torch.bfloat16) - .view(num_patches, 1, patch_size, patch_size, 3) - .expand(num_patches, 2, patch_size, patch_size, 3) - ) - - -class InklingImageProcessor(BaseImageProcessor): - r"""Turn raw images into ``vision_patches_bthwc`` for Inkling hMLP. - - ``rescale_image_frac`` scales the long edge while preserving aspect ratio. - ``rescale_image_max_upscaled_long_edge`` optionally caps only upscaling and - therefore requires a scale factor greater than one. The defaults, ``2.0`` and - ``2048``, grow images toward a 2048-pixel long edge by at most 2x, while leaving - images already at or above 2048 unchanged. - """ - - model_input_names = ["vision_patches_bthwc"] - - def __init__( - self, - patch_size: int = 40, - rescale_image_frac: float | None = 2.0, - rescale_image_max_upscaled_long_edge: int | None = 2048, - **kwargs, - ): - if patch_size <= 0: - raise ValueError("patch_size must be greater than zero") - _validate_image_rescale( - rescale_image_frac, - rescale_image_max_upscaled_long_edge, - ) - super().__init__(**kwargs) - self.patch_size = patch_size - self.rescale_image_frac = rescale_image_frac - self.rescale_image_max_upscaled_long_edge = rescale_image_max_upscaled_long_edge - - def _encode_one(self, image) -> torch.Tensor: - return _encode_image_bytes( - _load_image_bytes(image), - patch_size=self.patch_size, - rescale_image_frac=self.rescale_image_frac, - rescale_image_max_upscaled_long_edge=self.rescale_image_max_upscaled_long_edge, - ) - - def preprocess( - self, - images: ImageInput | list, - return_tensors: str | None = "pt", - **kwargs, - ) -> BatchFeature: - del return_tensors, kwargs - if not isinstance(images, (list, tuple)): - images = [images] - - per_image_patches: list[torch.Tensor] = [] - num_patches: list[int] = [] - num_tokens: list[int] = [] - for img in images: - vp = self._encode_one(img) - n_patches = int(vp.shape[0]) - per_image_patches.append(vp) - num_patches.append(n_patches) - num_tokens.append(n_patches) - - if len(per_image_patches) == 1: - vision_patches_bthwc = per_image_patches[0] - elif per_image_patches: - vision_patches_bthwc = torch.cat(per_image_patches, dim=0) - else: - vision_patches_bthwc = torch.empty(0) - - data = { - "vision_patches_bthwc": vision_patches_bthwc, - "num_patches": num_patches, - "num_tokens": num_tokens, - } - return BatchFeature(data=data, tensor_type=None) - - -# --------------------------------------------------------------------------- -# Audio feature extraction -# --------------------------------------------------------------------------- - - -@dataclass -class InklingAudioEncoderParams: - """Audio preprocessing parameters used to convert raw audio into dMel bins.""" - - sample_rate: int = 16_000 - window_size_multiplier: float = 2.0 - n_fft: int | None = None - n_mels: int = 80 - num_dmel_bins: int = 16 - dmel_min_value: float = -7.0 - dmel_max_value: float = 2.0 - audio_token_duration_s: float = 0.05 - - -def _to_exact_int(value: float, name: str, tolerance: float = 1e-6) -> int: - rounded = round(value) - if abs(value - rounded) > tolerance: - raise ValueError(f"{name} must resolve to an integer sample count, got {value}") - return int(rounded) - - -def _hz_to_mel(frequencies: np.ndarray) -> np.ndarray: - """Slaney mel scale, matching the librosa/torchaudio convention.""" - frequencies = np.asarray(frequencies, dtype=np.float64) - f_sp = 200.0 / 3.0 - min_log_hz = 1000.0 - min_log_mel = min_log_hz / f_sp - logstep = np.log(6.4) / 27.0 - linear = frequencies / f_sp - log = ( - min_log_mel + np.log(np.maximum(frequencies, min_log_hz) / min_log_hz) / logstep - ) - return np.where(frequencies >= min_log_hz, log, linear) - - -def _mel_to_hz(mels: np.ndarray) -> np.ndarray: - mels = np.asarray(mels, dtype=np.float64) - f_sp = 200.0 / 3.0 - min_log_hz = 1000.0 - min_log_mel = min_log_hz / f_sp - logstep = np.log(6.4) / 27.0 - linear = mels * f_sp - log = min_log_hz * np.exp(logstep * (mels - min_log_mel)) - return np.where(mels >= min_log_mel, log, linear) - - -_MEL_BASIS_CACHE: dict[tuple[int, int, int], torch.Tensor] = {} - - -def _mel_basis(sample_rate: int, n_fft: int, n_mels: int) -> torch.Tensor: - key = (sample_rate, n_fft, n_mels) - cached = _MEL_BASIS_CACHE.get(key) - if cached is not None: - return cached - - fft_bins = n_fft // 2 + 1 - fft_freqs = np.arange(fft_bins, dtype=np.float64) * sample_rate / n_fft - mel_edges = _mel_to_hz( - np.linspace( - _hz_to_mel(np.array([0.0]))[0], - _hz_to_mel(np.array([sample_rate / 2.0]))[0], - n_mels + 2, - dtype=np.float64, - ) - ) - mel_widths = np.diff(mel_edges) - lower = (fft_freqs[None, :] - mel_edges[:-2, None]) / mel_widths[:-1, None] - upper = (mel_edges[2:, None] - fft_freqs[None, :]) / mel_widths[1:, None] - weights = np.maximum(0.0, np.minimum(lower, upper)) - - # Slaney area normalization. - weights *= (2.0 / (mel_edges[2:] - mel_edges[:-2]))[:, None] - basis = torch.from_numpy(weights.astype(np.float32, copy=False)).contiguous() - _MEL_BASIS_CACHE[key] = basis - return basis - - -def _dmel_bins(audio: torch.Tensor, params: InklingAudioEncoderParams) -> torch.Tensor: - hop_length = _to_exact_int( - params.audio_token_duration_s * params.sample_rate, - "audio_token_duration_s * sample_rate", - ) - window_size = _to_exact_int( - params.audio_token_duration_s - * params.window_size_multiplier - * params.sample_rate, - "audio_token_duration_s * window_size_multiplier * sample_rate", - ) - n_fft = params.n_fft or window_size - if hop_length <= 0 or window_size <= 0 or n_fft <= 0: - raise ValueError("audio hop length, window size, and n_fft must be positive") - if audio.numel() == 0: - return torch.empty((0, params.n_mels), dtype=torch.int32) - - right_pad = math.ceil(audio.numel() / hop_length) * hop_length - audio.numel() - left_pad = max(n_fft - hop_length, 0) - audio = F.pad(audio, (left_pad, right_pad)) - - window = torch.hann_window(window_size, periodic=True, dtype=torch.float32) - spec = torch.stft( - audio.unsqueeze(0), - n_fft=n_fft, - hop_length=hop_length, - win_length=window_size, - window=window, - center=False, - normalized=False, - onesided=True, - return_complex=True, - ) - spec_ri = torch.view_as_real(spec) - magnitude = ( - (spec_ri[..., 0].square() + spec_ri[..., 1].square()) - .clamp_min(1e-10) - .sqrt() - .squeeze(0) - ) - - mel = ( - _mel_basis(params.sample_rate, n_fft, params.n_mels) - .matmul(magnitude) - .clamp_min(1e-10) - .log10() - ) - mel = mel.to(torch.float64).clamp( - min=params.dmel_min_value, max=params.dmel_max_value - ) - bin_centers = torch.linspace( - params.dmel_min_value, - params.dmel_max_value, - params.num_dmel_bins, - dtype=torch.float64, - ) - dmel_bins = (mel.unsqueeze(-1) - bin_centers).abs().argmin(dim=-1) - return dmel_bins.to(torch.int32).T.contiguous() - - -class InklingAudioFeatureExtractor(FeatureExtractionMixin): - """Convert raw audio into Inkling dMel bins in the HF feature-extractor API.""" - - model_input_names = ["dmel_bins"] - - def __init__(self, params: dict | None = None, **kwargs): - super().__init__(**kwargs) - merged = InklingAudioEncoderParams() - if params: - for k, v in params.items(): - if hasattr(merged, k): - setattr(merged, k, v) - # also accept flat kwargs (HF config style) - for k in list(kwargs.keys()): - if hasattr(merged, k): - setattr(merged, k, kwargs[k]) - self.params = merged - - def _decode_one(self, audio) -> torch.Tensor: - # vLLM hands the feature extractor numpy arrays (the dummy-input builder - # during profiling, and MultiModalDataParser after resampling to the - # target sample rate), so no other input types need to be supported. - return torch.from_numpy( - np.ascontiguousarray(audio.astype(np.float32, copy=False)) - ).flatten() - - def _encode_one(self, audio) -> torch.Tensor: - return _dmel_bins(self._decode_one(audio), self.params) - - def __call__( - self, - audios: Sequence | None, - return_tensors: str | None = None, - **kwargs, - ) -> BatchFeature: - del return_tensors, kwargs - if audios is None: - audios = [] - if not isinstance(audios, (list, tuple)): - audios = [audios] - - dmel_bins = [self._encode_one(a) for a in audios] - data = { - # per-clip feature: dmel bins as float32 [T, n_mels] - "dmel_bins": [bins.to(torch.float32) for bins in dmel_bins], - "num_audio_tokens": [int(bins.shape[0]) for bins in dmel_bins], - } - # return_tensors intentionally ignored: per-clip features have ragged T. - return BatchFeature(data=data, tensor_type=None) - - -# --------------------------------------------------------------------------- -# Composite processor -# --------------------------------------------------------------------------- - - -class InklingProcessor: - """Bundle Inkling image + audio preprocessing with the MM token ids.""" - - def __init__( - self, - image_processor: InklingImageProcessor | None = None, - audio_feature_extractor: InklingAudioFeatureExtractor | None = None, - tokenizer=None, - ): - self.image_processor = image_processor or InklingImageProcessor() - self.audio_feature_extractor = ( - audio_feature_extractor or InklingAudioFeatureExtractor() - ) - self.tokenizer = tokenizer - - def process_images(self, images: list): - """Raw images -> BatchFeature(vision_patches_bthwc, num_patches, num_tokens).""" - return self.image_processor.preprocess(images, return_tensors="pt") - - def process_audios(self, audios: list): - """Raw audios -> BatchFeature(dmel_bins, num_audio_tokens).""" - return self.audio_feature_extractor(audios) - - -__all__ = [ - "InklingImageProcessor", - "InklingAudioFeatureExtractor", - "InklingAudioEncoderParams", - "InklingProcessor", - "IMAGE_MARKER_ID", - "AUDIO_MARKER_ID", - "IMAGE_TOKEN_ID", - "AUDIO_TOKEN_ID", -] From 6b6a6c5c98e4a23c77bb07d344c515f48c157193 Mon Sep 17 00:00:00 2001 From: pkousha <43781676+pkousha@users.noreply.github.com> Date: Fri, 16 Oct 2026 18:41:17 -0700 Subject: [PATCH 34/49] [RL] Support per-parameter source layouts for NCCL M2N (#60934) Signed-off-by: pkousha Co-authored-by: OpenAI Codex Co-authored-by: aoshen02 --- docs/training/weight_transfer/m2n.md | 31 +- examples/rl/rlhf_m2n.py | 35 +- tests/distributed/test_weight_transfer_m2n.py | 1223 ++++++++++++++--- .../distributed/weight_transfer/m2n_common.py | 766 +++++++++-- .../distributed/weight_transfer/m2n_engine.py | 432 ++++-- .../distributed/weight_transfer/m2n_source.py | 124 +- .../weight_transfer/m2n_trainer.py | 619 ++++++--- .../weight_transfer/nccl_common.py | 55 +- 8 files changed, 2590 insertions(+), 695 deletions(-) diff --git a/docs/training/weight_transfer/m2n.md b/docs/training/weight_transfer/m2n.md index 61510978c8..a0e4e8089b 100644 --- a/docs/training/weight_transfer/m2n.md +++ b/docs/training/weight_transfer/m2n.md @@ -10,18 +10,14 @@ The `nccl_m2n` weight transfer engine uses [NCCL M2N](https://github.com/NVIDIA/ ## Requirements -This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. It needs NCCL 2.30.5 or newer, and `libnccl_m2n.so` is linked against a specific NCCL build, so vLLM must load the same library: - -```bash -export VLLM_NCCL_SO_PATH=/path/to/the/same/libnccl.so -``` +This backend depends on the `nccl-extensions` Python package (and its `nccl4py` dependency), which vLLM does not install. Install both against a compatible NCCL 2.30.5 or newer runtime. The backend creates its data-plane communicator through `nccl4py` and passes that communicator directly to `nccl.m2n`. The backend imports the runtime lazily, so vLLM is unaffected unless you select it. ## How It Works -1. The trainer ranks and all inference workers join **one** NCCL communicator via `StatelessProcessGroup`. Trainer ranks occupy `[0, T)` and workers `[T, T + N)` — m2n meshes are contiguous rank intervals, which is exactly what vLLM's existing `rank_offset` convention produces. -2. At the init handshake the trainer ships the full transfer plan: both meshes — one per side, shared by every parameter — then each parameter's name, dtype, shape and placement. The workers resolve the matching **destination placements** and publish them back over the shared group; the meshes themselves are declared, not negotiated. +1. A `StatelessProcessGroup` exchanges metadata and the NCCL unique ID. The trainer ranks and all inference workers then create **one** `nccl.core.Communicator`: trainer ranks occupy `[0, T)` and workers `[T, T + N)`. The communicator is passed directly to each M2N collective. +2. At the init handshake the trainer ships the full transfer plan: the session-wide destination mesh, then each parameter's name, dtype, shape, source mesh, and source placements. The workers resolve the matching **destination placements** and publish them over the data-plane communicator; the meshes themselves are declared, not negotiated. 3. Each round, both sides issue one `reshard` per parameter in the same order. Every rank in the communicator participates in every reshard, so — as with the NCCL engine — the worker must be inside `update_weights` while the trainer sends. The trainer engine handles that concurrency itself. ### Destination resolution @@ -83,19 +79,28 @@ engine.send_weights() `send_weights()` drives the whole round — `start_weight_update`, `update_weights`, `finish_weight_update`. Every rank in the communicator must be inside the same reshard, so `update_weights` has to be in flight while the trainer sends; the engine runs it on a helper thread so the trainer loop does not have to open-code that. -Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass `M2NWeightSource` and supply layouts explicitly. +Trainers with a custom producer (a Megatron export, MoE re-fusing) subclass +`M2NWeightSource` and supply layouts explicitly. A source may yield any +nccl4py-compatible CUDA buffer: a torch tensor, a CUDA Array Interface object, +or a raw device pointer. vLLM keeps buffer owners alive through stream +completion, using allocator stream tracking when the buffer provides it. +Sources must provide buffers on the communicator's current device; nccl4py +resolves their pointers but does not verify device ownership. Raw-pointer +sources are additionally responsible for pointer lifetime and matching their +declared dtype and local shape. See [`examples/rl/rlhf_m2n.py`](../../../examples/rl/rlhf_m2n.py) for a runnable FSDP → TP example. ## Limitations -- Tensor rank 1..3; 4-D parameters are rejected. -- dtypes: int8/uint8, fp8 e4m3/e5m2, fp16, bf16, int32/uint32, fp32, int64/uint64, fp64. No fp4. -- The trainer's device mesh must cover a contiguous rank interval starting at 0, and be 1-D or 2-D. +- vLLM validates tensor-rank and dtype support while agreeing on the complete + plan; the installed M2N runtime validates them again for each reshard. +- Each parameter's source mesh must cover the same contiguous trainer-rank interval starting at 0, and be 1-D or 2-D. - `Partial` placements are not supported. - Models with known model-level checkpoint preprocessing use the full-tensor fallback for the affected module. - A reshard cannot be issued from inside a CUDA graph capture (weight updates run outside capture, so this only matters for custom callers). -- The reshard plan is capped per *shard*, not per mesh: at most 16 source shards may feed one destination shard, and one source shard may feed at most 64 destination shards (`MAX_SOURCES` / `MAX_TARGETS` in the m2n build; raising either needs a rebuild). -- A directly-resharded parameter spreads over `tp_size` destination shards, so each one is fed by fewer source shards. A parameter on the fallback path is replicated, i.e. a single destination shard fed by every source shard, so a trainer that **sharded** its parameters is still limited to 16 ranks whenever any parameter falls back — which fused parameters always do. A replicated trainer is unaffected. +- vLLM checks every parameter's shard overlap against conservative 16-source + and 64-target capacities before committing the plan. The installed M2N + runtime remains authoritative and validates each reshard again. - Each communicator holds a cached staging pool (~256 MiB by default: 4 channels x 64 MiB). diff --git a/examples/rl/rlhf_m2n.py b/examples/rl/rlhf_m2n.py index e02229ed09..6fc64463e1 100644 --- a/examples/rl/rlhf_m2n.py +++ b/examples/rl/rlhf_m2n.py @@ -17,9 +17,8 @@ Every FSDP rank builds an ``M2NTrainerWeightTransferEngine`` and calls ``send_weights()``; all ranks run every reshard, and only rank 0 drives the inference side through its ``RayVLLMWeightSyncClient``. -Requires the ``nccl-extensions`` package (NCCL 2.30.5+) and a -``VLLM_NCCL_SO_PATH`` pointing at the same ``libnccl.so`` that -``libnccl_m2n.so`` was linked against. +Requires the ``nccl-extensions`` package (NCCL 2.30.5+). The backend creates +the data-plane communicator with ``nccl4py`` and passes it directly to M2N. This example was written for 4xH100. """ @@ -28,6 +27,7 @@ from __future__ import annotations import os +import nccl.core as nccl import pybase64 as base64 import ray import torch @@ -38,7 +38,6 @@ from transformers import AutoModelForCausalLM from vllm import LLM, SamplingParams from vllm.config import WeightTransferConfig -from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary from vllm.distributed.weight_transfer import ( RayVLLMWeightSyncClient, WeightTransferTrainerFactory, @@ -94,7 +93,15 @@ class FSDPTrainWorker: def ready(self): return True - def setup_engine(self, llm_handle, nccl_unique_id_b64, world_size, num_workers): + def setup_engine( + self, + llm_handle, + metadata_address, + metadata_port, + nccl_unique_id_b64, + world_size, + num_workers, + ): """Build the trainer engine on every FSDP rank. `DTensorModuleSource` reads each parameter's FSDP device mesh and @@ -103,6 +110,8 @@ class FSDPTrainWorker: """ self.engine = WeightTransferTrainerFactory.trainer_init( init_info=M2NTrainerInitInfo( + master_address=metadata_address, + master_port=metadata_port, nccl_unique_id_b64=nccl_unique_id_b64, world_size=world_size, num_trainer_ranks=FSDP_WORLD_SIZE, @@ -166,16 +175,22 @@ def main(): # Trainer ranks [0, FSDP_WORLD_SIZE) and inference ranks after them share # one communicator, so the two meshes are contiguous rank intervals. - nccl = NCCLLibrary() - nccl_unique_id_b64 = base64.b64encode( - bytes(nccl.ncclGetUniqueId().internal) - ).decode() + nccl_unique_id_b64 = base64.b64encode(nccl.get_unique_id().as_bytes).decode() world_size = FSDP_WORLD_SIZE + num_workers + metadata_address = get_ip() + metadata_port = get_open_port() print("[transfer] Initializing nccl_m2n weight transfer (all FSDP ranks)...") ray.get( [ - w.setup_engine.remote(llm, nccl_unique_id_b64, world_size, num_workers) + w.setup_engine.remote( + llm, + metadata_address, + metadata_port, + nccl_unique_id_b64, + world_size, + num_workers, + ) for w in fsdp_workers ] ) diff --git a/tests/distributed/test_weight_transfer_m2n.py b/tests/distributed/test_weight_transfer_m2n.py index 588ffeeb26..991852c034 100644 --- a/tests/distributed/test_weight_transfer_m2n.py +++ b/tests/distributed/test_weight_transfer_m2n.py @@ -2,36 +2,44 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Unit tests for the NCCL M2N weight transfer backend. -These cover how a transfer is *described* — layout encoding and the wire-type -validation that keeps a bad plan from reaching a collective, since a mismatch -there is a hang rather than an error. The transfer itself needs the m2n runtime -and multiple GPUs, so it is exercised separately. +These cover how a transfer is *described* — layout encoding, whole-plan +validation, and the vLLM-owned source contract. The transfer itself needs the +M2N runtime and multiple GPUs, so it is exercised separately. """ import logging +from contextlib import contextmanager from functools import wraps -from types import MethodType -from unittest.mock import Mock +from types import MethodType, SimpleNamespace +from unittest.mock import Mock, call import pybase64 as base64 import pytest import ray import torch +from vllm.config.weight_transfer import WeightTransferConfig from vllm.distributed.weight_transfer import ( WeightTransferEngineFactory, WeightTransferTrainerFactory, ) from vllm.distributed.weight_transfer.m2n_common import ( + M2N_WIRE_SCHEMA_VERSION, REPLICATE, REPLICATED, + M2NLayout, M2NMesh, M2NParamMeta, - check_placements, - check_transferable, - publish_destination_placements, + M2NWireParam, + check_complete_plan_agreement, + check_data_plane_agreement, + check_plan_limits, + check_runtime_ready_agreement, + check_source_plan_agreement, + create_m2n_communicator, + expected_local_shape, resolve_layout, - validate_layout, + source_plan_digest, ) from vllm.distributed.weight_transfer.m2n_engine import ( M2NWeightTransferEngine, @@ -43,6 +51,9 @@ from vllm.distributed.weight_transfer.m2n_layout import ( resolve_parameter_destinations, ) from vllm.distributed.weight_transfer.m2n_source import ( + M2NManifestEntry, + M2NWeightSource, + ManifestM2NWeightSource, mesh_from_tensor, placements_from_tensor, ) @@ -50,14 +61,16 @@ from vllm.distributed.weight_transfer.m2n_trainer import ( M2NTrainerInitInfo, M2NTrainerWeightTransferEngine, ) +from vllm.distributed.weight_transfer.nccl_common import ( + stateless_init_metadata_group, +) from vllm.model_executor.layers.linear import ( ColumnParallelLinear, RowParallelLinear, ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.platforms import current_platform - -VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() +from vllm.utils.network_utils import get_open_port @pytest.mark.parametrize( @@ -67,18 +80,114 @@ VALID_UID_B64 = base64.b64encode(b"\x00" * 128).decode() (M2NTrainerWeightTransferEngine, "group"), ], ) -def test_shutdown_destroys_communicator_once(engine_type, group_attribute): +def test_shutdown_releases_resources_once(monkeypatch, engine_type, group_attribute): engine = object.__new__(engine_type) - engine._handle = None - engine._executor = None + handle = Mock() + engine._handle = handle group = Mock() setattr(engine, group_attribute, group) + synchronize = Mock() + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + + if engine_type is M2NWeightTransferEngine: + engine._parameter_destinations = [Mock()] + else: + engine._dst_placements = [REPLICATED] + engine._metadata_group = Mock() engine.shutdown() engine.shutdown() - group.destroy.assert_called_once_with() + synchronize.assert_called_once_with() + handle.destroy.assert_called_once_with() + group.abort.assert_called_once_with() + group.destroy.assert_not_called() + assert engine._handle is None assert getattr(engine, group_attribute) is None + if engine_type is M2NWeightTransferEngine: + assert engine._parameter_destinations == [] + else: + assert engine._dst_placements == [] + assert engine._metadata_group is None + + +@pytest.mark.parametrize( + "engine_type,group_attribute", + [ + (M2NWeightTransferEngine, "model_update_group"), + (M2NTrainerWeightTransferEngine, "group"), + ], +) +@pytest.mark.parametrize("failure", ["synchronize", "handle"]) +def test_shutdown_attempts_all_cleanup_after_failure( + monkeypatch, engine_type, group_attribute, failure +): + engine = object.__new__(engine_type) + handle = Mock() + group = Mock() + engine._handle = handle + setattr(engine, group_attribute, group) + if engine_type is M2NTrainerWeightTransferEngine: + engine._metadata_group = Mock() + + synchronize = Mock() + if failure == "synchronize": + synchronize.side_effect = RuntimeError("synchronize failed") + else: + handle.destroy.side_effect = RuntimeError("handle destroy failed") + monkeypatch.setattr(torch.accelerator, "synchronize", synchronize) + + with pytest.raises(RuntimeError, match=failure): + engine.shutdown() + + synchronize.assert_called_once_with() + handle.destroy.assert_called_once_with() + group.abort.assert_called_once_with() + assert engine._handle is None + assert getattr(engine, group_attribute) is None + if engine_type is M2NTrainerWeightTransferEngine: + assert engine._metadata_group is None + + engine.shutdown() + synchronize.assert_called_once_with() + handle.destroy.assert_called_once_with() + group.abort.assert_called_once_with() + + +class _BufferLike: + """Minimal nccl4py buffer protocol without torch-only methods.""" + + def __init__(self, dtype=torch.float32, shape=(4,)): + self.dtype = dtype + self.shape = shape + + def data_ptr(self): + return 1234 + + +class _AllocatorAwareBuffer(_BufferLike): + def __init__(self): + super().__init__() + self.record_stream = Mock() + + +class _FakeDistTensor: + def __init__(self, buffer, local_shape=None, dtype=None, **kwargs): + del kwargs + self.local_shape = ( + tuple(local_shape) if local_shape is not None else tuple(buffer.shape) + ) + self.dtype = dtype if dtype is not None else buffer.dtype + + +def _fake_m2n(): + return SimpleNamespace( + DistTensor=_FakeDistTensor, + Mesh=lambda dims, start_rank=0: (tuple(dims), start_rank), + Replicate=lambda: REPLICATE, + Shard=lambda dim: dim, + reshard=Mock(), + ) class TestLayout: @@ -87,10 +196,10 @@ class TestLayout: dims are deliberately coprime with the mesh size: a layout that actually split the tensor would reject them.""" mesh = M2NMesh((2, 2), start_rank=1) - indivisible_shape = (7, 13) # neither dim divisible by any mesh axis resolved_mesh, placements = resolve_layout(mesh, REPLICATED) - validate_layout(resolved_mesh, placements, indivisible_shape, "destination") + assert resolved_mesh.dims == (4, 1) + assert placements == (REPLICATE, 0) def test_replicated_keeps_the_same_ranks(self): """Replication re-factors the mesh to get a size-1 axis for its no-op @@ -108,23 +217,27 @@ class TestLayout: assert resolved_mesh == mesh assert placements == (REPLICATE, 0) - def test_two_shard_axes_rejected(self): - """One axis has to replicate; a 2-D mesh that shards both is not - something a single reshard can express.""" - with pytest.raises(ValueError, match="shards both"): - check_placements((0, 1)) + @pytest.mark.parametrize( + ("placements", "match"), + [((REPLICATE, 1), "tensor dim 1"), ((REPLICATE, 0), "divide evenly")], + ) + def test_layout_validation_always_returns_shape_or_raises(self, placements, match): + meta = M2NParamMeta( + "w", + torch.float32, + (5,), + M2NLayout(M2NMesh((1, 2), 0), placements), + ) - def test_placement_below_replicate_rejected(self): - with pytest.raises(ValueError, match=r"parameter 'w'.*-2"): - check_placements((-2, REPLICATE), "parameter 'w' source placements") + with pytest.raises(ValueError, match=match): + expected_local_shape(meta) - def test_shard_dim_must_exist(self): - with pytest.raises(ValueError, match="rank 2"): - validate_layout(M2NMesh((1, 2), 0), (REPLICATE, 2), (8, 16), "source") + def test_cross_sharded_fanout_over_capacity_is_rejected(self): + src = (M2NMesh((1, 2), 0), (REPLICATE, 0)) + dst = (M2NMesh((1, 65), 2), (REPLICATE, 1)) - def test_shard_must_divide_evenly(self): - with pytest.raises(ValueError, match="does not divide evenly"): - validate_layout(M2NMesh((1, 3), 0), (REPLICATE, 0), (8, 16), "source") + with pytest.raises(ValueError, match=r"parameter 'w'.*MAX_TARGETS=64"): + check_plan_limits(src, dst, "w", (2, 65)) class TestSourceLayout: @@ -134,76 +247,398 @@ class TestSourceLayout: assert placements_from_tensor(torch.zeros(4)) is REPLICATED assert mesh_from_tensor(torch.zeros(4), 4) == M2NMesh((4, 1), 0) + def test_manifest_preserves_order_and_materializes_lazily(self): + calls: list[int] = [] -class TestTransferable: - def test_unsupported_dtype_names_the_parameter(self): - with pytest.raises(ValueError, match="'w'"): - check_transferable("w", torch.complex64, (4,)) + def provider(index, value): + def get_tensor(): + calls.append(index) + return value - def test_rank_four_rejected(self): - with pytest.raises(ValueError, match="rank 4"): - check_transferable("w", torch.bfloat16, (2, 2, 2, 2)) + return get_tensor + + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "second", + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + provider(0, torch.ones(4)), + ), + M2NManifestEntry( + "first", + torch.float32, + (2,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + provider(1, torch.full((2,), 2.0)), + ), + ] + ) + + assert [meta.name for meta in source.metadata()] == ["second", "first"] + assert calls == [] + assert [name for name, _ in source] == ["second", "first"] + assert calls == [0, 1] + + def test_trainer_boundary_rejects_duplicate_manifest_names(self): + entry = M2NManifestEntry( + "w", + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + lambda: torch.zeros(4), + ) + source = ManifestM2NWeightSource([entry, entry]) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match="duplicate parameter 'w'"): + engine._prepare_source_plan(source, num_trainer_ranks=1) + + @pytest.mark.parametrize( + ("dtype", "shape"), + [(torch.complex64, (4,)), (torch.float32, (2, 2, 2, 2))], + ) + def test_transport_capabilities_rejected_during_initialization(self, dtype, shape): + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "w", + dtype, + shape, + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + lambda: _BufferLike(dtype, shape), + ) + ] + ) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match=r"parameter 'w'.*(dtype|rank)"): + engine._prepare_source_plan(source, num_trainer_ranks=1) + + @pytest.mark.parametrize("placements", [(REPLICATE, 1), (REPLICATE, 0)]) + def test_layout_geometry_rejected_during_initialization(self, placements): + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "w", + torch.float32, + (5,), + M2NLayout(M2NMesh((1, 2), 0), placements), + lambda: _BufferLike(shape=(5,)), + ) + ] + ) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match=r"parameter 'w'.*source layout"): + engine._prepare_source_plan(source, num_trainer_ranks=2) + + def test_raw_pointer_invalid_layout_rejected_during_initialization(self): + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "pointer", + torch.float32, + (5,), + M2NLayout(M2NMesh((1, 2), 0), (REPLICATE, 0)), + lambda: 1234, + ) + ] + ) + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + + with pytest.raises(ValueError, match=r"parameter 'pointer'.*divide evenly"): + engine._prepare_source_plan(source, num_trainer_ranks=2) + + +class TestPreflightAgreement: + def test_source_error_is_reported_before_data_plane(self): + group = Mock(world_size=2) + group.all_gather_obj.return_value = [ + { + "phase": "source", + "digest": None, + "declared_digest": None, + "error": "ValueError: bad local manifest", + }, + { + "phase": "source", + "digest": None, + "declared_digest": None, + "error": "ValueError: bad local manifest", + }, + ] + + with pytest.raises(RuntimeError, match="bad local manifest"): + check_source_plan_agreement( + group, [], local_error="ValueError: bad local manifest" + ) + + def test_data_plane_identity_must_match(self): + group = Mock(world_size=2) + + def gather(envelope): + peer = dict(envelope) + peer["mode"] = "uid" + peer["uid_digest"] = "different" + return [envelope, peer] + + group.all_gather_obj.side_effect = gather + with pytest.raises(RuntimeError, match="data-plane identity disagrees"): + check_data_plane_agreement(group, None) + + def test_runtime_error_is_shared_before_nccl4py(self): + group = Mock(world_size=2) + group.all_gather_obj.return_value = [ + {"phase": "runtime_ready", "error": None}, + {"phase": "runtime_ready", "error": "ImportError: no nccl_m2n"}, + ] + + with pytest.raises(RuntimeError, match="no nccl_m2n"): + check_runtime_ready_agreement(group, None) + + def test_complete_plan_error_is_shared_before_first_reshard(self): + group = Mock(world_size=2) + group.all_gather_obj.return_value = [ + {"phase": "complete_plan", "digest": "same", "error": None}, + { + "phase": "complete_plan", + "digest": None, + "error": "ValueError: bad destination", + }, + ] + + with pytest.raises(RuntimeError, match="bad destination"): + check_complete_plan_agreement( + group, + "source", + M2NMesh((1, 1), 1), + [REPLICATED], + ) + + def test_complete_plan_digest_must_match(self): + group = Mock(world_size=2) + + def gather(envelope): + peer = dict(envelope) + peer["digest"] = "different" + return [envelope, peer] + + group.all_gather_obj.side_effect = gather + with pytest.raises(RuntimeError, match="complete-plan digest disagrees"): + check_complete_plan_agreement( + group, + "source", + M2NMesh((1, 1), 1), + [REPLICATED], + ) class TestWireTypes: + @staticmethod + def _param( + name="w", + dtype_name="bfloat16", + shape=(16, 16), + src_mesh_dims=(1, 1), + src_placements=None, + ): + return M2NWireParam( + name=name, + dtype_name=dtype_name, + shape=shape, + src_mesh_dims=src_mesh_dims, + src_placements=src_placements, + ) + def _init_info(self, **overrides): + param = self._param() fields = dict( + schema_version=M2N_WIRE_SCHEMA_VERSION, master_address="127.0.0.1", master_port=1234, rank_offset=1, world_size=3, - src_mesh_dims=[1, 1], dst_mesh_dims=[2, 1], - names=["w"], - dtype_names=["bfloat16"], - shapes=[[16, 16]], - src_placements=[None], + source_digest=source_plan_digest([param]), + params=[param.to_dict()], ) fields.update(overrides) + if "params" in overrides and "source_digest" not in overrides: + params = [M2NWireParam.from_dict(value) for value in fields["params"]] + fields["source_digest"] = source_plan_digest(params) return M2NWeightTransferInitInfo(**fields) def test_accepts_a_consistent_plan(self): - assert self._init_info().names == ["w"] + assert [param.name for param in self._init_info().parse_wire_params()] == ["w"] - def test_accepts_pre_shared_nccl_unique_id(self): + def test_m2n_communicator_uses_nccl4py_uid_under_unpinned_env(self, monkeypatch): + nccl = Mock() + group = Mock(rank=1, world_size=3) + unique_id = object() + communicator = object() + nccl.UniqueId.from_bytes.return_value = unique_id + active = False + + @contextmanager + def unpinned(): + nonlocal active + active = True + try: + yield + finally: + active = False + + def init(**kwargs): + assert active + assert kwargs == {"nranks": 3, "rank": 1, "unique_id": unique_id} + return communicator + + nccl.Communicator.init.side_effect = init + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_common.unpinned_nccl_env", + unpinned, + ) + + result = create_m2n_communicator(nccl, group, b"shared-id") + + assert result is communicator + nccl.UniqueId.from_bytes.assert_called_once_with(b"shared-id") + nccl.Communicator.init.assert_called_once_with( + nranks=3, rank=1, unique_id=unique_id + ) + group.broadcast_obj.assert_not_called() + + @pytest.mark.parametrize( + "mutation", + [ + lambda value: value.pop("shape"), + lambda value: value.update(extra="field"), + lambda value: value.update(src_mesh_dims=[1]), + lambda value: value.update(src_placements=[0, 1]), + ], + ) + def test_malformed_wire_parameter_rejected(self, mutation): + value = self._param().to_dict() + mutation(value) + with pytest.raises((TypeError, ValueError)): + M2NWireParam.from_dict(value) + + def test_source_digest_authenticates_ordered_plan(self): + params = [self._param("a"), self._param("b")] info = self._init_info( - master_address=None, - master_port=None, - nccl_unique_id_b64=VALID_UID_B64, + params=[param.to_dict() for param in params], + source_digest=source_plan_digest(list(reversed(params))), ) - assert info.nccl_unique_id_bytes == b"\x00" * 128 - assert VALID_UID_B64 not in repr(info) + with pytest.raises(ValueError, match="source digest"): + info.parse_wire_params() - def test_destination_plan_uses_uid_communicator_without_group(self): - comm = Mock(rank=0, device=torch.device("cpu"), group=None) - - def receive_plan(tensor, src): - assert src == 1 - tensor.copy_(torch.tensor([[-1, 0], [-2, -2]], dtype=torch.int8)) - - comm.broadcast.side_effect = receive_plan - - placements = publish_destination_placements(comm, 1, None, 2) - - assert placements == [(REPLICATE, 0), REPLICATED] - - def test_rejects_both_rendezvous_modes(self): - with pytest.raises(ValueError, match="not both"): - self._init_info(nccl_unique_id_b64=VALID_UID_B64) - - def test_trainer_accepts_pre_shared_nccl_unique_id(self): - info = M2NTrainerInitInfo( - nccl_unique_id_b64=VALID_UID_B64, - world_size=3, - num_trainer_ranks=1, - rank=0, + @pytest.fixture + def worker_init_env(self, monkeypatch): + metadata_group = Mock(rank=1, world_size=3) + metadata_group.all_gather_obj.side_effect = lambda envelope: [envelope] * 3 + init_metadata = Mock(return_value=metadata_group) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.worker_init_metadata_group", + init_metadata, + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.import_m2n", Mock() + ) + handle = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.prepare_m2n_local_runtime", + Mock(return_value=(Mock(), 0, handle)), + ) + communicator = Mock() + make_communicator = Mock(return_value=communicator) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine.create_m2n_communicator", + make_communicator, + ) + publish = Mock(return_value=[]) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "publish_destination_placements", + publish, + ) + vllm_config = SimpleNamespace( + parallel_config=Mock(pipeline_parallel_size=1), + model_config=SimpleNamespace(quantization=None), + ) + engine = M2NWeightTransferEngine( + WeightTransferConfig(backend="nccl_m2n"), + vllm_config, + torch.device("cuda:0"), + torch.nn.Module(), + ) + return SimpleNamespace( + engine=engine, + metadata_group=metadata_group, + init_metadata=init_metadata, + handle=handle, + communicator=communicator, + make_communicator=make_communicator, + publish=publish, ) - assert info.nccl_unique_id_bytes == b"\x00" * 128 - assert VALID_UID_B64 not in repr(info) - def test_ragged_plan_rejected(self): - with pytest.raises(ValueError, match="`shapes`"): - self._init_info(shapes=[]) + def test_trainer_disagreement_is_rejected_before_nccl4py(self, worker_init_env): + group = worker_init_env.metadata_group + + def gather(envelope): + peer = dict(envelope) + peer["digest"] = "different" + return [envelope, peer, envelope] + + group.all_gather_obj.side_effect = gather + with pytest.raises(RuntimeError, match="source digest disagrees"): + worker_init_env.engine.init_transfer_engine(self._init_info()) + + worker_init_env.make_communicator.assert_not_called() + + def test_worker_rejects_active_reinit_but_allows_init_after_shutdown( + self, monkeypatch, worker_init_env + ): + engine = worker_init_env.engine + old_handle = Mock() + old_group = Mock() + engine._handle = old_handle + engine.model_update_group = old_group + + with pytest.raises(RuntimeError, match="already initialized"): + engine.init_transfer_engine(self._init_info(params=[])) + worker_init_env.init_metadata.assert_not_called() + + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + engine.shutdown() + + engine.init_transfer_engine(self._init_info(params=[])) + + assert engine._handle is worker_init_env.handle + assert engine.model_update_group is worker_init_env.communicator + + def test_partial_worker_initialization_releases_transport_resources( + self, worker_init_env + ): + worker_init_env.publish.side_effect = RuntimeError("late init failure") + + with pytest.raises(RuntimeError, match="late init failure"): + worker_init_env.engine.init_transfer_engine(self._init_info(params=[])) + + worker_init_env.communicator.abort.assert_called_once_with() + worker_init_env.communicator.destroy.assert_not_called() + worker_init_env.handle.destroy.assert_called_once_with() + assert worker_init_env.engine.model_update_group is None + assert worker_init_env.engine._handle is None def test_destination_mesh_must_cover_the_workers(self): """The trainer declares the inference mesh, so one that does not cover @@ -217,69 +652,140 @@ class TestWireTypes: self._init_info(rank_offset=3, world_size=3) @pytest.mark.parametrize("dtype_name", ["not_a_dtype", "Tensor"]) - def test_invalid_dtype_name_names_parameter(self, monkeypatch, dtype_name): - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.import_m2n", - lambda: object(), - ) + def test_invalid_dtype_name_names_parameter(self, dtype_name): engine = object.__new__(M2NWeightTransferEngine) + param = self._param(dtype_name=dtype_name) with pytest.raises(ValueError, match=r"parameter 'w'.*dtype"): - engine.init_transfer_engine(self._init_info(dtype_names=[dtype_name])) + engine._prepare_source_plan((param,), num_trainer_ranks=1) def _init_32_to_4_plan(self, monkeypatch, destination): - m2n = Mock() - m2n.Handle.create.return_value = object() - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.import_m2n", lambda: m2n - ) monkeypatch.setattr( "vllm.distributed.weight_transfer.m2n_engine." "resolve_parameter_destinations", lambda *args, **kwargs: [destination], ) - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.worker_init_process_group", - lambda *args, **kwargs: object(), - ) - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine." - "publish_destination_placements", - lambda *args: args[2], - ) engine = object.__new__(M2NWeightTransferEngine) engine.parallel_config = Mock(pipeline_parallel_size=1) engine.model_config = Mock(quantization=None) engine.model = torch.nn.Module() - engine.init_transfer_engine( + param = self._param( + shape=(32, 8), + src_mesh_dims=(1, 32), + src_placements=(REPLICATE, 0), + ) + engine._prepare_source_plan((param,), num_trainer_ranks=32) + engine._prepare_destination_plan( self._init_info( rank_offset=32, world_size=36, - src_mesh_dims=[1, 32], dst_mesh_dims=[1, 4], - shapes=[[32, 8]], - src_placements=[[REPLICATE, 0]], + params=[param.to_dict()], ) ) - def test_limits_use_resolved_sharded_destination(self, monkeypatch): + def test_sharded_destination_within_capacity_initializes(self, monkeypatch): destination = M2NDestination("w", (REPLICATE, 0), torch.empty(8, 8)) self._init_32_to_4_plan(monkeypatch, destination) - def test_limits_still_reject_replicated_destination(self, monkeypatch): + def test_cross_sharded_destination_over_capacity_rejected_at_init( + self, monkeypatch + ): + destination = M2NDestination("w", (REPLICATE, 1), torch.empty(32, 2)) + + with pytest.raises(ValueError, match=r"parameter 'w'.*MAX_SOURCES=16"): + self._init_32_to_4_plan(monkeypatch, destination) + + def test_replicated_destination_over_capacity_rejected_at_init(self, monkeypatch): destination = M2NDestination("w", REPLICATED, None) - with pytest.raises(ValueError, match=r"32 source shards.*MAX_SOURCES=16"): + with pytest.raises(ValueError, match=r"parameter 'w'.*MAX_SOURCES=16"): self._init_32_to_4_plan(monkeypatch, destination) + def test_invalid_destination_placement_rejected_at_init(self, monkeypatch): + destination = M2NDestination("w", (REPLICATE, 2), torch.empty(8, 8)) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "resolve_parameter_destinations", + lambda *args, **kwargs: [destination], + ) + engine = object.__new__(M2NWeightTransferEngine) + engine.parallel_config = Mock(pipeline_parallel_size=1) + engine.model_config = Mock(quantization=None) + engine.model = torch.nn.Module() + param = self._param(shape=(32, 8)) + engine._prepare_source_plan((param,), num_trainer_ranks=1) + + with pytest.raises( + ValueError, match=r"parameter 'w'.*destination layout.*tensor dim 2" + ): + engine._prepare_destination_plan( + self._init_info( + rank_offset=1, + world_size=5, + dst_mesh_dims=[1, 4], + params=[param.to_dict()], + ) + ) + + def test_later_invalid_parameter_rejects_plan_before_any_reshard(self, monkeypatch): + destinations = [ + M2NDestination("valid", (REPLICATE, 0), torch.empty(8, 8)), + M2NDestination("invalid", (REPLICATE, 1), torch.empty(32, 2)), + ] + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_engine." + "resolve_parameter_destinations", + lambda *args, **kwargs: destinations, + ) + engine = object.__new__(M2NWeightTransferEngine) + engine.parallel_config = Mock(pipeline_parallel_size=1) + engine.model_config = Mock(quantization=None) + engine.model = torch.nn.Module() + params = ( + self._param( + name="valid", + shape=(32, 8), + src_mesh_dims=(8, 4), + src_placements=(REPLICATE, 0), + ), + self._param( + name="invalid", + shape=(32, 8), + src_mesh_dims=(1, 32), + src_placements=(REPLICATE, 0), + ), + ) + engine._prepare_source_plan(params, num_trainer_ranks=32) + engine._reshard = Mock() + + with pytest.raises(ValueError, match=r"parameter 'invalid'.*MAX_SOURCES=16"): + engine._prepare_destination_plan( + self._init_info( + rank_offset=32, + world_size=36, + dst_mesh_dims=[1, 4], + params=[param.to_dict() for param in params], + ) + ) + + engine._reshard.assert_not_called() + def test_update_preflights_all_names_before_reshard(self): engine = object.__new__(M2NWeightTransferEngine) engine._handle = object() engine.model_update_group = object() engine._index = {"valid": 0} - engine._metas = [M2NParamMeta("valid", torch.float32, (4,), REPLICATED)] + engine._metas = [ + M2NParamMeta( + "valid", + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + ) + ] engine._reshard = Mock() with pytest.raises(ValueError, match=r"parameter 'unknown'"): @@ -310,7 +816,12 @@ class TestWireTypes: engine.model = model engine._index = {"pair.weight": 0, "direct": 1, "pair.scale": 2} engine._metas = [ - M2NParamMeta(name, torch.float32, (1,), REPLICATED) + M2NParamMeta( + name, + torch.float32, + (1,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + ) for name in engine._index ] engine._parameter_destinations = [ @@ -326,9 +837,6 @@ class TestWireTypes: engine._reshard = reshard stream = Mock() - monkeypatch.setattr( - "vllm.distributed.weight_transfer.m2n_engine.comm_ptr", lambda group: 0 - ) monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) engine.receive_weights( @@ -385,6 +893,284 @@ class TestWireTypes: assert info.is_sender +def test_metadata_group_receives_reserved_socket(monkeypatch): + reservation = Mock() + create = Mock(return_value=Mock()) + monkeypatch.setattr("vllm.distributed.utils.StatelessProcessGroup.create", create) + + stateless_init_metadata_group( + "127.0.0.1", + 1234, + rank=0, + world_size=2, + listen_socket=reservation, + ) + + create.assert_called_once_with( + host="127.0.0.1", + port=1234, + rank=0, + world_size=2, + listen_socket=reservation, + ) + + +class _StaticM2NSource(M2NWeightSource): + def __init__(self, metadata, values): + self._metadata = metadata + self._values = values + + def metadata(self): + return list(self._metadata) + + def __iter__(self): + return iter(self._values) + + +class TestTrainerSourceContract: + @staticmethod + def _meta(name="weight"): + return M2NParamMeta( + name, + torch.float32, + (4,), + M2NLayout(M2NMesh((1, 1), 0), REPLICATED), + ) + + @staticmethod + def _engine(source): + engine = M2NTrainerWeightTransferEngine( + client=Mock(), source=source, is_sender=False + ) + engine._m2n = _fake_m2n() + engine._handle = object() + engine._metas = source.metadata() + engine._dst_mesh = M2NMesh((1, 1), 1) + engine._dst_placements = [REPLICATED] * len(engine._metas) + engine.group = Mock(comm=1) + return engine + + def test_partial_trainer_initialization_releases_transport_resources( + self, monkeypatch + ): + metadata_group = Mock(rank=0, world_size=2) + metadata_group.all_gather_obj.side_effect = lambda envelope: [envelope] * 2 + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer." + "stateless_init_metadata_group", + Mock(return_value=metadata_group), + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer.import_m2n", Mock() + ) + handle = Mock() + nccl = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer.prepare_m2n_local_runtime", + Mock(return_value=(nccl, 0, handle)), + ) + communicator = Mock() + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer.create_m2n_communicator", + Mock(return_value=communicator), + ) + monkeypatch.setattr( + "vllm.distributed.weight_transfer.m2n_trainer." + "publish_destination_placements", + Mock(return_value=[]), + ) + source = _StaticM2NSource([], []) + client = Mock() + client.init_weight_transfer_engine.side_effect = RuntimeError( + "worker init failed" + ) + engine = M2NTrainerWeightTransferEngine( + client=client, source=source, is_sender=True + ) + init_info = M2NTrainerInitInfo( + rank=0, + master_address="127.0.0.1", + master_port=1234, + world_size=2, + ) + + with pytest.raises(RuntimeError, match="worker init failed"): + engine._initialize(init_info, source) + + communicator.abort.assert_called_once_with() + communicator.destroy.assert_not_called() + handle.destroy.assert_called_once_with() + assert engine.group is None + assert engine._handle is None + + def test_non_sender_observes_rank_zero_finish_failure(self): + outcomes: dict[int, dict[str, object]] = {} + sender_group = Mock() + sender_group.send_obj.side_effect = lambda outcome, rank: outcomes.__setitem__( + rank, outcome + ) + sender_group.recv_obj.return_value = {"phase": "round_outcome_consumed"} + + source = _StaticM2NSource([], []) + sender = self._engine(source) + sender.is_sender = True + sender._metadata_group = sender_group + sender._num_trainer_ranks = 2 + sender._send = Mock() + sender.client = Mock() + sender.client.finish_weight_update.side_effect = RuntimeError("finish failed") + + with pytest.raises(RuntimeError, match="finish failed"): + sender.send_weights() + + sender._send.assert_called_once_with() + sender_group.recv_obj.assert_called_once_with(1) + receiver_group = Mock() + receiver_group.recv_obj.side_effect = lambda _source_rank: outcomes.pop(1) + receiver = self._engine(source) + receiver._metadata_group = receiver_group + receiver._num_trainer_ranks = 2 + receiver._send = Mock() + + with pytest.raises( + RuntimeError, + match=( + "worker update failed on trainer rank zero: RuntimeError: finish failed" + ), + ): + receiver.send_weights() + + receiver._send.assert_called_once_with() + receiver_group.recv_obj.assert_called_once_with(0) + receiver_group.send_obj.assert_called_once_with( + {"phase": "round_outcome_consumed"}, 0 + ) + assert not outcomes + + def test_sender_rejects_invalid_round_outcome_acknowledgement(self): + source = _StaticM2NSource([], []) + sender = self._engine(source) + sender.is_sender = True + sender._metadata_group = Mock() + sender._metadata_group.recv_obj.return_value = {"phase": "wrong"} + sender._num_trainer_ranks = 2 + sender._send = Mock() + sender.client = Mock() + + with pytest.raises(RuntimeError, match="invalid round-outcome acknowledgement"): + sender.send_weights() + + def test_non_sender_consumes_outcome_after_local_failure(self): + source = _StaticM2NSource([], []) + receiver = self._engine(source) + receiver._send = Mock(side_effect=[ValueError("local send failed"), None]) + receiver._metadata_group = Mock() + receiver._metadata_group.recv_obj.side_effect = [ + {"ok": False, "error": "RuntimeError: first round failed"}, + {"ok": True, "error": None}, + ] + + with pytest.raises(ValueError, match="local send failed"): + receiver.send_weights() + receiver.send_weights() + + assert receiver._send.call_count == 2 + assert receiver._metadata_group.recv_obj.call_count == 2 + assert receiver._metadata_group.send_obj.call_args_list == [ + call({"phase": "round_outcome_consumed"}, 0), + call({"phase": "round_outcome_consumed"}, 0), + ] + + @pytest.mark.parametrize( + ("case", "match"), + [ + ("missing", "first missing: 'weight'"), + ("extra", "first extra"), + ("reordered", "yielded 'other'.*declared 'weight'"), + ], + ) + def test_source_cardinality_and_order(self, monkeypatch, case, match): + buffer = _BufferLike() + values = { + "missing": [], + "extra": [("weight", buffer), ("extra", buffer)], + "reordered": [("other", buffer)], + }[case] + source = _StaticM2NSource([self._meta()], values) + engine = self._engine(source) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + with pytest.raises(RuntimeError, match=match): + engine._send() + + @pytest.mark.parametrize( + ("case", "match"), [("dtype", "dtype"), ("shape", "local shape")] + ) + def test_materialized_buffer_matches_declared_metadata( + self, monkeypatch, case, match + ): + buffer = _BufferLike() + if case == "dtype": + buffer.dtype = torch.float16 + else: + buffer.shape = (2,) + + source = _StaticM2NSource([self._meta()], [("weight", buffer)]) + engine = self._engine(source) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + with pytest.raises(ValueError, match=match): + engine._send() + + engine._m2n.reshard.assert_not_called() + + def test_m2n_compatible_buffer_needs_no_torch_methods(self, monkeypatch): + buffer = _BufferLike() + source = _StaticM2NSource([self._meta()], [("weight", buffer)]) + engine = self._engine(source) + stream = Mock() + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + engine._send() + + assert engine._m2n.reshard.call_args.args[0] is buffer + stream.synchronize.assert_called_once_with() + + def test_allocator_aware_buffer_records_send_stream(self, monkeypatch): + buffer = _AllocatorAwareBuffer() + source = _StaticM2NSource([self._meta()], [("weight", buffer)]) + engine = self._engine(source) + stream = Mock() + monkeypatch.setattr(torch.cuda, "current_stream", lambda: stream) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + engine._send() + + buffer.record_stream.assert_called_once_with(stream) + + def test_raw_pointer_uses_declared_shape_and_dtype(self, monkeypatch): + pointer = 1234 + source = _StaticM2NSource([self._meta()], [("weight", pointer)]) + engine = self._engine(source) + monkeypatch.setattr(torch.cuda, "current_stream", Mock()) + monkeypatch.setattr(torch.accelerator, "synchronize", Mock()) + monkeypatch.setattr(torch.accelerator, "current_device_index", lambda: 0) + + engine._send() + + call = engine._m2n.reshard.call_args + assert call.args[0] == pointer + assert call.kwargs["src_local_shape"] == (4,) + assert call.kwargs["src_dtype"] == torch.float32 + + class TestRegistration: def test_both_registries_expose_the_backend(self): """Both worker and trainer factories register nccl_m2n.""" @@ -433,8 +1219,15 @@ def _assigned_device() -> "torch.device": @ray.remote(num_gpus=1) -def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: - """Send one parameter through the real trainer engine.""" +def _m2n_trainer_send( + master_address: str, + master_port: int, + world_size: int, + nccl_unique_id_b64: str | None = None, + trainer_rank: int = 0, + mixed_layout: bool = False, +) -> bool: + """Send either the basic or mixed-layout manifest through the real engine.""" device = _assigned_device() from vllm.distributed.weight_transfer import WeightTransferTrainerFactory @@ -463,16 +1256,47 @@ def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: ).reshape(SHAPE) ) + num_trainer_ranks = 2 if mixed_layout else 1 + if mixed_layout: + row_full = torch.arange(16, dtype=torch.float32, device=device).reshape(4, 4) + column_full = row_full + 100 + row_local = row_full.chunk(2, dim=0)[trainer_rank].contiguous() + column_local = column_full.chunk(2, dim=1)[trainer_rank].contiguous() + source = ManifestM2NWeightSource( + [ + M2NManifestEntry( + "row_sharded", + torch.float32, + (4, 4), + M2NLayout(M2NMesh((1, 2), 0), (REPLICATE, 0)), + lambda: row_local, + ), + M2NManifestEntry( + "column_sharded", + torch.float32, + (4, 4), + M2NLayout(M2NMesh((2, 1), 0), (1, REPLICATE)), + lambda: column_local, + ), + ] + ) + dst_mesh_dims = (2, 1) + else: + source = DTensorModuleSource(Tiny(), num_trainer_ranks=1) + dst_mesh_dims = (1, world_size - 1) + engine = WeightTransferTrainerFactory.trainer_init( init_info=M2NTrainerInitInfo( + master_address=master_address, + master_port=master_port, nccl_unique_id_b64=nccl_unique_id_b64, world_size=world_size, - num_trainer_ranks=1, - dst_mesh_dims=(1, world_size - 1), - rank=0, + num_trainer_ranks=num_trainer_ranks, + dst_mesh_dims=dst_mesh_dims, + rank=trainer_rank, ), client=NoopClient(), - source=DTensorModuleSource(Tiny(), num_trainer_ranks=1), + source=source, ) engine.send_weights() torch.accelerator.synchronize() @@ -482,12 +1306,14 @@ def _m2n_trainer_send(nccl_unique_id_b64: str, world_size: int) -> bool: @ray.remote(num_gpus=1) def _m2n_worker_receive( - nccl_unique_id_b64: str, + master_address: str, + master_port: int, world_size: int, worker_rank: int = 0, + nccl_unique_id_b64: str | None = None, + mixed_layout: bool = False, ) -> dict: - """Receive that parameter through the real worker engine.""" - import contextlib + """Receive either the basic or mixed-layout manifest through the engine.""" from unittest.mock import MagicMock device = _assigned_device() @@ -503,8 +1329,9 @@ def _m2n_worker_receive( class Recorder(torch.nn.Module): def __init__(self): super().__init__() - self.received = [] - if world_size > 2: + self.received: list[tuple[str, torch.Tensor]] = [] + self.received_by_name: dict[str, torch.Tensor] = {} + if not mixed_layout and world_size > 2: self.weight = torch.nn.Parameter( torch.zeros( SHAPE[0] // (world_size - 1), @@ -521,14 +1348,19 @@ def _m2n_worker_receive( def load_weights(self, weights): for name, tensor in weights: - self.received.append((name, tensor.clone())) + if mixed_layout: + self.received_by_name[name] = tensor.clone() + else: + self.received.append((name, tensor.clone())) + num_trainer_ranks = 2 if mixed_layout else 1 + num_workers = world_size - num_trainer_ranks parallel_config = MagicMock(spec=ParallelConfig) parallel_config.rank = worker_rank - parallel_config.world_size = world_size - 1 + parallel_config.world_size = num_workers parallel_config.data_parallel_rank = 0 parallel_config.data_parallel_index = 0 - parallel_config.tensor_parallel_size = world_size - 1 + parallel_config.tensor_parallel_size = num_workers parallel_config.pipeline_parallel_size = 1 vllm_config = MagicMock() vllm_config.parallel_config = parallel_config @@ -539,37 +1371,52 @@ def _m2n_worker_receive( engine = M2NWeightTransferEngine( WeightTransferConfig(backend="nccl_m2n"), vllm_config, device, recorder ) - # Transport-only: receive_weights enters set_current_vllm_config, and - # vllm_config here is a mock. - import vllm.config as _vllm_config_mod - - _vllm_config_mod.set_current_vllm_config = lambda cfg: contextlib.nullcontext() - + if mixed_layout: + params = [ + M2NWireParam("row_sharded", "float32", (4, 4), (1, 2), (REPLICATE, 0)), + M2NWireParam("column_sharded", "float32", (4, 4), (2, 1), (1, REPLICATE)), + ] + names = [param.name for param in params] + dst_mesh_dims = [2, 1] + else: + params = [M2NWireParam("weight", DTYPE, tuple(SHAPE), (1, 1), REPLICATED)] + names = ["weight"] + dst_mesh_dims = [1, num_workers] engine.init_transfer_engine( M2NWeightTransferInitInfo( + schema_version=M2N_WIRE_SCHEMA_VERSION, + master_address=master_address, + master_port=master_port, nccl_unique_id_b64=nccl_unique_id_b64, - rank_offset=1, # trainer occupies rank 0 + rank_offset=num_trainer_ranks, world_size=world_size, - src_mesh_dims=[1, 1], - dst_mesh_dims=[1, world_size - 1], - names=["weight"], - dtype_names=[DTYPE], - shapes=[SHAPE], - src_placements=[None], # replicated on the single trainer rank + dst_mesh_dims=dst_mesh_dims, + source_digest=source_plan_digest(params), + params=[param.to_dict() for param in params], ) ) - engine.receive_weights(M2NWeightTransferUpdateInfo(names=["weight"])) + engine.start_weight_update() + engine.receive_weights(M2NWeightTransferUpdateInfo(names=names)) + engine.finish_weight_update() torch.accelerator.synchronize() + if mixed_layout: + result = { + name: tensor.cpu().tolist() + for name, tensor in recorder.received_by_name.items() + } + engine.shutdown() + return result + full = torch.arange( SHAPE[0] * SHAPE[1], dtype=torch.float32, device=device ).reshape(SHAPE) - direct = world_size > 2 + direct = num_workers > 1 name: str | None got: torch.Tensor | None if direct: name, got = "weight", recorder.weight - expected = full.chunk(world_size - 1, dim=0)[worker_rank] + expected = full.chunk(num_workers, dim=0)[worker_rank] else: name, got = recorder.received[0] if recorder.received else (None, None) expected = full @@ -588,7 +1435,8 @@ def _m2n_worker_receive( torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs: one trainer rank and one inference worker.", ) -def test_m2n_weight_transfer_between_processes(): +@pytest.mark.parametrize("data_plane", ["tcp", "uid"]) +def test_m2n_weight_transfer_between_processes(data_plane): """A parameter survives a real reshard from a trainer process to a worker. This is the only test here that moves bytes: it builds both engines, joins @@ -598,17 +1446,30 @@ def test_m2n_weight_transfer_between_processes(): pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") _init_ray() - from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary - - nccl = NCCLLibrary() - nccl_unique_id_b64 = base64.b64encode( - bytes(nccl.ncclGetUniqueId().internal) - ).decode() + master_address = "127.0.0.1" + master_port = get_open_port() world_size = 2 # 1 trainer + 1 inference worker + nccl_unique_id_b64 = None + if data_plane == "uid": + import nccl.core as nccl - worker = _m2n_worker_receive.remote(nccl_unique_id_b64, world_size) - trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) - trainer_ok, result = ray.get([trainer, worker]) + nccl_unique_id_b64 = base64.b64encode( + bytes(nccl.get_unique_id().as_bytes) + ).decode() + + worker = _m2n_worker_receive.remote( + master_address, + master_port, + world_size, + nccl_unique_id_b64=nccl_unique_id_b64, + ) + trainer = _m2n_trainer_send.remote( + master_address, + master_port, + world_size, + nccl_unique_id_b64, + ) + trainer_ok, result = ray.get([trainer, worker], timeout=300) assert trainer_ok, "trainer engine did not complete" assert result["count"] == 1, f"expected one parameter, got {result['count']}" @@ -626,20 +1487,16 @@ def test_m2n_weight_transfer_to_tp2_shards(): pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") _init_ray() - from vllm.distributed.device_communicators.pynccl_wrapper import NCCLLibrary - - nccl = NCCLLibrary() - nccl_unique_id_b64 = base64.b64encode( - bytes(nccl.ncclGetUniqueId().internal) - ).decode() + master_address = "127.0.0.1" + master_port = get_open_port() world_size = 3 workers = [ - _m2n_worker_receive.remote(nccl_unique_id_b64, world_size, rank) + _m2n_worker_receive.remote(master_address, master_port, world_size, rank) for rank in range(2) ] - trainer = _m2n_trainer_send.remote(nccl_unique_id_b64, world_size) - trainer_ok, *results = ray.get([trainer, *workers]) + trainer = _m2n_trainer_send.remote(master_address, master_port, world_size) + trainer_ok, *results = ray.get([trainer, *workers], timeout=300) assert trainer_ok, "trainer engine did not complete" for rank, result in enumerate(results): @@ -650,6 +1507,52 @@ def test_m2n_weight_transfer_to_tp2_shards(): assert result["exact"], f"worker {rank} received the wrong shard" +@pytest.mark.skipif( + torch.accelerator.device_count() < 4, + reason="Need 4 GPUs: two trainer ranks and two inference workers.", +) +def test_m2n_transfer_with_per_parameter_source_meshes(): + """One collective sequence accepts heterogeneous source factorizations.""" + pytest.importorskip("nccl.m2n", reason="nccl_m2n backend needs the m2n runtime") + _init_ray() + + master_address = "127.0.0.1" + master_port = get_open_port() + trainers = [ + _m2n_trainer_send.remote( + master_address, + master_port, + 4, + trainer_rank=rank, + mixed_layout=True, + ) + for rank in range(2) + ] + workers = [ + _m2n_worker_receive.remote( + master_address, + master_port, + 4, + worker_rank=rank, + mixed_layout=True, + ) + for rank in range(2) + ] + results = ray.get(trainers + workers, timeout=300) + trainer_results, worker_results = results[:2], results[2:] + + assert trainer_results == [True, True] + expected_row = torch.arange(16, dtype=torch.float32).reshape(4, 4).tolist() + expected_column = ( + torch.arange(16, dtype=torch.float32).reshape(4, 4) + 100 + ).tolist() + for result in worker_results: + assert result == { + "row_sharded": expected_row, + "column_sharded": expected_column, + } + + class _Model(torch.nn.Module): """Stand-in for a TP=2 shard of a tiny model. diff --git a/vllm/distributed/weight_transfer/m2n_common.py b/vllm/distributed/weight_transfer/m2n_common.py index a2c6b73844..13f7521408 100644 --- a/vllm/distributed/weight_transfer/m2n_common.py +++ b/vllm/distributed/weight_transfer/m2n_common.py @@ -9,16 +9,19 @@ runtime import, layout descriptors, and the conversion to `nccl.m2n` types — so the engine module stays about the transfer itself. """ -from collections.abc import Sequence +import hashlib +import json +from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, Any, cast import torch from vllm.distributed.weight_transfer.base import ParamMeta +from vllm.utils.nccl import unpinned_nccl_env if TYPE_CHECKING: - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup # Placement code for a mesh axis that replicates. Any other (non-negative) code # is the tensor dim that axis shards. Ints rather than `nccl.m2n` objects so @@ -35,15 +38,15 @@ MESH_NDIMS = 2 DESTINATION_REPLICA_AXIS = 0 DESTINATION_SHARD_AXIS = 1 -# m2n bounds how many source shards may feed one destination shard, and how many -# destination shards one source shard may feed, with compile-time arrays in -# `reshard_limits.h`. The bindings do not expose them, so they are mirrored -# here; a build with larger arrays makes these conservative, never wrong. +# Conservative capacities accepted by the supported nccl_m2n algorithms. They +# are not exposed by its Python API yet; remove these mirrors once nccl_m2n +# provides communication-free plan validation. MAX_SOURCE_SHARDS = 16 MAX_DEST_SHARDS = 64 -# Wire dtypes `ncclReshard` accepts. Notably excludes fp4 and any packed -# sub-byte type, so quantized checkpoints are out of scope for now. +# Increment when the JSON init payload changes incompatibly. +M2N_WIRE_SCHEMA_VERSION = 2 + SUPPORTED_DTYPES: frozenset[torch.dtype] = frozenset( dtype for dtype in ( @@ -70,12 +73,62 @@ _IMPORT_HINT = ( "The nccl_m2n weight transfer backend requires the `nccl-extensions` " "package (and its `nccl4py` dependency), which is not installed by vLLM. " "See https://github.com/NVIDIA/nccl-extensions for build " - "and install instructions. It needs NCCL 2.30.5 or newer, and " - "VLLM_NCCL_SO_PATH must point at the same libnccl.so that libnccl_m2n.so " - "was linked against." + "and install instructions. It needs NCCL 2.30.5 or newer." ) +def prepare_m2n_local_runtime(m2n: Any, max_cta: int | None) -> tuple[Any, int, Any]: + """Prepare rank-local resources before any rank enters NCCL init.""" + import nccl.core as nccl + + device = torch.accelerator.current_device_index() + if not isinstance(device, int) or isinstance(device, bool) or device < 0: + raise RuntimeError(f"invalid current accelerator device index: {device!r}") + handle = m2n.Handle.create(m2n.Config(max_cta=max_cta)) + return nccl, device, handle + + +def create_m2n_communicator( + nccl: Any, + group: "StatelessProcessGroup", + unique_id_bytes: bytes | None, +) -> Any: + """Create the M2N data-plane communicator through NCCL4Py.""" + if unique_id_bytes is None: + envelope: dict[str, bytes | str | None] | None = None + if group.rank == 0: + try: + envelope = { + "unique_id": bytes(nccl.get_unique_id().as_bytes), + "error": None, + } + except Exception as exc: + envelope = { + "unique_id": None, + "error": f"{type(exc).__name__}: {exc}", + } + envelope = group.broadcast_obj(envelope, src=0) + if not isinstance(envelope, Mapping) or set(envelope) != {"unique_id", "error"}: + raise RuntimeError("invalid NCCL4Py unique-id envelope") + error = envelope["error"] + received_unique_id = envelope["unique_id"] + if error is not None: + if not isinstance(error, str) or not error: + raise RuntimeError("invalid NCCL4Py unique-id error") + raise RuntimeError(f"NCCL4Py unique-id generation failed: {error}") + if not isinstance(received_unique_id, bytes): + raise RuntimeError("NCCL4Py unique-id broadcast did not return bytes") + unique_id_bytes = received_unique_id + + unique_id = nccl.UniqueId.from_bytes(unique_id_bytes) + with unpinned_nccl_env(): + return nccl.Communicator.init( + nranks=group.world_size, + rank=group.rank, + unique_id=unique_id, + ) + + def import_m2n() -> Any: """Import `nccl.m2n` lazily, with an actionable error when it is missing. @@ -97,17 +150,25 @@ class M2NMesh: `[start_rank, start_rank + dims[0] * dims[1])`. There is no 1-D mesh; a single-axis topology is spelled with a second axis of size 1. - One mesh describes every tensor on its side, so it is exchanged once at the - init handshake rather than per parameter. + A model may use a different factorization for each tensor. For example, + dense weights can use a DP x TP mesh while expert weights use EDP x EP. """ dims: tuple[int, int] start_rank: int def __post_init__(self) -> None: - if len(self.dims) != MESH_NDIMS or any(d <= 0 for d in self.dims): + object.__setattr__(self, "dims", tuple(self.dims)) + if len(self.dims) != MESH_NDIMS or any( + not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 + for dim in self.dims + ): raise ValueError(f"mesh dims must be {MESH_NDIMS} positive ints") - if self.start_rank < 0: + if ( + not isinstance(self.start_rank, int) + or isinstance(self.start_rank, bool) + or self.start_rank < 0 + ): raise ValueError( f"mesh start_rank must be non-negative, got {self.start_rank}" ) @@ -124,16 +185,26 @@ Placements = tuple[int, int] REPLICATED: Placements | None = None -def check_placements(placements: Placements, context: str = "placements") -> None: - """Reject placement pairs m2n cannot express. +@dataclass(frozen=True) +class M2NLayout: + """A tensor's mesh and placement on one side of a transfer.""" - The header requires exactly one SHARD axis and one REPLICATE axis. - `{REPLICATE, REPLICATE}` hits a degenerate prepare branch, which is why - full replication is carried as `REPLICATED` and resolved separately; - sharding both axes is not expressible at all. - """ + mesh: M2NMesh + placements: Placements | None + + def __post_init__(self) -> None: + if not isinstance(self.mesh, M2NMesh): + raise TypeError("M2N layout mesh must be an M2NMesh") + if self.placements is not None: + object.__setattr__(self, "placements", tuple(self.placements)) + + +def check_placements(placements: Placements, context: str = "placements") -> None: + """Reject placement pairs nccl_m2n cannot express.""" if len(placements) != MESH_NDIMS: raise ValueError(f"{context} must have {MESH_NDIMS} entries") + if any(not isinstance(code, int) or isinstance(code, bool) for code in placements): + raise TypeError(f"{context} must contain integer placement codes") invalid = [code for code in placements if code < REPLICATE] if invalid: raise ValueError( @@ -144,8 +215,7 @@ def check_placements(placements: Placements, context: str = "placements") -> Non if num_sharded == 0: raise ValueError( f"{context}: a fully replicated tensor is carried as REPLICATED, " - "not as two " - "REPLICATE axes" + "not as two REPLICATE axes" ) if num_sharded == MESH_NDIMS: raise ValueError( @@ -175,87 +245,46 @@ def resolve_layout( return mesh, placements -def shard_count(mesh: M2NMesh, placements: Placements) -> int: - """How many pieces this layout splits the tensor into.""" - return next( - (mesh.dims[axis] for axis, code in enumerate(placements) if code != REPLICATE), - 1, +def validate_layout( + mesh: M2NMesh, + placements: Placements, + shape: Sequence[int], + context: str, +) -> tuple[int, ...]: + """Validate one resolved layout and return its required local shape.""" + local_shape = list(shape) + for axis, tensor_dim in enumerate(placements): + if tensor_dim == REPLICATE: + continue + if tensor_dim >= len(shape): + raise ValueError( + f"{context} shards tensor dim {tensor_dim}, but parameter has " + f"rank {len(shape)}" + ) + factor = mesh.dims[axis] + if shape[tensor_dim] % factor: + raise ValueError( + f"{context} shards dim {tensor_dim} (size {shape[tensor_dim]}) " + f"over {factor} ranks, which does not divide evenly" + ) + local_shape[tensor_dim] //= factor + return tuple(local_shape) + + +def expected_local_shape(meta: "M2NParamMeta") -> tuple[int, ...]: + """Validate a parameter's source layout and return its local shape.""" + mesh, placements = resolve_layout( + meta.source_layout.mesh, + meta.source_layout.placements, + f"parameter '{meta.name}' source placements", + ) + return validate_layout( + mesh, placements, meta.shape, f"parameter '{meta.name}' source layout" ) -def check_plan_limits( - src: tuple[M2NMesh, Placements], - dst: tuple[M2NMesh, Placements], - name: str, -) -> None: - """Reject plans that exceed m2n's static per-shard fan-in / fan-out arrays. - - Only the unambiguous cases are checked: when one side is a single shard it - is fed by (or feeds) every shard on the other side, so the count is exactly - the other side's shard count. In the general sharded-to-sharded case the - overlap depends on the library's chunking, and reproducing that arithmetic - here would duplicate internals that can change under us. - - m2n re-checks authoritatively and, because the plan is derived from the - shared descriptors, fails identically on every rank -- so this is about - reporting at init with a message that names the parameter, not about - avoiding a hang. - """ - src_shards = shard_count(*src) - dst_shards = shard_count(*dst) - if dst_shards == 1 and src_shards > MAX_SOURCE_SHARDS: - raise ValueError( - f"parameter '{name}' would feed {src_shards} source shards into one " - f"replicated destination, over m2n's MAX_SOURCES={MAX_SOURCE_SHARDS}. " - "Shard the destination, replicate the source, or rebuild m2n with " - "larger arrays." - ) - if src_shards == 1 and dst_shards > MAX_DEST_SHARDS: - raise ValueError( - f"parameter '{name}' would fan one source shard out to {dst_shards} " - f"destination shards, over m2n's MAX_TARGETS={MAX_DEST_SHARDS}. " - "Rebuild m2n with larger arrays to raise the bound." - ) - - -def validate_layout( - mesh: M2NMesh, placements: Placements, shape: Sequence[int], side: str -) -> None: - """Check a resolved layout can describe `shape` before any collective runs. - - Called on both sides at init so a bad plan surfaces as an error from the - init RPC rather than as a hang inside the first reshard. - """ - for axis, code in enumerate(placements): - if code == REPLICATE: - continue - if code >= len(shape): - raise ValueError( - f"{side} layout shards tensor dim {code}, but the tensor " - f"has rank {len(shape)}" - ) - factor = mesh.dims[axis] - if shape[code] % factor: - raise ValueError( - f"{side} layout shards dim {code} (size {shape[code]}) over " - f"{factor} ranks, which does not divide evenly" - ) - - -@dataclass(frozen=True) -class M2NParamMeta(ParamMeta): - """`ParamMeta` extended with how the trainer places this tensor. - - The base class carries only name / dtype / full shape, which is not enough - to plan a reshard. `placements` is relative to its side's `M2NMesh`, or - `REPLICATED` when every rank holds the whole tensor. - """ - - placements: Placements | None - - def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> None: - """Reject tensors m2n cannot move at all, with the parameter named.""" + """Reject parameters nccl_m2n cannot transfer.""" if dtype not in SUPPORTED_DTYPES: raise ValueError( f"parameter '{name}' has dtype {dtype}, which nccl_m2n does not " @@ -268,6 +297,514 @@ def check_transferable(name: str, dtype: torch.dtype, shape: Sequence[int]) -> N ) +def _shard_spec(mesh: M2NMesh, placements: Placements) -> tuple[int, int]: + return next( + (code, mesh.dims[axis]) + for axis, code in enumerate(placements) + if code != REPLICATE + ) + + +def _max_partition_overlaps( + extent: int, + partitions: int, + other_partitions: int, +) -> int: + """Return the maximum other partitions intersected by any one partition.""" + if extent == 0: + return 0 + width = extent // partitions + other_width = extent // other_partitions + max_overlaps = 0 + for index in range(partitions): + start = index * width + end = start + width + first_other = start // other_width + last_other = (end - 1) // other_width + max_overlaps = max(max_overlaps, last_other - first_other + 1) + return max_overlaps + + +def check_plan_limits( + src: tuple[M2NMesh, Placements], + dst: tuple[M2NMesh, Placements], + name: str, + shape: Sequence[int], +) -> None: + """Reject plans whose per-shard fan-in or fan-out exceeds M2N capacity.""" + src_dim, src_shards = _shard_spec(*src) + dst_dim, dst_shards = _shard_spec(*dst) + if any(extent == 0 for extent in shape): + max_sources = max_targets = 0 + elif src_dim != dst_dim: + max_sources = src_shards + max_targets = dst_shards + else: + extent = shape[src_dim] + max_sources = _max_partition_overlaps(extent, dst_shards, src_shards) + max_targets = _max_partition_overlaps(extent, src_shards, dst_shards) + + if max_sources > MAX_SOURCE_SHARDS: + raise ValueError( + f"parameter '{name}' would feed up to {max_sources} source shards " + f"into one destination shard, over nccl_m2n MAX_SOURCES=" + f"{MAX_SOURCE_SHARDS}" + ) + if max_targets > MAX_DEST_SHARDS: + raise ValueError( + f"parameter '{name}' would fan one source shard out to up to " + f"{max_targets} destination shards, over nccl_m2n " + f"MAX_TARGETS={MAX_DEST_SHARDS}" + ) + + +def validate_transfer_plan( + metas: Sequence["M2NParamMeta"], + dst_mesh: M2NMesh, + dst_placements: Sequence[Placements | None], +) -> None: + """Validate the complete ordered source-to-destination transfer plan.""" + if len(dst_placements) != len(metas): + raise ValueError( + f"destination plan has {len(dst_placements)} parameters, expected " + f"{len(metas)}" + ) + for meta, placement in zip(metas, dst_placements): + check_transferable(meta.name, meta.dtype, meta.shape) + src = resolve_layout( + meta.source_layout.mesh, + meta.source_layout.placements, + f"parameter '{meta.name}' source placements", + ) + validate_layout(*src, meta.shape, f"parameter '{meta.name}' source layout") + dst = resolve_layout( + dst_mesh, + placement, + f"parameter '{meta.name}' destination placements", + ) + validate_layout(*dst, meta.shape, f"parameter '{meta.name}' destination layout") + check_plan_limits(src, dst, meta.name, meta.shape) + + +@dataclass(frozen=True) +class M2NParamMeta(ParamMeta): + """`ParamMeta` extended with how the trainer places this tensor. + + The base class carries only name / dtype / full shape, which is not enough + to plan a reshard. `source_layout` carries both the per-parameter mesh and + its placement on that mesh. + """ + + source_layout: M2NLayout + + def __post_init__(self) -> None: + if not isinstance(self.source_layout, M2NLayout): + raise TypeError("source_layout must be an M2NLayout") + + +@dataclass(frozen=True) +class M2NWireParam: + """JSON-safe source metadata for one parameter in the init handshake.""" + + name: str + dtype_name: str + shape: tuple[int, ...] + src_mesh_dims: tuple[int, int] + src_placements: Placements | None + + def __post_init__(self) -> None: + shape = tuple(self.shape) + dims = tuple(self.src_mesh_dims) + placements = None if self.src_placements is None else tuple(self.src_placements) + object.__setattr__(self, "shape", shape) + object.__setattr__(self, "src_mesh_dims", dims) + object.__setattr__(self, "src_placements", placements) + if not isinstance(self.name, str) or not self.name: + raise ValueError("M2N wire parameter name must not be empty") + if not isinstance(self.dtype_name, str) or not self.dtype_name: + raise ValueError(f"parameter '{self.name}' has an empty wire dtype name") + if not shape or any( + not isinstance(dim, int) or isinstance(dim, bool) or dim <= 0 + for dim in shape + ): + raise ValueError( + f"parameter '{self.name}' shape must contain positive integers" + ) + M2NMesh(cast(tuple[int, int], dims), 0) + if placements is not None: + check_placements( + cast(Placements, placements), + f"parameter '{self.name}' source placements", + ) + + @classmethod + def from_dict(cls, value: Mapping[str, Any]) -> "M2NWireParam": + """Parse one strict wire record from its JSON-decoded mapping.""" + fields = { + "name", + "dtype_name", + "shape", + "src_mesh_dims", + "src_placements", + } + missing = fields - value.keys() + extra = value.keys() - fields + if missing or extra: + details = [] + if missing: + details.append(f"missing {sorted(missing)}") + if extra: + details.append(f"unexpected {sorted(extra)}") + raise ValueError(f"invalid M2N wire parameter: {', '.join(details)}") + return cls(**value) + + def to_dict(self) -> dict[str, Any]: + """Serialize this record using only JSON-compatible values.""" + return { + "name": self.name, + "dtype_name": self.dtype_name, + "shape": list(self.shape), + "src_mesh_dims": list(self.src_mesh_dims), + "src_placements": ( + None if self.src_placements is None else list(self.src_placements) + ), + } + + +def source_plan_digest(params: Sequence[M2NWireParam]) -> str: + """Return a rank-independent digest of the ordered source wire plan.""" + payload = json.dumps( + { + "schema_version": M2N_WIRE_SCHEMA_VERSION, + "params": [param.to_dict() for param in params], + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def complete_plan_digest( + source_digest: str, + dst_mesh: M2NMesh, + dst_placements: Sequence[Placements | None], +) -> str: + """Return a digest of the agreed source and destination plan.""" + if not isinstance(source_digest, str) or not source_digest: + raise ValueError("source digest must be a non-empty string") + payload = json.dumps( + { + "schema_version": M2N_WIRE_SCHEMA_VERSION, + "source_digest": source_digest, + "destination_mesh": { + "dims": list(dst_mesh.dims), + "start_rank": dst_mesh.start_rank, + }, + "destination_placements": [ + None if placement is None else list(placement) + for placement in dst_placements + ], + }, + sort_keys=True, + separators=(",", ":"), + ).encode() + return hashlib.sha256(payload).hexdigest() + + +def check_source_plan_agreement( + group: "StatelessProcessGroup", + params: Sequence[M2NWireParam], + expected_digest: str | None = None, + local_error: str | None = None, +) -> str: + """Verify every participant entered with the same source plan.""" + prepared_error = local_error + digest: str | None = None + declared: str | None = expected_digest + if prepared_error is None: + try: + digest = source_plan_digest(params) + declared = digest if expected_digest is None else expected_digest + if not isinstance(declared, str) or not declared: + raise ValueError("declared source digest must be a non-empty string") + except Exception as exc: + prepared_error = f"{type(exc).__name__}: {exc}" + elif not isinstance(prepared_error, str) or not prepared_error: + prepared_error = "TypeError: local source error must be a non-empty string" + + gathered = group.all_gather_obj( + { + "phase": "source", + "digest": digest, + "declared_digest": declared, + "error": prepared_error, + } + ) + errors: list[str] = [] + expected_fields = {"phase", "digest", "declared_digest", "error"} + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + agreed: list[tuple[int, str, str]] = [] + for rank, item in enumerate(gathered): + if not isinstance(item, Mapping) or set(item) != expected_fields: + errors.append(f"rank {rank}: invalid source preflight envelope") + continue + if item["phase"] != "source": + errors.append(f"rank {rank}: invalid source preflight phase") + continue + error = item["error"] + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid source preflight error") + else: + errors.append(f"rank {rank}: {error}") + continue + computed = item["digest"] + claimed = item["declared_digest"] + if ( + not isinstance(computed, str) + or not computed + or not isinstance(claimed, str) + or not claimed + ): + errors.append(f"rank {rank}: invalid source digest") + continue + agreed.append((rank, computed, claimed)) + if errors: + raise RuntimeError( + "nccl_m2n source preflight rejected before NCCL init: " + "; ".join(errors) + ) + + canonical = agreed[0][1] + for rank, computed, claimed in agreed: + if computed != claimed: + errors.append( + f"rank {rank}: computed digest {computed} does not match " + f"declared digest {claimed}" + ) + if computed != canonical: + errors.append(f"rank {rank}: source digest disagrees with rank 0") + if errors: + raise RuntimeError( + "nccl_m2n source preflight rejected before NCCL init: " + "; ".join(errors) + ) + return canonical + + +def check_complete_plan_agreement( + group: "StatelessProcessGroup", + source_digest: str, + dst_mesh: M2NMesh, + dst_placements: Sequence[Placements | None], + local_error: str | None = None, +) -> str: + """Require every rank to accept the same complete transfer plan.""" + prepared_error = local_error + digest: str | None = None + if prepared_error is None: + try: + digest = complete_plan_digest(source_digest, dst_mesh, dst_placements) + except Exception as exc: + prepared_error = f"{type(exc).__name__}: {exc}" + elif not isinstance(prepared_error, str) or not prepared_error: + prepared_error = "TypeError: local plan error must be a non-empty string" + + gathered = group.all_gather_obj( + {"phase": "complete_plan", "digest": digest, "error": prepared_error} + ) + errors: list[str] = [] + valid: list[tuple[int, str]] = [] + expected_fields = {"phase", "digest", "error"} + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + for rank, item in enumerate(gathered): + if ( + not isinstance(item, Mapping) + or set(item) != expected_fields + or item.get("phase") != "complete_plan" + ): + errors.append(f"rank {rank}: invalid complete-plan envelope") + continue + error = item["error"] + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid complete-plan error") + else: + errors.append(f"rank {rank}: {error}") + continue + rank_digest = item["digest"] + if not isinstance(rank_digest, str) or not rank_digest: + errors.append(f"rank {rank}: invalid complete-plan digest") + continue + valid.append((rank, rank_digest)) + if not errors: + canonical = valid[0][1] + for rank, rank_digest in valid[1:]: + if rank_digest != canonical: + errors.append( + f"rank {rank}: complete-plan digest disagrees with rank 0" + ) + if errors: + raise RuntimeError( + "nccl_m2n complete-plan preflight rejected before first reshard: " + + "; ".join(errors) + ) + return valid[0][1] + + +def check_runtime_ready_agreement( + group: "StatelessProcessGroup", local_error: str | None +) -> None: + """Require every rank to prepare local resources before NCCL init.""" + prepared_error = local_error + if prepared_error is not None and ( + not isinstance(prepared_error, str) or not prepared_error + ): + prepared_error = "TypeError: local runtime error must be a non-empty string" + gathered = group.all_gather_obj({"phase": "runtime_ready", "error": prepared_error}) + errors: list[str] = [] + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + for rank, item in enumerate(gathered): + if ( + not isinstance(item, Mapping) + or set(item) != {"phase", "error"} + or item.get("phase") != "runtime_ready" + ): + errors.append(f"rank {rank}: invalid runtime readiness envelope") + continue + error = item.get("error") + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid runtime readiness error") + else: + errors.append(f"rank {rank}: {error}") + if errors: + raise RuntimeError( + "nccl_m2n local runtime preflight failed before NCCL init: " + + "; ".join(errors) + ) + + +def check_data_plane_agreement( + group: "StatelessProcessGroup", + unique_id_bytes: bytes | None, + local_error: str | None = None, +) -> None: + """Agree on the NCCL rendezvous identity before NCCL init.""" + prepared_error = local_error + uid_digest: str | None = None + if prepared_error is None: + try: + if unique_id_bytes is not None: + if not isinstance(unique_id_bytes, bytes): + raise TypeError("NCCL unique id must decode to bytes") + uid_digest = hashlib.sha256(unique_id_bytes).hexdigest() + except Exception as exc: + prepared_error = f"{type(exc).__name__}: {exc}" + elif not isinstance(prepared_error, str) or not prepared_error: + prepared_error = "TypeError: local data-plane error must be non-empty" + + gathered = group.all_gather_obj( + { + "phase": "data_plane", + "mode": "uid" if unique_id_bytes is not None else "tcp", + "uid_digest": uid_digest, + "error": prepared_error, + } + ) + expected_fields = { + "phase", + "mode", + "uid_digest", + "error", + } + errors: list[str] = [] + valid: list[tuple[int, str, str | None]] = [] + if len(gathered) != group.world_size: + errors.append( + "control-plane all-gather returned " + f"{len(gathered)} records for world size {group.world_size}" + ) + for rank, item in enumerate(gathered): + if ( + not isinstance(item, Mapping) + or set(item) != expected_fields + or item.get("phase") != "data_plane" + ): + errors.append(f"rank {rank}: invalid data-plane preflight envelope") + continue + error = item["error"] + if error is not None: + if not isinstance(error, str) or not error: + errors.append(f"rank {rank}: invalid data-plane preflight error") + else: + errors.append(f"rank {rank}: {error}") + continue + mode = item["mode"] + uid = item["uid_digest"] + if mode not in {"tcp", "uid"} or (mode == "uid") != isinstance(uid, str): + errors.append(f"rank {rank}: invalid data-plane identity") + continue + valid.append((rank, mode, uid)) + if not errors: + _, mode, uid = valid[0] + for rank, other_mode, other_uid in valid[1:]: + if (other_mode, other_uid) != (mode, uid): + errors.append(f"rank {rank}: NCCL data-plane identity disagrees") + if errors: + raise RuntimeError( + "nccl_m2n data-plane preflight rejected before NCCL init: " + + "; ".join(errors) + ) + + +def validate_local_buffer(m2n: Any, meta: M2NParamMeta, buffer: Any) -> None: + """Check source dtype/shape while delegating buffer validity to M2N.""" + # A raw pointer carries no inspectable metadata. The source owns its + # correctness; _send supplies the declared dtype and local shape to M2N. + if isinstance(buffer, int): + return + + mesh, placements = resolve_layout( + meta.source_layout.mesh, meta.source_layout.placements + ) + actual = m2n.DistTensor( + buffer, + mesh=to_mesh(m2n, mesh), + placements=to_placements(m2n, placements), + ) + declared = m2n.DistTensor( + None, + local_shape=actual.local_shape, + dtype=meta.dtype, + mesh=to_mesh(m2n, mesh), + placements=to_placements(m2n, placements), + ) + if actual.dtype != declared.dtype: + raise ValueError( + f"parameter '{meta.name}' source returned dtype {actual.dtype}, " + f"but its metadata declared {declared.dtype}" + ) + + expected = expected_local_shape(meta) + if actual.local_shape != expected: + raise ValueError( + f"parameter '{meta.name}' source returned local shape " + f"{actual.local_shape}, but layout {meta.source_layout} implies " + f"{expected} from global shape {meta.shape}" + ) + + def to_mesh(m2n: Any, mesh: M2NMesh) -> Any: return m2n.Mesh(mesh.dims, start_rank=mesh.start_rank) @@ -279,7 +816,9 @@ def to_placements(m2n: Any, placements: Placements) -> list[Any]: def publish_destination_placements( - comm: "PyNcclCommunicator", + comm: Any, + device: torch.device, + stream: torch.cuda.Stream, first_worker_rank: int, placements: "Sequence[Placements | None] | None", num_parameters: int, @@ -294,15 +833,15 @@ def publish_destination_placements( rather than deadlocking later. Only placements travel: the destination *mesh* is still derived from the - rank split, identically on both sides. Use the NCCL communicator directly - because unique-id rendezvous deliberately has no bootstrap process group. + rank split, identically on both sides. The plan is published on the data + plane before any M2N reshard collective. """ replicated_sentinel = REPLICATE - 1 encoded = torch.full( (num_parameters, MESH_NDIMS), replicated_sentinel, dtype=torch.int8, - device=comm.device, + device=device, ) if comm.rank == first_worker_rank: if placements is None or len(placements) != num_parameters: @@ -314,25 +853,12 @@ def publish_destination_placements( for placement in placements ] if rows: - encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=comm.device)) + encoded.copy_(torch.tensor(rows, dtype=torch.int8, device=device)) - comm.broadcast(encoded, src=first_worker_rank) + comm.broadcast(encoded, encoded, root=first_worker_rank, stream=stream) return [ REPLICATED if all(code == replicated_sentinel for code in row) else cast(Placements, tuple(row)) for row in encoded.tolist() ] - - -def comm_ptr(comm: "PyNcclCommunicator") -> int: - """Raw `ncclComm_t` behind vLLM's `PyNcclCommunicator`. - - m2n links its own NCCL, so the handle only means anything if vLLM loaded - the same `libnccl.so` — set `VLLM_NCCL_SO_PATH` accordingly. - """ - handle = comm.comm - ptr = getattr(handle, "value", handle) - if not ptr: - raise RuntimeError("PyNcclCommunicator has no live NCCL communicator") - return int(ptr) diff --git a/vllm/distributed/weight_transfer/m2n_engine.py b/vllm/distributed/weight_transfer/m2n_engine.py index e480e60eaf..66dbb6ba24 100644 --- a/vllm/distributed/weight_transfer/m2n_engine.py +++ b/vllm/distributed/weight_transfer/m2n_engine.py @@ -20,7 +20,7 @@ inside `receive_weights` while the trainer is sending. Driving that from the trainer is the trainer engine's job. """ -from collections.abc import Iterator +from collections.abc import Iterator, Mapping from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, cast @@ -34,20 +34,28 @@ from vllm.distributed.weight_transfer.base import ( ) from vllm.distributed.weight_transfer.m2n_common import ( DESTINATION_SHARD_AXIS, + M2N_WIRE_SCHEMA_VERSION, MESH_NDIMS, - REPLICATED, + M2NLayout, M2NMesh, M2NParamMeta, + M2NWireParam, Placements, - check_plan_limits, + check_complete_plan_agreement, + check_data_plane_agreement, + check_runtime_ready_agreement, + check_source_plan_agreement, check_transferable, - comm_ptr, + create_m2n_communicator, + expected_local_shape, import_m2n, + prepare_m2n_local_runtime, publish_destination_placements, resolve_layout, + source_plan_digest, to_mesh, to_placements, - validate_layout, + validate_transfer_plan, ) from vllm.distributed.weight_transfer.m2n_layout import ( M2NDestination, @@ -56,12 +64,14 @@ from vllm.distributed.weight_transfer.m2n_layout import ( from vllm.distributed.weight_transfer.nccl_common import ( NCCLWeightTransferInitInfo, decode_nccl_unique_id, - worker_init_process_group, + worker_init_metadata_group, ) +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.config import VllmConfig - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + +logger = init_logger(__name__) __all__ = [ "M2NWeightTransferInitInfo", @@ -84,59 +94,99 @@ class M2NWeightTransferInitInfo(WeightTransferInitInfo): plain JSON so the HTTP control plane carries it unchanged. """ - master_address: str | None = None - master_port: int | None = None + schema_version: int + master_address: str + master_port: int rank_offset: int """First worker rank, i.e. the number of trainer ranks.""" world_size: int """Trainer ranks + all inference workers.""" - src_mesh_dims: list[int] - """The trainer's mesh, shared by every parameter (`start_rank` is 0: the - trainer occupies the front of the communicator).""" dst_mesh_dims: list[int] """The inference mesh, starting at `rank_offset`. Declared rather than derived so both sides describe the destination identically.""" - names: list[str] - dtype_names: list[str] - shapes: list[list[int]] - src_placements: list[list[int] | None] - """Per parameter, relative to `src_mesh_dims`; `None` means replicated.""" + source_digest: str + """SHA-256 of the ordered, versioned source parameter plan.""" + params: list[dict[str, Any]] + """Ordered per-parameter source metadata and layouts.""" nccl_unique_id_b64: str | None = field(default=None, repr=False) max_cta: int | None = None def __post_init__(self) -> None: - _ = self.nccl_unique_id_bytes - num_params = len(self.names) - for label, values in ( - ("dtype_names", self.dtype_names), - ("shapes", self.shapes), - ("src_placements", self.src_placements), + if not isinstance(self.master_address, str) or not self.master_address: + raise ValueError("`master_address` must name the metadata rendezvous") + if ( + not isinstance(self.master_port, int) + or isinstance(self.master_port, bool) + or not 0 < self.master_port < 65536 + ): + raise ValueError("`master_port` must be a valid metadata port") + if ( + not isinstance(self.rank_offset, int) + or isinstance(self.rank_offset, bool) + or not isinstance(self.world_size, int) + or isinstance(self.world_size, bool) + or self.rank_offset < 1 + or self.rank_offset >= self.world_size ): - if len(values) != num_params: - raise ValueError( - f"`{label}` should be of the same size as `names`: " - f"got {len(values)} and {num_params}" - ) - if self.rank_offset < 1 or self.rank_offset >= self.world_size: raise ValueError( f"`rank_offset` ({self.rank_offset}) must leave at least one " f"trainer rank and one worker in world_size {self.world_size}" ) num_workers = self.world_size - self.rank_offset dst = tuple(self.dst_mesh_dims) - if len(dst) != MESH_NDIMS or dst[0] * dst[1] != num_workers: + try: + dst_mesh = M2NMesh(cast(tuple[int, int], dst), self.rank_offset) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid `dst_mesh_dims` {dst}: {exc}") from exc + if dst_mesh.size != num_workers: raise ValueError( f"`dst_mesh_dims` {dst} must be {MESH_NDIMS} dims covering the " f"{num_workers} inference workers" ) + def parse_wire_params(self) -> tuple[M2NWireParam, ...]: + """Parse and authenticate the ordered source plan.""" + if self.schema_version != M2N_WIRE_SCHEMA_VERSION: + raise ValueError( + "unsupported nccl_m2n init schema version " + f"{self.schema_version}; expected {M2N_WIRE_SCHEMA_VERSION}" + ) + if not isinstance(self.params, list): + raise TypeError("`params` must be a list of wire parameter mappings") + if not isinstance(self.source_digest, str) or not self.source_digest: + raise ValueError("`source_digest` must be a non-empty string") + wire_params: list[M2NWireParam] = [] + names: set[str] = set() + for index, value in enumerate(self.params): + if not isinstance(value, Mapping): + raise TypeError( + f"`params[{index}]` must be a mapping, got {type(value).__name__}" + ) + param = M2NWireParam.from_dict(value) + if param.name in names: + raise ValueError( + f"nccl_m2n init params contain duplicate parameter '{param.name}'" + ) + names.add(param.name) + wire_params.append(param) + parsed = tuple(wire_params) + computed_digest = source_plan_digest(parsed) + if self.source_digest != computed_digest: + raise ValueError( + "nccl_m2n init source digest does not match params: " + f"computed {computed_digest}, declared {self.source_digest}" + ) + return parsed + @property def nccl_unique_id_bytes(self) -> bytes | None: + if self.nccl_unique_id_b64 is None: + return None return decode_nccl_unique_id( - master_address=self.master_address, - master_port=self.master_port, + master_address=None, + master_port=None, nccl_unique_id_b64=self.nccl_unique_id_b64, - ctx="M2NWeightTransferInitInfo", + ctx="M2NWeightTransferInitInfo data plane", ) @@ -182,33 +232,172 @@ class M2NWeightTransferEngine( model: torch.nn.Module, ) -> None: super().__init__(config, vllm_config, device, model) - self.model_update_group: PyNcclCommunicator | None = None + self.model_update_group: Any = None self._m2n: Any = None self._handle: Any = None self._metas: list[M2NParamMeta] = [] - self._src_mesh: M2NMesh | None = None self._dst_mesh: M2NMesh | None = None self._parameter_destinations: list[M2NDestination] = [] self._index: dict[str, int] = {} self._uses_load_weights = False def init_transfer_engine(self, init_info: M2NWeightTransferInitInfo) -> None: - """Join the trainer's communicator and rebuild the transfer plan. - - Every precondition (dtype, tensor rank, divisibility) is checked here so - a bad plan fails the init RPC instead of hanging the first collective. - """ - self._m2n = import_m2n() - self._metas = [] - - self._src_mesh = M2NMesh( - cast(tuple[int, int], tuple(init_info.src_mesh_dims)), 0 - ) - if self._src_mesh.size != init_info.rank_offset: - raise ValueError( - f"source mesh covers {self._src_mesh.size} ranks, but there are " - f"{init_info.rank_offset} trainer ranks" + """Agree on source and local readiness before constructing NCCL.""" + if self._handle is not None or self.model_update_group is not None: + raise RuntimeError( + "nccl_m2n weight transfer is already initialized; call " + "shutdown_weight_transfer_engine() before initializing again" ) + rendezvous_info = NCCLWeightTransferInitInfo( + master_address=init_info.master_address, + master_port=init_info.master_port, + rank_offset=init_info.rank_offset, + world_size=init_info.world_size, + ) + metadata_group = worker_init_metadata_group( + rendezvous_info, self.parallel_config + ) + + wire_params: tuple[M2NWireParam, ...] = () + source_error: str | None = None + try: + self._m2n = import_m2n() + wire_params = init_info.parse_wire_params() + self._prepare_source_plan(wire_params, init_info.rank_offset) + except Exception as exc: + source_error = f"{type(exc).__name__}: {exc}" + self._metas = [] + check_source_plan_agreement( + metadata_group, + wire_params, + expected_digest=init_info.source_digest, + local_error=source_error, + ) + + destination_error: str | None = None + try: + self._prepare_destination_plan(init_info) + except Exception as exc: + destination_error = f"{type(exc).__name__}: {exc}" + self._parameter_destinations = [] + + unique_id_bytes: bytes | None = None + data_plane_error: str | None = None + try: + unique_id_bytes = init_info.nccl_unique_id_bytes + except Exception as exc: + data_plane_error = f"{type(exc).__name__}: {exc}" + check_data_plane_agreement( + metadata_group, + unique_id_bytes, + data_plane_error, + ) + + handle = None + group = None + try: + runtime_error = destination_error + nccl = None + device = None + if runtime_error is None: + try: + nccl, device, handle = prepare_m2n_local_runtime( + self._m2n, init_info.max_cta + ) + except Exception as exc: + runtime_error = f"{type(exc).__name__}: {exc}" + check_runtime_ready_agreement(metadata_group, runtime_error) + assert nccl is not None and device is not None and handle is not None + group = create_m2n_communicator(nccl, metadata_group, unique_id_bytes) + + # Keep the merged destination-placement publication unchanged: source + # agreement is independent of worker-side destination selection. + mine = [ + destination.placements for destination in self._parameter_destinations + ] + agreed = publish_destination_placements( + group, + self.device, + torch.cuda.current_stream(), + init_info.rank_offset, + mine, + len(self._metas), + ) + complete_plan_error: str | None = None + try: + if agreed != mine: + mismatched = next( + meta.name + for meta, ours, theirs in zip(self._metas, mine, agreed) + if ours != theirs + ) + raise RuntimeError( + "inference workers resolved different nccl_m2n " + f"destinations (first mismatch: '{mismatched}'); all " + "workers must run the same model and parallel config" + ) + assert self._dst_mesh is not None + validate_transfer_plan(self._metas, self._dst_mesh, agreed) + except Exception as exc: + complete_plan_error = f"{type(exc).__name__}: {exc}" + assert self._dst_mesh is not None + check_complete_plan_agreement( + metadata_group, + init_info.source_digest, + self._dst_mesh, + agreed, + complete_plan_error, + ) + except BaseException: + if group is not None: + try: + # Transport cleanup only; reload validity belongs to the + # model-update lifecycle, not this backend. + group.abort() + except BaseException: + logger.exception( + "failed to abort partially initialized nccl_m2n communicator" + ) + if handle is not None: + try: + handle.destroy() + except BaseException: + logger.exception( + "failed to destroy partially initialized nccl_m2n handle" + ) + raise + + self._handle = handle + self.model_update_group = group + + def _prepare_source_plan( + self, + wire_params: tuple[M2NWireParam, ...], + num_trainer_ranks: int, + ) -> None: + """Validate worker-visible source metadata without constructing NCCL.""" + self._metas = [] + for param in wire_params: + dtype = getattr(torch, param.dtype_name, None) + if not isinstance(dtype, torch.dtype): + raise ValueError( + f"parameter '{param.name}' has invalid dtype name " + f"'{param.dtype_name}'; expected the name of a torch.dtype" + ) + check_transferable(param.name, dtype, param.shape) + layout = M2NLayout(M2NMesh(param.src_mesh_dims, 0), param.src_placements) + if layout.mesh.size != num_trainer_ranks: + raise ValueError( + f"parameter '{param.name}' source mesh covers " + f"{layout.mesh.size} ranks, but there are " + f"{num_trainer_ranks} trainer ranks" + ) + meta = M2NParamMeta(param.name, dtype, param.shape, layout) + expected_local_shape(meta) + self._metas.append(meta) + + def _prepare_destination_plan(self, init_info: M2NWeightTransferInitInfo) -> None: + """Resolve the merged-#51520 in-place/full-fallback destination plan.""" # The inference topology, as declared by the trainer. This is the mesh # a sharded destination is placed over; a replicated one is described # by the alternate descriptor `resolve_layout` derives from it, which @@ -217,33 +406,6 @@ class M2NWeightTransferEngine( cast(tuple[int, int], tuple(init_info.dst_mesh_dims)), init_info.rank_offset, ) - - for name, name_dtype, shape, placements in zip( - init_info.names, - init_info.dtype_names, - init_info.shapes, - init_info.src_placements, - ): - dtype = getattr(torch, name_dtype, None) - if not isinstance(dtype, torch.dtype): - raise ValueError( - f"parameter '{name}' has invalid dtype name " - f"'{name_dtype}'; expected the name of a torch.dtype" - ) - check_transferable(name, dtype, shape) - codes = ( - REPLICATED - if placements is None - else cast(Placements, tuple(placements)) - ) - src_mesh, src_placements = resolve_layout( - self._src_mesh, codes, f"parameter '{name}' source placements" - ) - validate_layout(src_mesh, src_placements, shape, "source") - # Destination checks are deferred until loader-aware resolution; - # assuming replication here can reject a valid sharded plan. - self._metas.append(M2NParamMeta(name, dtype, tuple(shape), codes)) - parallel_config = self.parallel_config shard_axis_size = self._dst_mesh.dims[DESTINATION_SHARD_AXIS] quantization = getattr(self.model_config, "quantization", None) @@ -266,27 +428,11 @@ class M2NWeightTransferEngine( shard_axis_size=shard_axis_size, allow_direct=allow_direct, ) - # Check the resolved layout rather than a provisional replicated one. - # For example, 32 source shards feeding 4 destination shards may need - # only 8 sources per destination, while a replicated plan needs all 32. - for meta, destination in zip(self._metas, self._parameter_destinations): - src_mesh, src_placements = resolve_layout( - self._src_mesh, - meta.placements, - f"parameter '{meta.name}' source placements", - ) - dst_mesh, dst_placements = resolve_layout( - self._dst_mesh, - destination.placements, - f"parameter '{meta.name}' destination placements", - ) - validate_layout(dst_mesh, dst_placements, meta.shape, "destination") - check_plan_limits( - (src_mesh, src_placements), - (dst_mesh, dst_placements), - meta.name, - ) - + validate_transfer_plan( + self._metas, + self._dst_mesh, + [destination.placements for destination in self._parameter_destinations], + ) # Update requests carry names only, so cache their plan indices. Track # whether any fallback entry requires the `load_weights` lifecycle. self._index = {meta.name: i for i, meta in enumerate(self._metas)} @@ -294,38 +440,6 @@ class M2NWeightTransferEngine( not destination.direct for destination in self._parameter_destinations ) - self.model_update_group = worker_init_process_group( - NCCLWeightTransferInitInfo( - master_address=init_info.master_address, - master_port=init_info.master_port, - nccl_unique_id_b64=init_info.nccl_unique_id_b64, - rank_offset=init_info.rank_offset, - world_size=init_info.world_size, - ), - parallel_config, - ) - # Destinations are per-parameter and depend on the inference model, so - # the trainer cannot derive them. The first worker publishes the plan - # and every other rank checks its own against it — a silent - # disagreement between workers would mean mismatched collectives. - mine = [destination.placements for destination in self._parameter_destinations] - agreed = publish_destination_placements( - self.model_update_group, init_info.rank_offset, mine, len(self._metas) - ) - if agreed != mine: - mismatched = next( - meta.name for meta, a, b in zip(self._metas, mine, agreed) if a != b - ) - raise RuntimeError( - "inference workers resolved different nccl_m2n destinations " - f"(first mismatch: '{mismatched}'); all workers must run the " - "same model and parallel config" - ) - - self._handle = self._m2n.Handle.create( - self._m2n.Config(max_cta=init_info.max_cta) - ) - def start_weight_update(self) -> None: """Set up layerwise reloading, but only if some parameter needs it. @@ -335,7 +449,9 @@ class M2NWeightTransferEngine( """ if not self._uses_load_weights: return - from vllm.model_executor.model_loader.reload import initialize_layerwise_reload + from vllm.model_executor.model_loader.reload import ( + initialize_layerwise_reload, + ) initialize_layerwise_reload(self.model) @@ -354,6 +470,10 @@ class M2NWeightTransferEngine( "nccl_m2n weight transfer not initialized. " "Call init_transfer_engine() first." ) + if not isinstance(update_info.names, list) or any( + not isinstance(name, str) for name in update_info.names + ): + raise TypeError("nccl_m2n update names must be a list of strings") requested = [] for name in update_info.names: @@ -369,7 +489,6 @@ class M2NWeightTransferEngine( disable_mtp_completeness_check, ) - comm = comm_ptr(self.model_update_group) stream = torch.cuda.current_stream() # Reshard lazily so one `load_weights` invocation sees the complete @@ -385,7 +504,13 @@ class M2NWeightTransferEngine( meta.shape, dtype=meta.dtype, device=self.device ) - self._reshard(comm, stream, meta, destination.placements, buffer) + self._reshard( + self.model_update_group, + stream, + meta, + destination.placements, + buffer, + ) if destination.tensor is None: # `load_weights` reads on the host stream, so the transfer @@ -410,7 +535,7 @@ class M2NWeightTransferEngine( def _reshard( self, - comm: int, + comm: Any, stream: torch.cuda.Stream, src_meta: M2NParamMeta, dst_placements: Placements | None, @@ -421,11 +546,10 @@ class M2NWeightTransferEngine( `dst_placements` describes the buffer's logical placement over the worker mesh; `dst_buffer` is this rank's physical storage. """ - assert self._src_mesh is not None assert self._dst_mesh is not None m2n = self._m2n src_mesh, resolved_src_placements = resolve_layout( - self._src_mesh, src_meta.placements + src_meta.source_layout.mesh, src_meta.source_layout.placements ) dst_mesh, resolved_dst_placements = resolve_layout( self._dst_mesh, dst_placements @@ -444,12 +568,32 @@ class M2NWeightTransferEngine( def shutdown(self) -> None: """Finish pending GPU work and release M2N communication state.""" - if self._handle is not None: - # M2N does not synchronize caller streams on finalize. - torch.accelerator.synchronize() - self._handle.destroy() - self._handle = None - if self.model_update_group is not None: - self.model_update_group.destroy() - self.model_update_group = None + handle, self._handle = self._handle, None + group, self.model_update_group = self.model_update_group, None + self._m2n = None + self._metas = [] + self._dst_mesh = None self._parameter_destinations = [] + self._index = {} + self._uses_load_weights = False + + first_error: BaseException | None = None + + def attempt(action: Any, message: str) -> None: + nonlocal first_error + try: + action() + except BaseException as exc: + if first_error is None: + first_error = exc + logger.exception(message) + + if handle is not None or group is not None: + # M2N does not synchronize caller streams on finalize. + attempt(torch.accelerator.synchronize, "failed to synchronize nccl_m2n") + if handle is not None: + attempt(handle.destroy, "failed to destroy nccl_m2n handle") + if group is not None: + attempt(group.abort, "failed to abort nccl_m2n communicator") + if first_error is not None: + raise first_error diff --git a/vllm/distributed/weight_transfer/m2n_source.py b/vllm/distributed/weight_transfer/m2n_source.py index dd08d24b56..d9c9e9194a 100644 --- a/vllm/distributed/weight_transfer/m2n_source.py +++ b/vllm/distributed/weight_transfer/m2n_source.py @@ -5,21 +5,22 @@ m2n plans a transfer from both sides' layouts, so the trainer has to say how it holds each parameter — which the base `WeightSource` / `ParamMeta` pair does not express. This module supplies the m2n flavor of both, and the DTensor-backed -implementation that covers the common trainer. - -A source declares its mesh once (`mesh()`) and a placement per parameter, since -one topology describes every tensor on the side. +implementation that covers the common trainer. Each parameter owns its source +mesh and placements because dense and expert weights can use different +factorizations over the same trainer ranks. """ -from collections.abc import Iterator +from collections.abc import Callable, Iterator, Sequence +from dataclasses import dataclass from typing import Any import torch -from vllm.distributed.weight_transfer.base import ParamMeta, WeightSource +from vllm.distributed.weight_transfer.base import WeightSource from vllm.distributed.weight_transfer.m2n_common import ( REPLICATE, REPLICATED, + M2NLayout, M2NMesh, M2NParamMeta, Placements, @@ -27,6 +28,8 @@ from vllm.distributed.weight_transfer.m2n_common import ( __all__ = [ "M2NWeightSource", + "M2NManifestEntry", + "ManifestM2NWeightSource", "DTensorModuleSource", "mesh_from_tensor", "placements_from_tensor", @@ -40,17 +43,76 @@ class M2NWeightSource(WeightSource): materialized full tensor: gathering is exactly the cost m2n removes. """ - def mesh(self) -> M2NMesh: - """The trainer's rank topology, shared by every parameter.""" + def metadata(self) -> list[M2NParamMeta]: # type: ignore[override] + """Name, dtype, full shape, and source layout for each parameter.""" raise NotImplementedError - def metadata(self) -> list[ParamMeta]: - """Name, dtype, full shape, and trainer placement for each parameter.""" + def __iter__(self) -> Iterator[tuple[str, Any]]: # type: ignore[override] + """Yield `(name, local buffer)` pairs in metadata order. + + Providers must enqueue materialization on the caller's current CUDA + stream, or return a buffer whose writes are already visible to it. + Buffers may be torch tensors, CUDA Array Interface objects, or raw + device pointers accepted by nccl4py. Providers own correct + communicator-device placement. Raw-pointer providers additionally own + allocation lifetime, dtype, and shape correctness. + """ raise NotImplementedError - def __iter__(self) -> Iterator[tuple[str, torch.Tensor]]: - """Yield `(name, local shard)` pairs in the same order as `metadata()`.""" - raise NotImplementedError + +@dataclass(frozen=True) +class M2NManifestEntry: + """One immutable source descriptor with a lazy rank-local tensor view.""" + + key: str + dtype: torch.dtype + global_shape: tuple[int, ...] + source_layout: M2NLayout + local_tensor: Callable[[], Any] + + def __post_init__(self) -> None: + object.__setattr__(self, "global_shape", tuple(self.global_shape)) + if not self.key: + raise ValueError("M2N manifest key must not be empty") + if not isinstance(self.dtype, torch.dtype): + raise TypeError( + f"manifest entry '{self.key}' has invalid dtype {self.dtype!r}" + ) + if not isinstance(self.source_layout, M2NLayout): + raise TypeError( + f"manifest entry '{self.key}' source_layout must be an M2NLayout" + ) + if not callable(self.local_tensor): + raise TypeError( + f"manifest entry '{self.key}' local_tensor must be callable" + ) + + @property + def metadata(self) -> M2NParamMeta: + """Return the serializable portion consumed during initialization.""" + return M2NParamMeta( + self.key, + self.dtype, + self.global_shape, + self.source_layout, + ) + + +class ManifestM2NWeightSource(M2NWeightSource): + """An ordered M2N source backed by lazy, rank-local manifest entries.""" + + def __init__(self, entries: Sequence[M2NManifestEntry]) -> None: + self._entries = tuple(entries) + self._metadata = tuple(entry.metadata for entry in self._entries) + + def metadata(self) -> list[M2NParamMeta]: # type: ignore[override] + """Return a fresh list while preserving the immutable manifest order.""" + return list(self._metadata) + + def __iter__(self) -> Iterator[tuple[str, Any]]: # type: ignore[override] + """Materialize one rank-local buffer at a time.""" + for entry in self._entries: + yield entry.key, entry.local_tensor() def _placement_code(placement: Any) -> int: @@ -69,7 +131,8 @@ def _placement_code(placement: Any) -> int: def mesh_from_tensor(tensor: torch.Tensor, num_trainer_ranks: int) -> M2NMesh: """The mesh a parameter lives on, as an `M2NMesh`. - A plain tensor is identical on every trainer rank, so it spans all of them. + Tensors without DeviceMesh metadata are treated as replicated. Implicitly + sharded tensors require a custom source with an explicit layout. """ device_mesh = getattr(tensor, "device_mesh", None) if device_mesh is None: @@ -122,29 +185,18 @@ class DTensorModuleSource(M2NWeightSource): self._module = module self._num_trainer_ranks = num_trainer_ranks - def mesh(self) -> M2NMesh: - """The mesh every sharded parameter agrees on. - - Replicated parameters span the trainer without a device mesh of their - own, so they never decide this; a model whose *sharded* parameters - disagree cannot be described by one side mesh and is rejected. - """ - meshes = { - mesh_from_tensor(p, self._num_trainer_ranks) - for _, p in self._module.named_parameters() - if placements_from_tensor(p) is not REPLICATED - } - if len(meshes) > 1: - raise ValueError( - "nccl_m2n needs one mesh per side, but this module's sharded " - f"parameters span several: {sorted(m.dims for m in meshes)}" - ) - return meshes.pop() if meshes else M2NMesh((self._num_trainer_ranks, 1), 0) - - def metadata(self) -> list[ParamMeta]: - """Read global shape/dtype and placements without gathering shards.""" + def metadata(self) -> list[M2NParamMeta]: # type: ignore[override] + """Read each global shape and source layout without gathering shards.""" return [ - M2NParamMeta(name, p.dtype, tuple(p.shape), placements_from_tensor(p)) + M2NParamMeta( + name, + p.dtype, + tuple(p.shape), + M2NLayout( + mesh_from_tensor(p, self._num_trainer_ranks), + placements_from_tensor(p), + ), + ) for name, p in self._module.named_parameters() ] diff --git a/vllm/distributed/weight_transfer/m2n_trainer.py b/vllm/distributed/weight_transfer/m2n_trainer.py index 8b9b73f9f9..19dd973c73 100644 --- a/vllm/distributed/weight_transfer/m2n_trainer.py +++ b/vllm/distributed/weight_transfer/m2n_trainer.py @@ -13,6 +13,8 @@ engine -- where the `TrainerWeightTransferEngine` contract puts it -- rather tha in every caller. """ +import socket +from collections.abc import Mapping from concurrent.futures import Future, ThreadPoolExecutor from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, ClassVar, Self, cast @@ -26,39 +28,43 @@ from vllm.distributed.weight_transfer.base import ( WeightSource, ) from vllm.distributed.weight_transfer.m2n_common import ( + M2N_WIRE_SCHEMA_VERSION, M2NMesh, M2NParamMeta, + M2NWireParam, Placements, + check_complete_plan_agreement, + check_data_plane_agreement, + check_runtime_ready_agreement, + check_source_plan_agreement, check_transferable, - comm_ptr, + create_m2n_communicator, + expected_local_shape, import_m2n, + prepare_m2n_local_runtime, publish_destination_placements, resolve_layout, + source_plan_digest, to_mesh, to_placements, - validate_layout, + validate_local_buffer, + validate_transfer_plan, ) from vllm.distributed.weight_transfer.m2n_source import M2NWeightSource from vllm.distributed.weight_transfer.nccl_common import ( decode_nccl_unique_id, - trainer_init, - uid_init_process_group, + stateless_init_metadata_group, ) from vllm.logger import init_logger if TYPE_CHECKING: - from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup logger = init_logger(__name__) __all__ = ["M2NTrainerInitInfo", "M2NTrainerWeightTransferEngine"] -def _dtype_name(dtype: torch.dtype) -> str: - """Short name for a dtype, e.g. `bfloat16`, for the worker-side init dict.""" - return str(dtype).split(".")[-1] - - @dataclass(kw_only=True) class M2NTrainerInitInfo(TrainerInitInfo): """Trainer-side init info for nccl_m2n. @@ -70,8 +76,8 @@ class M2NTrainerInitInfo(TrainerInitInfo): backend: ClassVar[str] = "nccl_m2n" - master_address: str | None = None - master_port: int | None = None + master_address: str + master_port: int world_size: int """Trainer ranks + all inference workers.""" num_trainer_ranks: int = 1 @@ -79,20 +85,34 @@ class M2NTrainerInitInfo(TrainerInitInfo): """How the inference ranks are laid out: axis 0 replicates and axis 1 shards. The trainer declares it so both sides describe the destination identically; defaults to a flat `(num_workers, 1)`, which is all a - replicated destination needs.""" - max_cta: int | None = None + replicated destination needs. + + Destination meshes remain common to every parameter; per-parameter + destination meshes are intentionally future work. + """ nccl_unique_id_b64: str | None = field(default=None, repr=False) + """Optional UID for the NCCL data plane; TCP remains the metadata plane.""" + max_cta: int | None = None + listen_socket: socket.socket | None = field(default=None, repr=False, compare=False) + """Optional rank-zero TCPStore reservation; never sent to workers.""" @property def destination_mesh_dims(self) -> tuple[int, int]: """`dst_mesh_dims`, or a flat mesh over every inference worker.""" if self.dst_mesh_dims is not None: - return self.dst_mesh_dims + return cast(tuple[int, int], tuple(self.dst_mesh_dims)) return (self.world_size - self.num_trainer_ranks, 1) def __post_init__(self) -> None: """Reject rank counts or destination meshes that cannot form a group.""" - _ = self.nccl_unique_id_bytes + if not isinstance(self.master_address, str) or not self.master_address: + raise ValueError("`master_address` must name the metadata rendezvous") + if ( + not isinstance(self.master_port, int) + or isinstance(self.master_port, bool) + or not 0 < self.master_port < 65536 + ): + raise ValueError("`master_port` must be a valid metadata port") if not 0 < self.num_trainer_ranks < self.world_size: raise ValueError( f"`num_trainer_ranks` ({self.num_trainer_ranks}) must leave at " @@ -101,7 +121,11 @@ class M2NTrainerInitInfo(TrainerInitInfo): num_workers = self.world_size - self.num_trainer_ranks if self.dst_mesh_dims is not None: dims = tuple(self.dst_mesh_dims) - if len(dims) != 2 or dims[0] * dims[1] != num_workers: + try: + mesh = M2NMesh(cast(tuple[int, int], dims), self.num_trainer_ranks) + except (TypeError, ValueError) as exc: + raise ValueError(f"invalid `dst_mesh_dims` {dims}: {exc}") from exc + if mesh.size != num_workers: raise ValueError( f"`dst_mesh_dims` {dims} must be two dims covering the " f"{num_workers} inference workers" @@ -111,14 +135,23 @@ class M2NTrainerInitInfo(TrainerInitInfo): f"trainer `rank` ({self.rank}) must be below " f"`num_trainer_ranks` ({self.num_trainer_ranks})" ) + if self.listen_socket is not None: + if self.rank != 0: + raise ValueError("only trainer rank zero may reserve the TCPStore") + if not isinstance(self.listen_socket, socket.socket): + raise TypeError("listen_socket must be a socket.socket") + if self.listen_socket.fileno() < 0: + raise ValueError("listen_socket is already closed") @property def nccl_unique_id_bytes(self) -> bytes | None: + if self.nccl_unique_id_b64 is None: + return None return decode_nccl_unique_id( - master_address=self.master_address, - master_port=self.master_port, + master_address=None, + master_port=None, nccl_unique_id_b64=self.nccl_unique_id_b64, - ctx="M2NTrainerInitInfo", + ctx="M2NTrainerInitInfo data plane", ) @@ -136,22 +169,20 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI *, client: VLLMWeightSyncClient, source: M2NWeightSource, - is_controller: bool, - num_trainer_ranks: int, + is_sender: bool, ) -> None: """Hold the client and source; `trainer_init` fills in the group.""" - super().__init__(client=client, source=source, is_sender=is_controller) - self.source: M2NWeightSource = source - self.is_controller = is_controller - self.num_trainer_ranks = num_trainer_ranks - self.group: PyNcclCommunicator | None = None + super().__init__(client=client, source=source, is_sender=is_sender) + self.group: Any = None self._m2n: Any = None self._handle: Any = None self._metas: list[M2NParamMeta] = [] - self._src_mesh: M2NMesh | None = None + self._wire_params: list[M2NWireParam] = [] + self._source_digest = "" self._dst_mesh: M2NMesh | None = None self._dst_placements: list[Placements | None] = [] - self._executor: ThreadPoolExecutor | None = None + self._metadata_group: StatelessProcessGroup | None = None + self._num_trainer_ranks = 0 @classmethod def trainer_init( @@ -168,204 +199,351 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI build local state and join the communicator. Every rank participates in every reshard and sends its local shard. - The ordering here is not arbitrary -- see the comments inline. In - short: validate before anything blocks; start the init RPC without - waiting (workers join the communicator inside that RPC, so it cannot - return until this side has joined too); then join the NCCL - communicator; then create the M2N handle; and only then wait for the RPC - to finish. + Build a local validation envelope, start the worker RPC without + waiting, join the metadata rendezvous, and require unanimous source, + data-plane, and runtime readiness before constructing NCCL. """ - # The base WeightSource carries name/dtype/shape but not layout, which - # is not enough to plan a reshard. Fail here rather than at the first - # missing attribute deep inside the send loop. - if not isinstance(source, M2NWeightSource): - raise TypeError( - "nccl_m2n needs per-parameter layouts, so its source must be an " - f"M2NWeightSource (e.g. DTensorModuleSource); got " - f"{type(source).__name__}" - ) - - # Trainer-side weight transfer engine for this rank. `__init__` only - # holds the client and source; the communicator, meshes, and M2N handle - # are filled in below. engine = cls( client=client, - source=source, - is_controller=init_info.rank == 0, - num_trainer_ranks=init_info.num_trainer_ranks, + source=cast(M2NWeightSource, source), + is_sender=init_info.is_sender, ) - engine._m2n = import_m2n() - - engine._src_mesh = source.mesh() - if engine._src_mesh.size != init_info.num_trainer_ranks: - raise ValueError( - f"source mesh covers {engine._src_mesh.size} ranks, but there " - f"are {init_info.num_trainer_ranks} trainer ranks" - ) - - # Validate the whole plan before any rendezvous. A layout m2n cannot - # express is a mismatched collective at send time, which hangs rather - # than raises -- so every precondition is checked while an exception can - # still propagate to the caller. - engine._metas = cast(list[M2NParamMeta], source.metadata()) - for meta in engine._metas: - check_transferable(meta.name, meta.dtype, meta.shape) - mesh, placements = resolve_layout(engine._src_mesh, meta.placements) - validate_layout(mesh, placements, meta.shape, "source") - - rendezvous: Future | None = None - if engine.is_controller: - engine._executor = ThreadPoolExecutor( - max_workers=1, thread_name_prefix="m2n-weight-sync" - ) - # Only trainer rank 0 starts the control-plane RPC. Workers join - # the communicator inside it, so submit it first, join below, and - # wait for it at the end. - rendezvous = engine._executor.submit( - client.init_weight_transfer_engine, - engine._worker_init_info(init_info), - ) - - # Every trainer rank joins -- reshard is collective over the whole - # communicator. Trainer ranks take - # [0, num_trainer_ranks) and the workers follow, which is the - # contiguous-interval layout m2n requires of both meshes. - unique_id_bytes = init_info.nccl_unique_id_bytes - if unique_id_bytes is not None: - engine.group = uid_init_process_group( - unique_id_bytes, - rank=init_info.rank, - world_size=init_info.world_size, - device=torch.accelerator.current_device_index(), - ) - else: - assert init_info.master_address is not None - assert init_info.master_port is not None - engine.group = trainer_init( - { - "master_address": init_info.master_address, - "master_port": init_info.master_port, - "world_size": init_info.world_size, - }, - rank=init_info.rank, - ) - - # The trainer declares the inference topology; the worker uses exactly - # this, so neither side has to infer the other's factorization. - engine._dst_mesh = M2NMesh( - init_info.destination_mesh_dims, init_info.num_trainer_ranks - ) - engine._dst_placements = engine._receive_destination_plan( - init_info.num_trainer_ranks - ) - logger.info( - "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), " - "inference ranks [%d, %d)", - len(engine._metas), - init_info.num_trainer_ranks, - init_info.num_trainer_ranks, - init_info.world_size, - ) - - # Created after the group so a failed rendezvous does not leave a live - # handle behind; shutdown() releases it once caller streams are synced. - engine._handle = engine._m2n.Handle.create( - engine._m2n.Config(max_cta=init_info.max_cta) - ) - - # Now safe to collect: both sides have joined, so the RPC has returned - # or raised. Doing this last means a worker-side init failure surfaces - # here, as an exception from trainer_init, rather than at the first send. - if rendezvous is not None: - rendezvous.result() + engine._initialize(init_info, source) return engine - def _receive_destination_plan( - self, first_worker_rank: int - ) -> list[Placements | None]: - """Receive and validate the inference-side per-parameter placements.""" - if self.group is None or self._dst_mesh is None: - raise RuntimeError( - "destination plan requires the communicator and destination mesh" + def _initialize( + self, + init_info: M2NTrainerInitInfo, + source: WeightSource | None, + ) -> None: + """Run source validation, metadata agreement, and NCCL rendezvous.""" + self._num_trainer_ranks = init_info.num_trainer_ranks + local_source_error: str | None = None + try: + if not isinstance(source, M2NWeightSource): + raise TypeError( + "nccl_m2n needs per-parameter layouts, so its source must " + f"be an M2NWeightSource; got {type(source).__name__}" + ) + self._m2n = import_m2n() + self._prepare_source_plan(source, init_info.num_trainer_ranks) + except Exception as exc: + local_source_error = f"{type(exc).__name__}: {exc}" + self._metas = [] + self._wire_params = [] + self._source_digest = source_plan_digest([]) + + executor: ThreadPoolExecutor | None = None + init_future: Future[None] | None = None + if self.is_sender: + executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="m2n-weight-sync" + ) + init_future = executor.submit( + self.client.init_weight_transfer_engine, + self._worker_init_info(init_info), + ) + handle = None + group = None + try: + reservation = init_info.listen_socket + try: + metadata_group = stateless_init_metadata_group( + init_info.master_address, + init_info.master_port, + init_info.rank, + init_info.world_size, + listen_socket=reservation, + ) + finally: + init_info.listen_socket = None + if reservation is not None: + reservation.close() + plan_digest = check_source_plan_agreement( + metadata_group, + self._wire_params, + expected_digest=self._source_digest, + local_error=local_source_error, ) + # Retain the common destination mesh and the merged NCCL publication + # of per-parameter placements. + self._dst_mesh = M2NMesh( + init_info.destination_mesh_dims, init_info.num_trainer_ranks + ) + unique_id_bytes: bytes | None = None + data_plane_error: str | None = None + try: + unique_id_bytes = init_info.nccl_unique_id_bytes + except Exception as exc: + data_plane_error = f"{type(exc).__name__}: {exc}" + check_data_plane_agreement( + metadata_group, + unique_id_bytes, + data_plane_error, + ) + + runtime_error: str | None = None + nccl = None + device = None + try: + nccl, device, handle = prepare_m2n_local_runtime( + self._m2n, init_info.max_cta + ) + except Exception as exc: + runtime_error = f"{type(exc).__name__}: {exc}" + check_runtime_ready_agreement(metadata_group, runtime_error) + assert nccl is not None and device is not None and handle is not None + group = create_m2n_communicator(nccl, metadata_group, unique_id_bytes) + destination_placements = self._receive_destination_plan( + group, + init_info.num_trainer_ranks, + torch.device("cuda", device), + ) + + complete_plan_error: str | None = None + try: + assert self._dst_mesh is not None + validate_transfer_plan( + self._metas, self._dst_mesh, destination_placements + ) + except Exception as exc: + complete_plan_error = f"{type(exc).__name__}: {exc}" + assert self._dst_mesh is not None + check_complete_plan_agreement( + metadata_group, + plan_digest, + self._dst_mesh, + destination_placements, + complete_plan_error, + ) + + if init_future is not None: + init_future.result() + + logger.info( + "nccl_m2n trainer ready: %d parameters, trainer ranks [0, %d), " + "inference ranks [%d, %d)", + len(self._metas), + init_info.num_trainer_ranks, + init_info.num_trainer_ranks, + init_info.world_size, + ) + logger.info("nccl_m2n source plan digest: %s", plan_digest) + except BaseException: + if group is not None: + try: + # Transport cleanup only; reload validity belongs to the + # model-update lifecycle, not this backend. + group.abort() + except BaseException: + logger.exception( + "failed to abort partially initialized nccl_m2n communicator" + ) + if handle is not None: + try: + handle.destroy() + except BaseException: + logger.exception( + "failed to destroy partially initialized nccl_m2n handle" + ) + raise + finally: + if executor is not None: + executor.shutdown(wait=False, cancel_futures=True) + + self._metadata_group = metadata_group + self._handle = handle + self.group = group + self._dst_placements = destination_placements + + def _receive_destination_plan( + self, group: Any, first_worker_rank: int, device: torch.device + ) -> list[Placements | None]: + """Receive the inference-side per-parameter placements.""" # The first inference worker publishes one placement per parameter. # Every trainer rank receives the complete plan in this single, # initialization-time broadcast. placements = publish_destination_placements( - self.group, first_worker_rank, None, len(self._metas) + group, + device, + torch.cuda.current_stream(), + first_worker_rank, + None, + len(self._metas), ) - # The send loop matches placements to metadata by position, so reject - # an incomplete or oversized plan before entering any M2N collective. - if len(placements) != len(self._metas): - raise RuntimeError( - f"inference side planned {len(placements)} parameters, but " - f"this trainer announced {len(self._metas)}" - ) - # Check every worker-selected destination against the corresponding - # global parameter shape before the first transfer. - for meta, placement in zip(self._metas, placements): - mesh, resolved = resolve_layout(self._dst_mesh, placement) - validate_layout(mesh, resolved, meta.shape, "destination") return placements def _worker_init_info(self, init_info: M2NTrainerInitInfo) -> dict[str, Any]: """Handshake payload: rendezvous, both meshes, and the transfer plan.""" - assert self._src_mesh is not None payload: dict[str, Any] = { + "schema_version": M2N_WIRE_SCHEMA_VERSION, "master_address": init_info.master_address, "master_port": init_info.master_port, "nccl_unique_id_b64": init_info.nccl_unique_id_b64, "rank_offset": init_info.num_trainer_ranks, "world_size": init_info.world_size, - "src_mesh_dims": list(self._src_mesh.dims), "dst_mesh_dims": list(init_info.destination_mesh_dims), - "names": [m.name for m in self._metas], - "dtype_names": [_dtype_name(m.dtype) for m in self._metas], - "shapes": [list(m.shape) for m in self._metas], - "src_placements": [ - None if m.placements is None else list(m.placements) - for m in self._metas - ], + "source_digest": self._source_digest, + "params": [param.to_dict() for param in self._wire_params], "max_cta": init_info.max_cta, } return {key: value for key, value in payload.items() if value is not None} + def _prepare_source_plan( + self, source: M2NWeightSource, num_trainer_ranks: int + ) -> None: + """Build and validate this rank's immutable source plan.""" + self._metas = source.metadata() + names: set[str] = set() + for meta in self._metas: + if not isinstance(meta, M2NParamMeta): + raise TypeError( + "nccl_m2n source metadata must contain M2NParamMeta " + f"entries, got {type(meta).__name__}" + ) + if meta.name in names: + raise ValueError( + "nccl_m2n source metadata contains duplicate parameter " + f"'{meta.name}'" + ) + names.add(meta.name) + check_transferable(meta.name, meta.dtype, meta.shape) + layout = meta.source_layout + if layout.mesh.start_rank != 0: + raise ValueError( + f"parameter '{meta.name}' source mesh must start at rank 0, " + f"got {layout.mesh.start_rank}" + ) + if layout.mesh.size != num_trainer_ranks: + raise ValueError( + f"parameter '{meta.name}' source mesh covers " + f"{layout.mesh.size} ranks, but there are " + f"{num_trainer_ranks} trainer ranks" + ) + expected_local_shape(meta) + self._wire_params = [ + M2NWireParam( + name=meta.name, + dtype_name=str(meta.dtype).split(".")[-1], + shape=meta.shape, + src_mesh_dims=meta.source_layout.mesh.dims, + src_placements=meta.source_layout.placements, + ) + for meta in self._metas + ] + self._source_digest = source_plan_digest(self._wire_params) + def send_weights(self) -> None: """Drive one update round: start, reshard concurrently, then finish.""" if self._handle is None or self.group is None: raise RuntimeError("nccl_m2n trainer engine is not initialized") - update: Future | None = None - if self.is_controller: - self.client.start_weight_update() - # The workers reshard from inside `update_weights`, concurrently - # with the sends below. - assert self._executor is not None - update = self._executor.submit( - self.client.update_weights, - {"names": [m.name for m in self._metas]}, - ) + if not self.is_sender: + try: + self._send() + except BaseException: + try: + self._receive_round_outcome() + except BaseException: + logger.exception( + "failed to consume nccl_m2n round outcome after local failure" + ) + raise + self._receive_round_outcome() + return try: - self._send() - finally: - if update is not None: + self.client.start_weight_update() + executor = ThreadPoolExecutor( + max_workers=1, thread_name_prefix="m2n-weight-sync" + ) + try: + update = executor.submit( + self.client.update_weights, + {"names": [m.name for m in self._metas]}, + ) + # Surface an immediate worker-side rejection before trainers + # enter the collective sequence. + if update.done(): + update.result() + self._send() update.result() - - if self.is_controller: + finally: + executor.shutdown(wait=False, cancel_futures=True) self.client.finish_weight_update() + except BaseException as exc: + try: + self._publish_round_outcome(exc) + except BaseException: + logger.exception("failed to publish nccl_m2n round failure") + raise + + self._publish_round_outcome(None) + + def _publish_round_outcome(self, failure: BaseException | None) -> None: + """Tell non-sender trainer ranks whether worker finalization passed.""" + if self._num_trainer_ranks <= 1: + return + if self._metadata_group is None: + raise RuntimeError("nccl_m2n metadata group is unavailable") + envelope = { + "ok": failure is None, + "error": ( + None if failure is None else f"{type(failure).__name__}: {failure}" + ), + } + for rank in range(1, self._num_trainer_ranks): + self._metadata_group.send_obj(envelope, rank) + # Keep rank zero's TCPStore alive until every trainer consumed its + # outcome. This also keeps the per-rank message counters aligned. + expected_ack = {"phase": "round_outcome_consumed"} + for rank in range(1, self._num_trainer_ranks): + ack = self._metadata_group.recv_obj(rank) + if ack != expected_ack: + raise RuntimeError( + f"trainer rank {rank} sent an invalid round-outcome acknowledgement" + ) + + def _receive_round_outcome(self) -> None: + """Wait until trainer rank zero observed every worker's finish RPC.""" + if self._metadata_group is None: + raise RuntimeError("nccl_m2n metadata group is unavailable") + outcome = self._metadata_group.recv_obj(0) + self._metadata_group.send_obj({"phase": "round_outcome_consumed"}, 0) + if not isinstance(outcome, Mapping) or set(outcome) != {"ok", "error"}: + raise RuntimeError("trainer rank zero sent an invalid round outcome") + ok = outcome["ok"] + error = outcome["error"] + if not isinstance(ok, bool) or ok != (error is None): + raise RuntimeError("trainer rank zero sent an invalid round outcome") + if not ok: + if not isinstance(error, str) or not error: + raise RuntimeError("trainer rank zero sent an invalid round error") + raise RuntimeError( + "nccl_m2n worker update failed on trainer rank zero: " + error + ) def _send(self) -> None: """Reshard each local shard into its worker-planned destination.""" assert self.group is not None - assert self._src_mesh is not None assert self._dst_mesh is not None + assert self.source is not None m2n = self._m2n - comm = comm_ptr(self.group) stream = torch.cuda.current_stream() - for index, (name, tensor) in enumerate(self.source): + device_index = torch.accelerator.current_device_index() + # Optimizer work may have used another CUDA stream. Make prior writes + # visible before a provider materializes current-stream local shards. + torch.accelerator.synchronize(device_index) + inflight_buffers: list[Any] = [] + sent = 0 + for index, (name, buffer) in enumerate(self.source): + if index >= len(self._metas): + raise RuntimeError( + "weight source yielded more parameters than metadata() " + f"declared (first extra parameter: '{name}')" + ) meta = self._metas[index] if name != meta.name: raise RuntimeError( @@ -373,28 +551,47 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI f"its metadata declared '{meta.name}'; the source must " "iterate in a stable order" ) - shard = tensor.detach().contiguous() - src_mesh, src_placements = resolve_layout(self._src_mesh, meta.placements) + validate_local_buffer(m2n, meta, buffer) + src_mesh, src_placements = resolve_layout( + meta.source_layout.mesh, meta.source_layout.placements + ) dst_mesh, dst_placements = resolve_layout( self._dst_mesh, self._dst_placements[index] ) + raw_pointer_args = {} + if isinstance(buffer, int): + raw_pointer_args = { + "src_local_shape": expected_local_shape(meta), + "src_dtype": meta.dtype, + } m2n.reshard( - shard, + buffer, None, - comm, + self.group, stream, src_mesh=to_mesh(m2n, src_mesh), src_placements=to_placements(m2n, src_placements), dst_mesh=to_mesh(m2n, dst_mesh), dst_placements=to_placements(m2n, dst_placements), handle=self._handle, + **raw_pointer_args, + ) + # Let allocator-aware buffers recycle storage in stream order. + # Other object owners must stay alive until the final synchronize; + # raw-pointer lifetime remains the source's responsibility. + record_stream = getattr(buffer, "record_stream", None) + if callable(record_stream): + record_stream(stream) + elif not isinstance(buffer, int): + inflight_buffers.append(buffer) + sent += 1 + + if sent != len(self._metas): + missing = self._metas[sent].name + raise RuntimeError( + f"weight source yielded {sent} parameters, but metadata() " + f"declared {len(self._metas)} (first missing: '{missing}')" ) - # A custom source may yield temporary storage allocated on another - # stream. M2N retains only its device pointer for asynchronous work, - # so tell the allocator not to reuse that storage until this stream - # has finished with it. For storage allocated on this stream, - # record_stream is effectively a no-op. - shard.record_stream(stream) # As in the other weight-transfer backends, wait before returning # rather than relying on same-stream ordering. The caller may then @@ -402,14 +599,30 @@ class M2NTrainerWeightTransferEngine(TrainerWeightTransferEngine[M2NTrainerInitI stream.synchronize() def shutdown(self) -> None: - """Destroy the m2n handle, join the helper thread, and drop the group.""" - if self._handle is not None: - torch.accelerator.synchronize() - self._handle.destroy() - self._handle = None - if self._executor is not None: - self._executor.shutdown() - self._executor = None - if self.group is not None: - self.group.destroy() - self.group = None + """Finish pending GPU work and release M2N communication state.""" + handle, self._handle = self._handle, None + group, self.group = self.group, None + self._metadata_group = None + self._dst_placements = [] + self._dst_mesh = None + self._m2n = None + + first_error: BaseException | None = None + + def attempt(action: Any, message: str) -> None: + nonlocal first_error + try: + action() + except BaseException as exc: + if first_error is None: + first_error = exc + logger.exception(message) + + if handle is not None or group is not None: + attempt(torch.accelerator.synchronize, "failed to synchronize nccl_m2n") + if handle is not None: + attempt(handle.destroy, "failed to destroy nccl_m2n handle") + if group is not None: + attempt(group.abort, "failed to abort nccl_m2n communicator") + if first_error is not None: + raise first_error diff --git a/vllm/distributed/weight_transfer/nccl_common.py b/vllm/distributed/weight_transfer/nccl_common.py index 1cd026fb07..b0ac285cf1 100644 --- a/vllm/distributed/weight_transfer/nccl_common.py +++ b/vllm/distributed/weight_transfer/nccl_common.py @@ -8,6 +8,7 @@ The dense (`NCCLWeightTransferEngine`) and sparse sparse engine does not have to subclass the dense one. """ +import socket from dataclasses import asdict, dataclass, field from typing import TYPE_CHECKING, Protocol @@ -17,6 +18,7 @@ import torch if TYPE_CHECKING: from vllm.config.parallel import ParallelConfig from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator + from vllm.distributed.utils import StatelessProcessGroup from vllm.distributed.device_communicators.pynccl_wrapper import ( NCCL_UNIQUE_ID_BYTES, @@ -136,6 +138,26 @@ class NCCLRendezvous(Protocol): world_size: int +def stateless_init_metadata_group( + master_address: str, + master_port: int, + rank: int, + world_size: int, + *, + listen_socket: socket.socket | None = None, +) -> "StatelessProcessGroup": + """Create only the TCPStore group used for metadata agreement.""" + from vllm.distributed.utils import StatelessProcessGroup + + return StatelessProcessGroup.create( + host=master_address, + port=master_port, + rank=rank, + world_size=world_size, + listen_socket=listen_socket, + ) + + def stateless_init_process_group( master_address: str, master_port: int, @@ -181,6 +203,29 @@ def uid_init_process_group( ) +def _worker_group_rank(rank_offset: int, parallel_config: "ParallelConfig") -> int: + """Return this worker's rank in a trainer-to-worker process group.""" + dp_rank = parallel_config.data_parallel_index + world_size_per_dp = parallel_config.world_size # TP * PP + rank_within_dp = parallel_config.rank + return rank_offset + dp_rank * world_size_per_dp + rank_within_dp + + +def worker_init_metadata_group( + init_info: NCCLWeightTransferInitInfo, + parallel_config: "ParallelConfig", +) -> "StatelessProcessGroup": + """Join the trainer-to-worker TCP metadata group without constructing NCCL.""" + assert init_info.master_address is not None + assert init_info.master_port is not None + return stateless_init_metadata_group( + init_info.master_address, + init_info.master_port, + _worker_group_rank(init_info.rank_offset, parallel_config), + init_info.world_size, + ) + + def worker_init_process_group( init_info: NCCLWeightTransferInitInfo, parallel_config: "ParallelConfig", @@ -190,15 +235,7 @@ def worker_init_process_group( Computes a unique rank for this worker across all data-parallel groups and joins the trainer via whichever rendezvous mode `init_info` carries. """ - # Calculate the global rank in the trainer-worker process group. - # Must account for data parallel to get unique ranks across all workers. - dp_rank = parallel_config.data_parallel_index - world_size_per_dp = parallel_config.world_size # TP * PP - rank_within_dp = parallel_config.rank - - # Unique rank across all DP groups - worker_rank = dp_rank * world_size_per_dp + rank_within_dp - rank = worker_rank + init_info.rank_offset + rank = _worker_group_rank(init_info.rank_offset, parallel_config) device = torch.accelerator.current_device_index() unique_id_bytes = init_info.nccl_unique_id_bytes From b3a7e2ff6e7d182cd9c230cdb04dd66043b9f413 Mon Sep 17 00:00:00 2001 From: Shuolei Wang <46160365+ShuoleiWang@users.noreply.github.com> Date: Sat, 17 Oct 2026 10:08:07 +0800 Subject: [PATCH 35/49] [RL][Frontend] Return weight-version spans on OpenAI completions and chat (#61347) Signed-off-by: Shuolei Wang Co-authored-by: aoshen02 --- docs/training/async_rl.md | 4 + .../chat_completion/test_serving_chat.py | 9 +++ .../completion/test_completion_error.py | 73 ++++++++++++++++++- .../openai/chat_completion/protocol.py | 14 ++++ .../openai/chat_completion/serving.py | 18 ++++- .../entrypoints/openai/completion/protocol.py | 14 ++++ vllm/entrypoints/openai/completion/serving.py | 11 +++ 7 files changed, 139 insertions(+), 4 deletions(-) diff --git a/docs/training/async_rl.md b/docs/training/async_rl.md index b5ae067e3c..6442595999 100644 --- a/docs/training/async_rl.md +++ b/docs/training/async_rl.md @@ -76,6 +76,10 @@ for span in completion.weight_versions: The offsets cover the whole output, so with streamed delta outputs, collect the tokens before slicing. `/inference/v1/generate` returns the same list on each choice; see [Tokens In <> Tokens Out API](../serving/online_serving/token_in_token_out.md#weight-versions). +`/v1/completions` and `/v1/chat/completions` return it on each choice when the request sets `return_weight_versions: true`. The offsets count the choice's output tokens, which `return_token_ids` returns as `token_ids`, and a streamed choice gets the list on its last chunk. + +An aborted request returns spans over its partial output. If you resubmit that output as part of a new prompt, the new request's spans start again at 0, so shift them by the number of tokens already generated before joining the two lists. + ## Example The [async RLHF example](../../examples/rl/rlhf_async_new_apis.py) demonstrates this pattern with `vllm.AsyncLLMEngine`, NCCL weight transfer, and mid-flight pause/resume with validation. diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 747b84a938..85404af06c 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -1188,6 +1188,15 @@ async def test_streaming_reasoning_usage_lifecycle_edges( await iterator.aclose() error_chunks = [c for c in chunks if isinstance(c, dict) and "error" in c] + if name == "empty-finish": + # A choice that finishes without output still gets its last chunk, unparsed. + assert any( + choice.get("finish_reason") == "stop" + for item in chunks + if isinstance(item, dict) + for choice in item.get("choices", []) + ) + parsers[0].parse_delta.assert_not_called() call_count = sum(parser.count_reasoning_tokens.call_count for parser in parsers) if expected_tokens is not None: assert outcome == "exhausted" and chunks[-1] == "[DONE]" and not error_chunks diff --git a/tests/entrypoints/openai/completion/test_completion_error.py b/tests/entrypoints/openai/completion/test_completion_error.py index 17934026ef..0217d8a733 100644 --- a/tests/entrypoints/openai/completion/test_completion_error.py +++ b/tests/entrypoints/openai/completion/test_completion_error.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from dataclasses import dataclass, field from typing import Any from unittest.mock import AsyncMock, MagicMock @@ -17,7 +18,7 @@ from vllm.entrypoints.scale_out.render.serving import ServingRender from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.exceptions import GenerationError, VLLMValidationError from vllm.logprobs import Logprob -from vllm.outputs import CompletionOutput, RequestOutput +from vllm.outputs import CompletionOutput, RequestOutput, WeightVersionSpan from vllm.renderers.deepseek_v4 import DeepseekV4Renderer from vllm.renderers.hf import HfRenderer from vllm.renderers.online_renderer import OnlineRenderer @@ -310,6 +311,76 @@ def test_completion_metrics_carries_both_timing_and_spec_decode(): assert response.metrics.speculative_decoding.num_spec_steps == 2 +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +@pytest.mark.parametrize("opt_in", [False, True]) +async def test_completion_weight_versions(stream: bool, opt_in: bool): + """Each choice gets its spans on its final output, even with no tokens.""" + expected = [{"version": "v1", "start": 0, "end": 2}] + spans = [WeightVersionSpan(**span) for span in expected] + + def output(index, token_ids, finish_reason=None, weight_versions=None): + return CompletionOutput( + index=index, + text="x" * len(token_ids), + token_ids=token_ids, + cumulative_logprob=None, + logprobs=None, + finish_reason=finish_reason, + weight_versions=weight_versions, + ) + + # Choice 1 is aborted before its first token. + final = [output(1, [], "abort", [])] + if stream: + steps = [[output(0, [100])], [output(0, [101], "length", spans), *final]] + else: + steps = [[output(0, [100, 101], "length", spans), *final]] + + async def mock_generate(*args, **kwargs): + for i, outputs in enumerate(steps): + yield RequestOutput( + request_id="test-id", + prompt="Test prompt", + prompt_token_ids=[1, 2, 3], + prompt_logprobs=None, + outputs=outputs, + finished=i == len(steps) - 1, + ) + + mock_engine = MagicMock(spec=AsyncLLM) + mock_engine.errored = False + mock_engine.model_config = MockModelConfig() + mock_engine.input_processor = MagicMock() + mock_engine.renderer = _build_renderer(mock_engine.model_config) + mock_engine.generate = MagicMock(side_effect=mock_generate) + request = CompletionRequest( + model=MODEL_NAME, + prompt="Test prompt", + max_tokens=10, + n=2, + stream=stream, + return_weight_versions=opt_in, + ) + response = await _build_serving_completion(mock_engine).create_completion(request) + + if stream: + finals = {} + async for line in response: + payload = line.removeprefix("data: ").strip() + if payload == "[DONE]": + continue + for choice in json.loads(payload)["choices"]: + if choice.get("finish_reason"): + finals[choice["index"]] = choice + else: + assert choice.get("weight_versions") is None + got = [finals[i].get("weight_versions") for i in (0, 1)] + else: + got = [choice["weight_versions"] for choice in response.model_dump()["choices"]] + assert got == ([expected, []] if opt_in else [None, None]) + + def _build_online_renderer_for_completion( model_config: MockModelConfig, renderer: Any | None = None, diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index 79dd9605e6..071f73f043 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -40,6 +40,7 @@ from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.logprobs import Logprob +from vllm.outputs import WeightVersionSpan from vllm.renderers import ChatParams, TokenizeParams, merge_kwargs from vllm.renderers.chat_utils import ( ChatCompletionMessageParam, @@ -129,6 +130,8 @@ class ChatCompletionResponseChoice(OpenAIBaseModel): # ``None`` if (a) the request was aborted before any forward pass, # or (b) ``enable_return_routed_experts`` is off server-side. routed_experts: str | None = None + # Contiguous weight version spans over output-token indexes. + weight_versions: list[WeightVersionSpan] | None = None class ChatCompletionResponse(OpenAIBaseModel): @@ -168,6 +171,8 @@ class ChatCompletionResponseStreamChoice(OpenAIBaseModel): stop_reason: int | str | None = None # not part of the OpenAI spec but for tracing the tokens token_ids: list[int] | None = None + # Set on the final chunk only; covers the whole output. + weight_versions: list[WeightVersionSpan] | None = None class ChatCompletionStreamResponse(OpenAIBaseModel): @@ -430,6 +435,15 @@ class ChatCompletionRequest(OpenAIBaseModel): "need to map generated text back to input tokens." ), ) + return_weight_versions: bool = Field( + default=False, + description=( + "If true, each choice includes weight_versions: the weight-version " + "label that sampled each half-open range of its output tokens. " + "When streaming, only the last chunk of each choice includes it. " + "Beam search does not return it." + ), + ) routed_experts_prompt_start: int = Field( default=0, ge=0, diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 9ddc20ed30..7e1782e241 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -651,17 +651,21 @@ class OpenAIServingChat(GenerateBaseServing): delta_text = output.text - if ( + nothing_generated = ( not delta_text and not output.token_ids and not previous_num_tokens[i] - ): + ) + if nothing_generated and output.finish_reason is None: # Chunked prefill case, don't return empty chunks continue delta_message: DeltaMessage | None - if parser is not None: + if nothing_generated: + # Finished without output: send the final chunk unparsed. + delta_message = DeltaMessage() + elif parser is not None: delta_message = parser.parse_delta( delta_text=delta_text, delta_token_ids=as_list(output.token_ids), @@ -773,6 +777,11 @@ class OpenAIServingChat(GenerateBaseServing): token_ids=( as_list(output.token_ids) if include_token_ids else None ), + weight_versions=( + output.weight_versions + if request.return_weight_versions + else None + ), ) finish_reason_sent[i] = True @@ -1097,6 +1106,9 @@ class OpenAIServingChat(GenerateBaseServing): else None ), routed_experts=routed_experts_b64, + weight_versions=( + output.weight_versions if request.return_weight_versions else None + ), ) choice_data = maybe_filter_parallel_tool_calls(choice_data, request) diff --git a/vllm/entrypoints/openai/completion/protocol.py b/vllm/entrypoints/openai/completion/protocol.py index 4da67e607a..9e8ffb496b 100644 --- a/vllm/entrypoints/openai/completion/protocol.py +++ b/vllm/entrypoints/openai/completion/protocol.py @@ -26,6 +26,7 @@ from vllm.entrypoints.serve.engine.protocol import OpenAIBaseModel, UsageInfo from vllm.exceptions import VLLMValidationError from vllm.logger import init_logger from vllm.logprobs import Logprob +from vllm.outputs import WeightVersionSpan from vllm.renderers import TokenizeParams from vllm.sampling_params import ( BeamSearchParams, @@ -178,6 +179,15 @@ class CompletionRequest(OpenAIBaseModel): "need to map generated text back to input tokens." ), ) + return_weight_versions: bool = Field( + default=False, + description=( + "If true, each choice includes weight_versions: the weight-version " + "label that sampled each half-open range of its output tokens. " + "When streaming, only the last chunk of each choice includes it. " + "Beam search does not return it." + ), + ) routed_experts_prompt_start: int = Field( default=0, ge=0, @@ -677,6 +687,8 @@ class CompletionResponseChoice(OpenAIBaseModel): # ``None`` if (a) the request was aborted before any forward pass, # or (b) ``enable_return_routed_experts`` is off server-side. routed_experts: str | None = None + # Contiguous weight version spans over output-token indexes. + weight_versions: list[WeightVersionSpan] | None = None class CompletionResponse(OpenAIBaseModel): @@ -716,6 +728,8 @@ class CompletionResponseStreamChoice(OpenAIBaseModel): # prompt tokens is put into choice to align with CompletionResponseChoice prompt_token_ids: list[int] | None = None token_ids: list[int] | None = None + # Set on the final chunk only; covers the whole output. + weight_versions: list[WeightVersionSpan] | None = None class CompletionStreamResponse(OpenAIBaseModel): diff --git a/vllm/entrypoints/openai/completion/serving.py b/vllm/entrypoints/openai/completion/serving.py index 0306909237..944da27530 100644 --- a/vllm/entrypoints/openai/completion/serving.py +++ b/vllm/entrypoints/openai/completion/serving.py @@ -376,6 +376,7 @@ class OpenAIServingCompletion(GenerateBaseServing): not delta_text and not delta_token_ids and not previous_num_tokens[i] + and finish_reason is None ): # Chunked prefill case, don't return empty chunks continue @@ -421,6 +422,11 @@ class OpenAIServingCompletion(GenerateBaseServing): if request.return_token_ids else None ), + weight_versions=( + output.weight_versions + if request.return_weight_versions + else None + ), ) ], ) @@ -594,6 +600,11 @@ class OpenAIServingCompletion(GenerateBaseServing): as_list(output.token_ids) if request.return_token_ids else None ), routed_experts=routed_experts_b64, + weight_versions=( + output.weight_versions + if request.return_weight_versions + else None + ), ) choices.append(choice_data) From 71c20f73f2f1764576fed799c152319c9fd1c820 Mon Sep 17 00:00:00 2001 From: Yang Xu <38851819+YangXu1990uiuc@users.noreply.github.com> Date: Fri, 16 Oct 2026 19:08:28 -0700 Subject: [PATCH 36/49] [FlashInfer] Migrate autotune warmup to flashinfer.autotune_v2 managed cache (#55671) Signed-off-by: Yang Xu Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Signed-off-by: YangXu Co-authored-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Kimi Code Co-authored-by: Codex Co-authored-by: Misha Goin Co-authored-by: vllm-agent --- docs/usage/security.md | 4 +- .../test_flashinfer_autotune_warmup.py | 271 +++++++++++++++++- .../warmup/flashinfer_autotune_cache.py | 50 ++++ .../warmup/flashinfer_sparse_mla_warmup.py | 62 +++- vllm/model_executor/warmup/kernel_warmup.py | 72 +++-- vllm/utils/flashinfer.py | 14 + 6 files changed, 436 insertions(+), 37 deletions(-) diff --git a/docs/usage/security.md b/docs/usage/security.md index f43ad0a6e3..4a704578fa 100644 --- a/docs/usage/security.md +++ b/docs/usage/security.md @@ -423,9 +423,9 @@ Most cache paths default to subdirectories under a single root. Changing `VLLM_C | Environment Variable | Default | Description | | --- | --- | --- | -| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base cache directory. Respects `XDG_CACHE_HOME` if set. All paths below inherit from this unless explicitly overridden. | +| `VLLM_CACHE_ROOT` | `~/.cache/vllm` | Base directory for vLLM-owned caches. Respects `XDG_CACHE_HOME` if set. | | *(torch.compile)* | `$VLLM_CACHE_ROOT/torch_compile_cache/` | Compilation cache for AOT-compiled models, Inductor graphs, and Triton kernels. Controlled by `VLLM_DISABLE_COMPILE_CACHE` (set to `1` to disable). | -| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | `$VLLM_CACHE_ROOT/flashinfer_autotune_cache////` | FlashInfer autotune config cache. | +| `VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR` | Managed: `~/.cache/flashinfer/autotune/`; legacy: `$VLLM_CACHE_ROOT/flashinfer_autotune_cache////` | Overrides cache placement for both FlashInfer autotune cache APIs. | | `VLLM_ASSETS_CACHE` | `$VLLM_CACHE_ROOT/assets/` | Downloaded assets (e.g., tokenizer files). | | `VLLM_XLA_CACHE_PATH` | `$VLLM_CACHE_ROOT/xla_cache/` | XLA/TPU compilation cache. | | `VLLM_MEDIA_CACHE` | *(disabled)* | Optional cache for downloaded media (images, video, audio). Not enabled unless explicitly set. | diff --git a/tests/model_executor/test_flashinfer_autotune_warmup.py b/tests/model_executor/test_flashinfer_autotune_warmup.py index 59046a5f3e..f8c86cd1b5 100644 --- a/tests/model_executor/test_flashinfer_autotune_warmup.py +++ b/tests/model_executor/test_flashinfer_autotune_warmup.py @@ -4,7 +4,7 @@ import json import sys from collections import defaultdict -from contextlib import nullcontext +from contextlib import contextmanager, nullcontext from inspect import signature from pathlib import Path from types import ModuleType, SimpleNamespace @@ -12,8 +12,10 @@ from typing import Any from unittest.mock import Mock, call, patch import pytest +import torch -from vllm.model_executor.warmup import kernel_warmup as warmup +import vllm.model_executor.warmup.kernel_warmup as warmup +import vllm.utils.flashinfer as fi_utils from vllm.model_executor.warmup.kernel_warmup import ( _flashinfer_autotune_token_counts, _run_flashinfer_autotune_dummy_runs, @@ -45,7 +47,12 @@ def _make_runner(modules, *, max_tokens: int = 8192): scheduler_config=SimpleNamespace(max_num_batched_tokens=max_tokens), vllm_config=SimpleNamespace( attention_config=SimpleNamespace(hisparse_config=None), - parallel_config=SimpleNamespace(data_parallel_rank=0), + parallel_config=SimpleNamespace( + enable_elastic_ep=False, + nnodes=1, + pipeline_parallel_size=1, + data_parallel_rank=0, + ), ), get_model=Mock( return_value=SimpleNamespace(modules=Mock(return_value=modules)) @@ -133,6 +140,258 @@ def test_flashinfer_autotune_uses_token_buckets_for_each_dummy_run(skip_attn): ] +@pytest.mark.parametrize( + "has_v2,elastic_ep,node_count", + [(False, False, 1), (True, False, 1), (True, True, 1), (True, False, 2)], + ids=["v1", "v2", "elastic-v1", "multi-node-v1"], +) +@pytest.mark.parametrize( + "hisparse,world_size,rank,skip_ops,failure", + [ + (False, 1, 0, [], None), + (False, 2, 0, [], None), + (False, 2, 1, [], None), + (True, 2, 0, ["fp4_gemm", "bmm_fp8"], None), + (False, 2, 1, ["fp4_gemm"], "dummy"), + (False, 2, 1, [], "replayssm"), + (True, 2, 0, [], "kimi"), + ], +) +def test_flashinfer_autotune_lifecycle( + monkeypatch, + tmp_path, + has_v2, + elastic_ep, + node_count, + hisparse, + world_size, + rank, + skip_ops, + failure, +): + use_v2 = has_v2 and not elastic_ep and node_count == 1 and world_size == 1 + runner = _make_runner([]) + runner.vllm_config.parallel_config.enable_elastic_ep = elastic_ep + runner.vllm_config.attention_config.hisparse_config = object() if hisparse else None + model = runner.get_model() + cache_path = tmp_path / "autotune.json" + rank_cache_path = cache_path.with_name(f"autotune_dp0_rank{rank}.json") + rank_cache_path.write_text(json.dumps({"_metadata": {}, "op": 1})) + events = Mock() + events.tuner.load_configs.return_value = True + events.tuner._dirty = True + world = SimpleNamespace( + rank_in_group=rank, + world_size=world_size, + cpu_group=object(), + broadcast_object=events.broadcast, + barrier=events.barrier, + ) + + @contextmanager + def autotune_context(**kwargs): + assert torch.is_inference_mode_enabled() + events.enter() + try: + yield + finally: + events.exit() + + events.autotune.side_effect = autotune_context + if failure is not None: + getattr(events, failure).side_effect = RuntimeError("warmup failed") + monkeypatch.setitem( + sys.modules, + "flashinfer.autotuner", + SimpleNamespace( + AutoTuner=SimpleNamespace(get=lambda: events.tuner), + set_autotune_process_group=events.set_group, + ), + ) + monkeypatch.setitem( + sys.modules, + "flashinfer", + SimpleNamespace(autotune_v2=events.autotune, autotune_v2_reload=events.reload), + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_world_group", lambda: world + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_node_count", lambda: node_count + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_pp_group", + lambda: SimpleNamespace(world_size=1), + ) + monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: has_v2) + monkeypatch.setattr(fi_utils, "autotune", events.autotune) + monkeypatch.setattr(warmup.envs, "VLLM_FLASHINFER_AUTOTUNE_SKIP_OPS", skip_ops) + resolve_file = Mock(return_value=cache_path) + resolve_root = Mock(return_value=str(tmp_path)) + monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_file", resolve_file) + monkeypatch.setattr(warmup, "resolve_flashinfer_autotune_v2_root", resolve_root) + + def gather(out, value, *, group): + out[:] = [value] * world_size + + monkeypatch.setattr(torch.distributed, "all_gather_object", gather) + monkeypatch.setattr(warmup, "_run_flashinfer_autotune_dummy_runs", events.dummy) + monkeypatch.setattr(warmup, "replayssm_autotune_warmup", events.replayssm) + monkeypatch.setattr(warmup, "_autotune_kimi_k3_kda_qkvg", events.kimi) + monkeypatch.setattr( + warmup, "autotune_hisparse_flashinfer_attention", events.hisparse + ) + + if failure is None: + warmup.flashinfer_autotune(runner) + else: + with pytest.raises(RuntimeError, match="warmup failed"): + warmup.flashinfer_autotune(runner) + + kwargs = {"skip_ops": set(skip_ops)} if skip_ops else {} + expected = [] + if use_v2: + resolve_file.assert_not_called() + resolve_root.assert_called_once_with() + expected.append(call.autotune(mode="tune", cache_root=str(tmp_path), **kwargs)) + else: + resolve_root.assert_not_called() + resolve_file.assert_called_once_with(runner) + expected.extend( + [ + call.tuner.load_configs(str(rank_cache_path)), + call.autotune(tune_mode=True, **kwargs), + ] + ) + expected.extend( + [call.set_group(world.cpu_group if world_size > 1 else None), call.enter()] + ) + if hisparse: + expected.append(call.hisparse(runner)) + for name, arg, kwargs in [ + ("dummy", runner, {"skip_attn": hisparse}), + ("replayssm", runner, {}), + ("kimi", model, {}), + ]: + expected.append(getattr(call, name)(arg, **kwargs)) + if failure == name: + break + expected.append(call.exit()) + expected.append(call.set_group(None)) + if failure is None: + if world_size > 1: + expected.append(call.barrier()) + if use_v2 and world_size > 1: + expected.extend([call.reload(), call.barrier()]) + elif not use_v2: + expected.append(call.tuner.save_configs(str(rank_cache_path))) + assert events.mock_calls == expected + + +@pytest.mark.parametrize( + "node_count,elastic_ep,pp_size", + [(1, False, 1), (2, False, 1), (1, True, 1), (1, False, 2)], +) +@pytest.mark.parametrize("world_size,rank", [(1, 0), (2, 0), (2, 1)]) +@pytest.mark.parametrize("model_runner_v2,max_reqs", [(False, 2), (True, 1), (True, 2)]) +def test_sparse_autotune_respects_cache_selection( + monkeypatch, + tmp_path, + node_count, + elastic_ep, + pp_size, + world_size, + rank, + model_runner_v2, + max_reqs, +): + import vllm.model_executor.warmup.flashinfer_sparse_mla_warmup as sparse + + runner = _make_runner([]) + runner.max_num_reqs = max_reqs + runner.vllm_config.use_v2_model_runner = model_runner_v2 + runner.vllm_config.kernel_config = SimpleNamespace(enable_flashinfer_autotune=True) + runner.vllm_config.parallel_config.enable_elastic_ep = elastic_ep + runner.vllm_config.parallel_config.pipeline_parallel_size = pp_size + worker = SimpleNamespace( + model_runner=runner, + vllm_config=runner.vllm_config, + execute_model=Mock(), + sample_tokens=Mock(), + ) + world = SimpleNamespace( + rank_in_group=rank, + world_size=world_size, + broadcast_object=Mock(return_value=b"cached"), + barrier=Mock(), + ) + managed = Mock(return_value=nullcontext()) + legacy = Mock(return_value=nullcontext()) + reload = Mock() + tuner = Mock() + monkeypatch.setitem( + sys.modules, + "flashinfer", + SimpleNamespace(autotune_v2=managed, autotune_v2_reload=reload), + ) + monkeypatch.setitem( + sys.modules, + "flashinfer.autotuner", + SimpleNamespace(AutoTuner=SimpleNamespace(get=lambda: tuner)), + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_world_group", lambda: world + ) + monkeypatch.setattr( + "vllm.distributed.parallel_state.get_node_count", lambda: node_count + ) + monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: True) + monkeypatch.setattr(sparse, "has_flashinfer", lambda: True) + monkeypatch.setattr( + sparse.current_platform, "is_device_capability_family", lambda _: True + ) + monkeypatch.setattr( + sparse, "_flashinfer_sparse_mla_decode_label", lambda *_: "DSv3.2" + ) + monkeypatch.setattr(sparse, "resolve_flashinfer_autotune_v2_root", lambda: tmp_path) + cache_path = tmp_path / "autotune.json" + cache_path.write_bytes(b"cached") + monkeypatch.setattr( + sparse, "resolve_flashinfer_autotune_file", lambda _: cache_path + ) + monkeypatch.setattr(sparse, "flashinfer_autotune", legacy) + + def mixed_warmup(*args, mixed_step_context=None, **kwargs): + context = ( + mixed_step_context if mixed_step_context is not None else nullcontext() + ) + with context: + assert torch.is_inference_mode_enabled() + return True + + mixed = Mock(side_effect=mixed_warmup) + monkeypatch.setattr(sparse, "run_mixed_prefill_decode_warmup", mixed) + assert sparse._run_flashinfer_sparse_mla_decode_autotune(worker, 16, frozenset()) + + if node_count == 1 and not elastic_ep and pp_size == 1 and world_size == 1: + managed.assert_called_once_with(mode="tune", cache_root=tmp_path) + legacy.assert_not_called() + world.broadcast_object.assert_not_called() + tuner.load_configs.assert_not_called() + reload.assert_not_called() + world.barrier.assert_not_called() + else: + managed.assert_not_called() + reload.assert_not_called() + assert legacy.call_count == (1 if rank == 0 else 0) + world.broadcast_object.assert_called_once_with( + b"cached" if rank == 0 else None, src=0 + ) + tuner.load_configs.assert_called_once_with(str(cache_path)) + assert mixed.call_count == (1 if model_runner_v2 and max_reqs >= 2 else 0) + assert runner._dummy_run.call_count == (0 if mixed.called else 1) + + def test_flashinfer_autotune_buckets_cover_drafter_tokens(): """The drafter's M can exceed the pass size; buckets must include it.""" runner = _make_runner([]) @@ -266,7 +525,9 @@ class _AutotuneRun: for rank in range(self.pp * self.tp): self.rank = rank self.tuners[rank] = _AutotuneTuner(self) - flashinfer_autotune(_make_runner([])) + runner = _make_runner([]) + runner.vllm_config.parallel_config.pipeline_parallel_size = self.pp + flashinfer_autotune(runner) return self def assert_collectives_match(self): @@ -283,6 +544,8 @@ def autotune_run(monkeypatch, tmp_path): def make_run(*, pp=2, tp=4): run = _AutotuneRun(pp, tp, tmp_path) + monkeypatch.setattr(fi_utils, "has_flashinfer_autotune_v2", lambda: pp > 1) + monkeypatch.setattr(parallel_state, "get_node_count", lambda: 1) autotuner = ModuleType("flashinfer.autotuner") monkeypatch.setattr( autotuner, diff --git a/vllm/model_executor/warmup/flashinfer_autotune_cache.py b/vllm/model_executor/warmup/flashinfer_autotune_cache.py index ca2c4d1421..3c41743033 100644 --- a/vllm/model_executor/warmup/flashinfer_autotune_cache.py +++ b/vllm/model_executor/warmup/flashinfer_autotune_cache.py @@ -11,11 +11,14 @@ from typing import TYPE_CHECKING import vllm.envs as envs from vllm.config import CacheConfig, VllmConfig, replace +from vllm.logger import init_logger if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator from vllm.v1.worker.gpu_model_runner import GPUModelRunner +logger = init_logger(__name__) + def _normalize_cache_config_for_hash(cache_config: CacheConfig) -> CacheConfig: """Normalize cache config for hashing""" @@ -55,6 +58,53 @@ def resolve_flashinfer_autotune_file(runner: "GPUModelRunner") -> Path: return output_dir / "autotune_configs.json" +def use_flashinfer_autotune_v2(runner: "GPUModelRunner") -> bool: + from vllm.distributed.parallel_state import get_node_count, get_world_group + from vllm.utils.flashinfer import has_flashinfer_autotune_v2 + + # Elastic EP transfers tuning state to joining workers with v1 + # save_configs/load_configs, which do not export managed-store winners. + if runner.vllm_config.parallel_config.enable_elastic_ep: + return False + # PP stages profile independently. Keep their group-scoped v1 files until + # managed-cache reads and publishes can be isolated across tuning groups. + if runner.vllm_config.parallel_config.pipeline_parallel_size > 1: + return False + if not has_flashinfer_autotune_v2(): + return False + # Use the initialized topology: Ray may span nodes with config.nnodes=1. + # A cache hit skips profiling collectives, so node-local stores with + # different contents cannot safely participate in synchronized tuning. + if get_node_count() != 1: + logger.info_once( + "Using legacy FlashInfer autotune cache synchronization for " + "multi-node deployments; managed-cache synchronization is " + "currently supported only within one node." + ) + return False + # Lift this gate after pinning a release with the fix below and validating + # distributed managed-cache tuning: https://github.com/flashinfer-ai/flashinfer/pull/6185 + if get_world_group().world_size > 1: + logger.info_once( + "Using legacy FlashInfer autotune cache synchronization for " + "multi-rank deployments until FlashInfer coordinates cache hits." + ) + return False + return True + + +def resolve_flashinfer_autotune_v2_root() -> Path | None: + """Placement root for FlashInfer's managed autotune store. + + Unlike :func:`resolve_flashinfer_autotune_file`, no vLLM-side identity + (config hash, flashinfer version, arch) is folded into the path: the + managed store keys entries per op and hashes the software/hardware + environment itself. ``None`` selects FlashInfer's default root. + """ + override_dir = envs.VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR + return Path(override_dir).expanduser() if override_dir else None + + def sync_flashinfer_autotune_cache( runner: "GPUModelRunner", group: "GroupCoordinator", diff --git a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py index 5c01f4ca2e..fee6a84762 100644 --- a/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py +++ b/vllm/model_executor/warmup/flashinfer_sparse_mla_warmup.py @@ -9,6 +9,8 @@ import torch from vllm.logger import init_logger from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, + resolve_flashinfer_autotune_v2_root, + use_flashinfer_autotune_v2, write_flashinfer_autotune_cache, ) from vllm.platforms import current_platform @@ -145,7 +147,6 @@ def _run_flashinfer_sparse_mla_decode_autotune( world = get_world_group() is_leader = world.rank_in_group == 0 - cache_path = resolve_flashinfer_autotune_file(runner) dummy_run_kwargs = dict( num_tokens=num_tokens, @@ -155,6 +156,13 @@ def _run_flashinfer_sparse_mla_decode_autotune( create_mixed_batch=True, ) + if use_flashinfer_autotune_v2(runner): + return _run_sparse_mla_decode_autotune_v2( + worker, num_tokens, log_label, dummy_run_kwargs + ) + + cache_path = resolve_flashinfer_autotune_file(runner) + if is_leader: logger.info( "Autotuning FlashInfer SM120 sparse MLA %s decode with cache: %s", @@ -223,6 +231,58 @@ def _run_flashinfer_sparse_mla_decode_autotune( return True +def _run_sparse_mla_decode_autotune_v2( + worker: "Worker", + num_tokens: int, + log_label: str, + dummy_run_kwargs: dict, +) -> bool: + """Run every rank's sparse warmup in the managed autotune context. + + FlashInfer owns persistence; some sparse kernels use a separate calibration + cache instead of managed winner entries. + """ + from flashinfer import autotune_v2, autotune_v2_reload + + from vllm.distributed.parallel_state import get_world_group + + runner = worker.model_runner + world = get_world_group() + cache_root = resolve_flashinfer_autotune_v2_root() + if world.rank_in_group == 0: + logger.info( + "Autotuning FlashInfer SM120 sparse MLA %s decode with the managed " + "autotune context (root=%s)", + log_label, + cache_root if cache_root is not None else "flashinfer default", + ) + + with torch.inference_mode(): + if _uses_v2_model_runner(runner) and runner.max_num_reqs >= 2: + v2_runner = cast("V2GPUModelRunner", runner) + warmup_executed = run_mixed_prefill_decode_warmup( + v2_runner, + worker.execute_model, + worker.sample_tokens, + num_tokens, + mixed_step_context=autotune_v2(mode="tune", cache_root=cache_root), + req_id_prefix="_sparse_mla_v2_warmup", + ) + else: + warmup_executed = True + with autotune_v2(mode="tune", cache_root=cache_root): + runner._dummy_run(**dummy_run_kwargs) + + if not warmup_executed: + return False + + if world.world_size > 1: + world.barrier() + autotune_v2_reload() + world.barrier() + return True + + def _flashinfer_sparse_mla_decode_autotune( worker: "Worker", num_tokens: int, diff --git a/vllm/model_executor/warmup/kernel_warmup.py b/vllm/model_executor/warmup/kernel_warmup.py index b7bf50865f..ba97d2dfa3 100644 --- a/vllm/model_executor/warmup/kernel_warmup.py +++ b/vllm/model_executor/warmup/kernel_warmup.py @@ -23,6 +23,8 @@ from vllm.model_executor.warmup.cutedsl_warmup import cutedsl_warmup from vllm.model_executor.warmup.deep_gemm_warmup import deep_gemm_warmup from vllm.model_executor.warmup.flashinfer_autotune_cache import ( resolve_flashinfer_autotune_file, + resolve_flashinfer_autotune_v2_root, + use_flashinfer_autotune_v2, ) from vllm.model_executor.warmup.flashinfer_sparse_mla_warmup import ( autotune_hisparse_flashinfer_attention, @@ -419,7 +421,7 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> otherwise the world group tunes together. Per-tactic timings are averaged over the tuning group so all its ranks select the same tactic. - Results are persisted per rank: FlashInfer keys MoE entries by tp/ep rank + Legacy results are persisted per rank: FlashInfer keys MoE entries by tp/ep rank (``MoERunner.get_cache_key_extras``), so one rank's file only hits on that rank. A rank with a cache hit skips the per-tactic reduce the others block in, so ranks keep loaded configs only if every rank in the tuning group @@ -450,38 +452,45 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> ) autotune_kwargs["skip_ops"] = skip_ops - cache_path = resolve_flashinfer_autotune_file(runner) - # The world group only spans DP ranks when vLLM folds DP into it. - dp_rank = runner.vllm_config.parallel_config.data_parallel_rank - cache_path = cache_path.with_name( - f"{cache_path.stem}_dp{dp_rank}_rank{world.rank_in_group}{cache_path.suffix}" - ) - if is_leader: - logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) + use_autotune_v2 = use_flashinfer_autotune_v2(runner) + if use_autotune_v2: + from flashinfer import autotune_v2, autotune_v2_reload - # We skip EPLB here since we don't want to record dummy metrics. - # Randomize inputs to avoid every token pick the same experts, - # which lead to some EP ranks receiving no tokens and skipping their - # MoE kernel entirely, and cause hang due to all-reduce collective - # during synchronized autotuning. - if _all_ranks_have_matching_cache(cache_path, tune_group): - loaded = tuner.load_configs(str(cache_path)) - if tune_group.world_size > 1: - loaded_by_rank: list[bool | None] = [None] * tune_group.world_size - torch.distributed.all_gather_object( - loaded_by_rank, loaded, group=tune_group.cpu_group + cache_root = resolve_flashinfer_autotune_v2_root() + if is_leader: + logger.info_once( + "Using FlashInfer managed autotune cache (root=%s)", + cache_root if cache_root is not None else "flashinfer default", ) - loaded = all(loaded_by_rank) - if not loaded: - tuner.clear_cache() + autotune_context = autotune_v2( + mode="tune", cache_root=cache_root, **autotune_kwargs + ) + else: + cache_path = resolve_flashinfer_autotune_file(runner) + # The world group only spans DP ranks when vLLM folds DP into it. + dp_rank = runner.vllm_config.parallel_config.data_parallel_rank + cache_path = cache_path.with_name( + f"{cache_path.stem}_dp{dp_rank}_rank{world.rank_in_group}{cache_path.suffix}" + ) + if is_leader: + logger.info_once("Using FlashInfer autotune cache file: %s", cache_path) + + if _all_ranks_have_matching_cache(cache_path, tune_group): + loaded = tuner.load_configs(str(cache_path)) + if tune_group.world_size > 1: + loaded_by_rank: list[bool | None] = [None] * tune_group.world_size + torch.distributed.all_gather_object( + loaded_by_rank, loaded, group=tune_group.cpu_group + ) + loaded = all(loaded_by_rank) + if not loaded: + tuner.clear_cache() + autotune_context = fi_utils.autotune(tune_mode=True, **autotune_kwargs) group = tune_group.cpu_group if tune_group.world_size > 1 else None set_autotune_process_group(group) try: - with ( - torch.inference_mode(), - fi_utils.autotune(tune_mode=True, **autotune_kwargs), - ): + with torch.inference_mode(), autotune_context: hisparse_enabled = ( runner.vllm_config.attention_config.hisparse_config is not None ) @@ -498,7 +507,10 @@ def flashinfer_autotune(runner: "GPUModelRunner", *, skip_attn: bool = False) -> if world.world_size > 1: world.barrier() - # Skip the rewrite when nothing was tuned this start (every entry came from - # the file). FlashInfer gates its own autotune(cache=...) save the same way. - if tuner._dirty: + if use_autotune_v2: + if world.world_size > 1: + # All ranks serve the shared store's final state after publishing. + autotune_v2_reload() + world.barrier() + elif tuner._dirty: tuner.save_configs(str(cache_path)) diff --git a/vllm/utils/flashinfer.py b/vllm/utils/flashinfer.py index 69401febea..54fa01ed69 100644 --- a/vllm/utils/flashinfer.py +++ b/vllm/utils/flashinfer.py @@ -358,6 +358,19 @@ autotune = _lazy_import_wrapper( ) +@functools.cache +def has_flashinfer_autotune_v2() -> bool: + """Return ``True`` if FlashInfer ships the managed-cache autotuner + (``flashinfer.autotune_v2``, flashinfer-ai/flashinfer#3861).""" + if not has_flashinfer(): + return False + mod = _get_submodule("flashinfer") + return mod is not None and all( + callable(getattr(mod, name, None)) + for name in ("autotune_v2", "autotune_v2_reload") + ) + + @functools.cache def has_flashinfer_comm() -> bool: """Return `True` if FlashInfer comm module is available.""" @@ -1269,6 +1282,7 @@ __all__ = [ "flashinfer_fused_kda_decode", "autotune", "has_flashinfer_moe", + "has_flashinfer_autotune_v2", "has_flashinfer_comm", "has_flashinfer_nvlink_two_sided", "has_flashinfer_nvlink_one_sided", From ce8853d27a977bd5d955c7e03c3d1dd9c735ef86 Mon Sep 17 00:00:00 2001 From: Sherif Waly Date: Sat, 17 Oct 2026 03:11:12 +0100 Subject: [PATCH 37/49] [Bugfix][KV Offload] Allow explicit SimpleCPUOffload block counts (#61674) Signed-off-by: Sherif Waly Co-authored-by: vllm-agent --- docs/features/kv_offloading_usage.md | 29 +++++++++ tests/v1/simple_kv_offload/test_scheduler.py | 49 +++++++++++++++ tests/v1/simple_kv_offload/test_worker.py | 63 +++++++++++++++++++ .../v1/simple_cpu_offload_connector.py | 18 ++++++ vllm/v1/simple_kv_offload/manager.py | 18 +++++- vllm/v1/simple_kv_offload/worker.py | 23 ++++++- 6 files changed, 195 insertions(+), 5 deletions(-) diff --git a/docs/features/kv_offloading_usage.md b/docs/features/kv_offloading_usage.md index e3575af63d..98159e38a7 100644 --- a/docs/features/kv_offloading_usage.md +++ b/docs/features/kv_offloading_usage.md @@ -411,6 +411,35 @@ Example (OpenAI-compatible completions request): } ``` +## Explicit pool sizing for `SimpleCPUOffloadConnector` + +`SimpleCPUOffloadConnector` accepts an optional `num_cpu_blocks` in +`kv_connector_extra_config`. The scheduler and every worker use this exact +number of offload blocks instead of automatic sizing from +`cpu_bytes_to_use_per_rank` or `cpu_bytes_to_use / world_size`. +Use the same positive integer on all ranks, chosen +to fit the smallest per-worker capacity. This supports pipeline stages with +different KV cache sizes without adding a collective during initialization. + +```json +{ + "kv_connector": "SimpleCPUOffloadConnector", + "kv_role": "kv_both", + "kv_connector_extra_config": { + "cpu_bytes_to_use_per_rank": 1073741824, + "num_cpu_blocks": 64, + "lazy_offload": false + } +} +``` + +The example count must be adjusted for the model, KV layout, and parallelism. +The configured byte budget remains an upper bound: startup fails if the +requested count exceeds a worker's capacity. For the `disk` backend the count +selects disk slots and is checked against `disk_capacity_bytes`. Omitting +`num_cpu_blocks` preserves automatic per-worker sizing. This option applies +to `SimpleCPUOffloadConnector`, separately from `OffloadingConnector` above. + ## Further Reading - [vLLM blog: KV Offloading Connector](https://vllm.ai/blog/2026-01-08-kv-offloading-connector) — motivation, architecture (DMA-based async transfer), and benchmarks (TTFT and throughput). diff --git a/tests/v1/simple_kv_offload/test_scheduler.py b/tests/v1/simple_kv_offload/test_scheduler.py index 527f7bee05..31aa75d31f 100644 --- a/tests/v1/simple_kv_offload/test_scheduler.py +++ b/tests/v1/simple_kv_offload/test_scheduler.py @@ -245,6 +245,55 @@ def test_connector_reports_configured_cache_source( assert connector.get_external_cache_hit_sources(None, 0) == {} # type: ignore[arg-type] +@pytest.mark.parametrize("num_layers", [1, 2]) +@pytest.mark.parametrize("backend", ["cpu", "disk"]) +def test_connector_explicit_cpu_blocks_align_uneven_stages( + tmp_path, num_layers: int, backend: str +): + """The same configured pool size must fit lighter and heavier PP stages.""" + config = _make_kv_cache_config(num_blocks=16) + names = [f"layer_{i}" for i in range(num_layers)] + config.kv_cache_groups[0].layer_names = names + config.kv_cache_tensors[0].layers = names + config.kv_cache_tensors[0].size *= num_layers + vllm_config = _make_vllm_config() + extra = { + "cpu_bytes_to_use_per_rank": _BYTES_PER_BLOCK * 8, + "num_cpu_blocks": 4, + "kv_offload_backend": backend, + } + if backend == "disk": + extra.update( + disk_path=str(tmp_path / "offload.bin"), + disk_capacity_bytes=_BYTES_PER_BLOCK * 8, + ) + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config = extra + connector = SimpleCPUOffloadConnector( + vllm_config, KVConnectorRole.SCHEDULER, config + ) + assert connector.scheduler_manager is not None + assert connector.scheduler_manager.num_cpu_blocks == 4 + assert len(connector.scheduler_manager.cpu_block_pool.blocks) == 4 + + +@pytest.mark.parametrize( + "count, capacity_bytes", + [(count, _BYTES_PER_BLOCK * 8) for count in [0, -1, True, 1.5, "4", 9]] + + [(1, _BYTES_PER_BLOCK // 2)], +) +def test_connector_rejects_invalid_explicit_cpu_blocks(count, capacity_bytes): + config = _make_kv_cache_config(num_blocks=16) + vllm_config = _make_vllm_config() + assert vllm_config.kv_transfer_config is not None + vllm_config.kv_transfer_config.kv_connector_extra_config = { + "cpu_bytes_to_use_per_rank": capacity_bytes, + "num_cpu_blocks": count, + } + with pytest.raises(ValueError, match="num_cpu_blocks"): + SimpleCPUOffloadConnector(vllm_config, KVConnectorRole.SCHEDULER, config) + + def make_scheduler( num_cpu_blocks: int = 8, num_gpu_blocks: int = 16, diff --git a/tests/v1/simple_kv_offload/test_worker.py b/tests/v1/simple_kv_offload/test_worker.py index 4f9a7c451f..c58cea6788 100644 --- a/tests/v1/simple_kv_offload/test_worker.py +++ b/tests/v1/simple_kv_offload/test_worker.py @@ -371,6 +371,69 @@ def test_register_shared_kv_cache_storage(monkeypatch, layout: KVCacheLayout): } +@pytest.mark.parametrize("num_layers", [1, 2]) +def test_explicit_cpu_blocks_bound_live_worker_allocation(monkeypatch, num_layers): + """Different live KV strides must honor the same explicit CPU pool size.""" + num_blocks = 4 + names = [f"layer.{i}" for i in range(num_layers)] + raw = torch.zeros( + num_layers * num_blocks * BLOCK_BYTES, dtype=torch.int8, device="cuda" + ) + config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[ + KVCacheTensor( + size=raw.nbytes, + layers=names, + layer_stride=num_blocks * BLOCK_BYTES, + block_stride=BLOCK_BYTES, + ) + ], + kv_cache_groups=[], + ) + caches = dict(zip(names, raw.view(num_layers, num_blocks, BLOCK_BYTES))) + worker = SimpleCPUOffloadWorker(None, config, 8 * BLOCK_BYTES, num_cpu_blocks=4) + monkeypatch.setattr("vllm.v1.simple_kv_offload.worker.PIN_MEMORY", False) + try: + worker.register_kv_caches(caches) + assert worker.num_cpu_blocks == 4 + assert worker.cpu_kv_caches is not None + assert all(t.shape[0] == 4 for t in worker.cpu_kv_caches.values()) + assert sum(t.nbytes for t in worker.cpu_kv_caches.values()) <= 8 * BLOCK_BYTES + finally: + if worker._backend is not None: + worker._backend.shutdown() + + +@pytest.mark.parametrize( + "capacity_bytes, count, local_capacity", + [(8 * BLOCK_BYTES, 5, 4), (BLOCK_BYTES, 1, 0)], +) +def test_explicit_cpu_blocks_reject_heavier_worker_before_allocation( + capacity_bytes, count, local_capacity +): + """A count that fits the scheduler's stage must fail on a smaller worker.""" + names = ["layer.0", "layer.1"] + raw = torch.zeros(2 * 4 * BLOCK_BYTES, dtype=torch.int8, device="cuda") + config = KVCacheConfig( + num_blocks=4, + kv_cache_tensors=[ + KVCacheTensor( + size=raw.nbytes, + layers=names, + layer_stride=4 * BLOCK_BYTES, + block_stride=BLOCK_BYTES, + ) + ], + kv_cache_groups=[], + ) + worker = SimpleCPUOffloadWorker(None, config, capacity_bytes, num_cpu_blocks=count) + caches = dict(zip(names, raw.view(2, 4, BLOCK_BYTES))) + with pytest.raises(ValueError, match=f"local offload capacity {local_capacity}"): + worker.register_kv_caches(caches) + assert worker.cpu_kv_caches is None + + def test_register_kv_cache_storage_with_trailing_padding(monkeypatch): num_blocks = 4 block_bytes = 32 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py index 0bd8ebe3cd..49afa017f6 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/simple_cpu_offload_connector.py @@ -101,6 +101,16 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): ) cpu_capacity_per_rank = explicit + # When set, num_cpu_blocks overrides automatic pool sizing from byte budgets + # (including cpu_bytes_to_use and cpu_bytes_to_use_per_rank), allowing ranks + # with different KV layouts (e.g. PP stages) to use the same pool size. + # Byte budgets still bound capacity. + num_cpu_blocks = extra_config.get("num_cpu_blocks") + if num_cpu_blocks is not None and ( + type(num_cpu_blocks) is not int or num_cpu_blocks <= 0 + ): + raise ValueError("num_cpu_blocks must be a positive integer") + lazy_offload = bool(extra_config.get("lazy_offload", False)) kv_offload_backend = str(extra_config.get("kv_offload_backend", "cpu")) @@ -160,6 +170,12 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): kv_offload_backend, disk_path or "none", ) + if num_cpu_blocks is not None: + logger.info( + "Using explicit num_cpu_blocks=%d for offload pool sizing " + "instead of automatic sizing from byte budgets.", + num_cpu_blocks, + ) if role == KVConnectorRole.SCHEDULER: from vllm.v1.core.kv_cache_utils import resolve_kv_cache_block_sizes @@ -177,6 +193,7 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): lazy_offload=lazy_offload, disk_capacity_bytes=disk_capacity_bytes if disk_mode else 0, use_page_cache=use_page_cache if disk_mode else False, + num_cpu_blocks=num_cpu_blocks, ) elif role == KVConnectorRole.WORKER: self.worker_handler = SimpleCPUOffloadWorker( @@ -188,6 +205,7 @@ class SimpleCPUOffloadConnector(KVConnectorBase_V1, SupportsHMA): disk_capacity_bytes=disk_capacity_bytes, disk_buffer_slots=disk_buffer_slots, use_page_cache=use_page_cache, + num_cpu_blocks=num_cpu_blocks, ) # --- Worker-side methods --- diff --git a/vllm/v1/simple_kv_offload/manager.py b/vllm/v1/simple_kv_offload/manager.py index dfc24e420a..8cf5f4f985 100644 --- a/vllm/v1/simple_kv_offload/manager.py +++ b/vllm/v1/simple_kv_offload/manager.py @@ -151,6 +151,7 @@ class SimpleCPUOffloadScheduler: lazy_offload: bool = False, disk_capacity_bytes: int = 0, use_page_cache: bool = False, + num_cpu_blocks: int | None = None, ): self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config @@ -170,7 +171,7 @@ class SimpleCPUOffloadScheduler: # Derive a CPU KVCacheConfig from the GPU config and build a coordinator assert kv_cache_config is not None self.cpu_kv_cache_config = self._derive_cpu_config( - kv_cache_config, offload_capacity + kv_cache_config, offload_capacity, num_cpu_blocks ) self.num_cpu_blocks = self.cpu_kv_cache_config.num_blocks self.prefix_cacheable_group_ids = ( @@ -278,7 +279,9 @@ class SimpleCPUOffloadScheduler: @staticmethod def _derive_cpu_config( - gpu_config: "KVCacheConfig", cpu_capacity_bytes: int + gpu_config: "KVCacheConfig", + cpu_capacity_bytes: int, + num_cpu_blocks: int | None = None, ) -> "KVCacheConfig": """Derive a CPU KVCacheConfig from the GPU config. Same kv_cache_groups, num_blocks scaled by CPU/GPU memory ratio.""" @@ -291,7 +294,16 @@ class SimpleCPUOffloadScheduler: # so its size is the total GPU KV cache size. gpu_total_bytes = gpu_config.kv_cache_tensors[0].size num_gpu_blocks = gpu_config.num_blocks - num_cpu_blocks = max(1, num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes) + local_capacity = num_gpu_blocks * cpu_capacity_bytes // gpu_total_bytes + if num_cpu_blocks is None: + num_cpu_blocks = max(1, local_capacity) + elif ( + type(num_cpu_blocks) is not int or not 0 < num_cpu_blocks <= local_capacity + ): + raise ValueError( + f"num_cpu_blocks={num_cpu_blocks!r} must be between 1 and " + f"local offload capacity {local_capacity}" + ) # Create CPU kv_cache_tensors mirroring GPU by scaling size proportionally. cpu_tensors = [ KVCacheTensor( diff --git a/vllm/v1/simple_kv_offload/worker.py b/vllm/v1/simple_kv_offload/worker.py index c120951eed..2ad318b987 100644 --- a/vllm/v1/simple_kv_offload/worker.py +++ b/vllm/v1/simple_kv_offload/worker.py @@ -36,10 +36,12 @@ class SimpleCPUOffloadWorker: disk_capacity_bytes: int = 0, disk_buffer_slots: int = 2, use_page_cache: bool = False, + num_cpu_blocks: int | None = None, ): self.vllm_config = vllm_config self.kv_cache_config = kv_cache_config self.cpu_capacity_bytes = cpu_capacity_bytes + self._num_cpu_blocks_override = num_cpu_blocks self.disk_path = disk_path self.disk_capacity_bytes = disk_capacity_bytes self.disk_buffer_slots = disk_buffer_slots @@ -159,7 +161,10 @@ class SimpleCPUOffloadWorker: ] total_bytes_per_block = sum(per_tensor_bpb) - self.num_cpu_blocks = max(1, self.cpu_capacity_bytes // total_bytes_per_block) + self.num_cpu_blocks = self._resolve_num_cpu_blocks( + self.disk_capacity_bytes if self.disk_mode else self.cpu_capacity_bytes, + total_bytes_per_block, + ) # Use lowest priority so KV cache I/O yields to compute streams. low_pri, _ = torch.cuda.Stream.priority_range() @@ -173,13 +178,27 @@ class SimpleCPUOffloadWorker: else: self._init_cpu_mode(unique_gpu_caches, total_bytes_per_block, self.device) + def _resolve_num_cpu_blocks(self, capacity_bytes: int, bytes_per_block: int) -> int: + local_capacity = capacity_bytes // bytes_per_block + requested = self._num_cpu_blocks_override + if requested is None: + return max(1, local_capacity) + if type(requested) is not int or not 0 < requested <= local_capacity: + raise ValueError( + f"num_cpu_blocks={requested!r} must be between 1 and " + f"local offload capacity {local_capacity}" + ) + return requested + def _init_disk_mode( self, unique_gpu_caches: dict[str, torch.Tensor], total_bytes_per_block: int, device: torch.device, ) -> None: - num_disk_slots = max(1, self.disk_capacity_bytes // total_bytes_per_block) + num_disk_slots = self._resolve_num_cpu_blocks( + self.disk_capacity_bytes, total_bytes_per_block + ) self.num_cpu_blocks = num_disk_slots logger.info( From ce49c47cbb241b59a32ef110079a5261dbd08e49 Mon Sep 17 00:00:00 2001 From: Akshat Anand Date: Sat, 17 Oct 2026 07:56:57 +0530 Subject: [PATCH 38/49] [Bugfix][Sampler] Carry the boundary tie count across blocks in the split top-p mask kernel (#62046) Signed-off-by: cipheraxat Co-authored-by: cipheraxat Co-authored-by: Bowen Wang --- tests/v1/sample/test_topk_topp_sampler.py | 30 +++++++++++++++++++++++ vllm/v1/sample/ops/topk_topp_triton.py | 1 + 2 files changed, 31 insertions(+) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index 190a4e3fca..f78f48b3e8 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -821,6 +821,36 @@ class TestTritonTopkTopp: ) assert len(kept_idx) == pytorch_kept + def test_topp_ties_across_blocks(self): + """The split top-p tie budget must carry across a slice's blocks. + + Regression: the split top-p mask kernel restarted its count of + boundary-value copies at every 8192-entry block, so each block of a + slice could keep up to `numkeep` copies of its own. + """ + from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton + + batch_size, vocab_size, num_ties = 64, 151936, 4096 + # 4096 tied tokens spread over many blocks, between 64 tokens above + # and a low tail; the p=0.75 boundary lands mid-copy in the tie group. + tie_idx = torch.arange(num_ties) * 37 + logits = torch.full((batch_size, vocab_size), -8.0) + logits[:, tie_idx] = 0.0 + logits[:, -64:] = 4.0 + p = torch.full((batch_size,), 0.75, dtype=torch.float32) + + kept = apply_top_k_top_p_triton(logits.clone(), None, p) != float("-inf") + expected = apply_top_k_top_p_pytorch(logits.clone(), None, p) != float("-inf") + num_kept, num_expected = kept.sum(dim=-1), expected.sum(dim=-1) + assert torch.equal(num_kept, num_expected), ( + f"kept {num_kept.unique().tolist()}, " + f"expected {num_expected.unique().tolist()}" + ) + # Kept copies of the boundary value are the first ones in index order. + kept_ties = kept[:, tie_idx] + num_kept_ties = kept_ties.sum(dim=-1, keepdim=True) + assert torch.equal(kept_ties, torch.arange(num_ties) < num_kept_ties) + # ============================================================================= # FlashInfer top-k/top-p robustness tests diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index 4aa837e9d7..d5849a8d67 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -1513,6 +1513,7 @@ def _topp_sb_mask_kernel( cum = tl.cumsum(dmask.to(tl.int32)).to(tl.float32) + num_preceding keep_dup = dmask & (cum <= numkeep) keep = keep & (~dmask | keep_dup) + num_preceding += tl.sum(dmask.to(tl.float32)) tl.store(ROW + i + offs, tl.where(keep, x, MASK_VALUE), mask=mask_n) From dcd07967caa2537c68aeeaddca4315b1319512b1 Mon Sep 17 00:00:00 2001 From: Gabriel Wu <13583761+lucifer1004@users.noreply.github.com> Date: Sat, 17 Oct 2026 12:06:03 +0800 Subject: [PATCH 39/49] [CPU] Keep the hybrid block size across the engine core's config rerun (#61805) Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Co-authored-by: Claude Opus 5.5 --- .../test_cpu_linear_attn_chunked_prefix.py | 22 +++++++++++++++++ vllm/platforms/cpu.py | 24 ++++++++++++------- 2 files changed, 38 insertions(+), 8 deletions(-) diff --git a/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py b/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py index 1b6a48b7d2..d56df03dbc 100644 --- a/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py +++ b/tests/v1/e2e/test_cpu_linear_attn_chunked_prefix.py @@ -8,6 +8,7 @@ import pytest from tests.models.utils import check_logprobs_close from vllm import LLM, SamplingParams +from vllm.inputs import TokensPrompt from vllm.platforms import current_platform pytestmark = pytest.mark.cpu_model @@ -112,3 +113,24 @@ def test_prefix_cache_hit_matches_cold_cache(full_prefill_refs): name_0="cold_cache", name_1="warm_cache", ) + + +def test_prefix_cache_hit_reaches_last_full_block(): + """A repeated prompt must hit every full block of the hybrid block size. + + The workers raise the block size to fit the GDN state pages, and the + engine core reruns its config checks after its handshake; the raised size + must survive them, or the prefill's state checkpoint misses the block + boundaries and nothing is reused. One chunk holds the whole prompt. + """ + llm = _make_llm(enable_prefix_caching=True, max_num_batched_tokens=2048) + block_size = llm.llm_engine.vllm_config.cache_config.block_size + # Ends mid-block: only the last full block's state can be reused. + num_tokens = block_size + block_size // 2 + assert num_tokens <= 2048 + prompt = TokensPrompt(prompt_token_ids=[1000 + i for i in range(num_tokens)]) + sp = SamplingParams(max_tokens=1) + hits = [llm.generate(prompt, sp)[0].num_cached_tokens for _ in range(2)] + del llm + + assert hits == [0, block_size] diff --git a/vllm/platforms/cpu.py b/vllm/platforms/cpu.py index 542a7a405d..91ff96ab33 100644 --- a/vllm/platforms/cpu.py +++ b/vllm/platforms/cpu.py @@ -289,10 +289,9 @@ class CpuPlatform(Platform): cache_config.block_size, ) cache_config.block_size = 256 - elif not cache_config.user_specified_block_size: - cache_config.block_size = 128 - - if not reference_cpu_mla_enabled and cache_config.block_size % 32 != 0: + elif ( + cache_config.user_specified_block_size and cache_config.block_size % 32 != 0 + ): logger.warning( "CPU backend prefers block_size is multiples of 32, " "otherwise the performance is not optimized." @@ -532,15 +531,24 @@ class CpuPlatform(Platform): @classmethod def update_block_size_for_backend(cls, vllm_config: "VllmConfig") -> None: model_config = vllm_config.model_config - if model_config is None or not model_config.is_hybrid: + if model_config is None: return - - # reconcile attention and mamba page sizes backend_classes = cls._find_non_ssm_backends(vllm_config) if not backend_classes: return - pre_block_size = vllm_config.cache_config.block_size + # Not in check_and_update_config(), which reruns with every + # VllmConfig.__post_init__ and would undo the hybrid alignment below. + cache_config = vllm_config.cache_config + if not cache_config.user_specified_block_size: + cache_config.block_size = cls._preferred_block_size_for_backends( + backend_classes, 128, vllm_config + ) + if not model_config.is_hybrid: + return + + # reconcile attention and mamba page sizes + pre_block_size = cache_config.block_size cls._align_hybrid_block_size(vllm_config, backend_classes[0]) cls._check_aligned_block_size(vllm_config, backend_classes, pre_block_size) From 556ac15eb930875693acc5489cf860ac509708a9 Mon Sep 17 00:00:00 2001 From: Fadi Arafeh <115173828+fadara01@users.noreply.github.com> Date: Sat, 17 Oct 2026 05:11:33 +0100 Subject: [PATCH 40/49] [Bugfix] Remove extra args from CPU's compute slot mapping kernel (#61992) Signed-off-by: Fadi Arafeh Co-authored-by: vllm-agent --- .../model_executor/test_triton_dispatcher.py | 41 +++++++++++++++++++ vllm/utils/cpu_triton_utils.py | 8 +--- 2 files changed, 43 insertions(+), 6 deletions(-) diff --git a/tests/model_executor/test_triton_dispatcher.py b/tests/model_executor/test_triton_dispatcher.py index 909230c40f..09faa74a2d 100644 --- a/tests/model_executor/test_triton_dispatcher.py +++ b/tests/model_executor/test_triton_dispatcher.py @@ -220,3 +220,44 @@ def test_real_triton_kernel_override_end_to_end() -> None: "out_ptr", "BLOCK", ) + + +def test_cpu_slot_mapping_override_forwards_kernel_arguments( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from vllm.utils import cpu_triton_utils + from vllm.v1.worker.block_table import ComputeSlotMappingKernel + + calls: list[tuple[Any, ...]] = [] + fake_torch = types.SimpleNamespace( + ops=types.SimpleNamespace( + _C=types.SimpleNamespace( + compute_slot_mapping_kernel_impl=lambda *args: calls.append(args) + ) + ) + ) + monkeypatch.setattr(cpu_triton_utils, "torch", fake_torch) + + override = KernelOverride( + ComputeSlotMappingKernel.kernel, + cpu_triton_utils._compute_slot_mapping_kernel_impl, + ) + override[(2,)]( + num_tokens=2, + max_num_tokens=4, + query_start_loc_ptr="query_start_loc", + positions_ptr="positions", + block_table_ptr="block_table", + block_table_stride=128, + slot_mapping_ptr="slot_mapping", + KV_CACHE_BLOCK_SIZE=16, + TOTAL_CP_WORLD_SIZE=1, + TOTAL_CP_RANK=0, + CP_KV_CACHE_INTERLEAVE_SIZE=1, + PAD_ID=-1, + BLOCK_SIZE=1024, + ) + + assert calls == [ + ("query_start_loc", "positions", "block_table", "slot_mapping", 16) + ] diff --git a/vllm/utils/cpu_triton_utils.py b/vllm/utils/cpu_triton_utils.py index d8e4cd389d..3950e4c2ac 100644 --- a/vllm/utils/cpu_triton_utils.py +++ b/vllm/utils/cpu_triton_utils.py @@ -16,10 +16,8 @@ def _compute_slot_mapping_kernel_impl( positions: torch.Tensor, # [num_tokens], int64 block_table: torch.Tensor, # [max_num_reqs, max_num_blocks_per_req], int32 block_table_stride: int, # max_num_blocks_per_req - block_size: int, slot_mapping: torch.Tensor, # [max_num_tokens], int64 - KV_CACHE_BLOCK_SIZE: int | None = None, - BLOCKS_PER_KV_BLOCK: int = 1, + KV_CACHE_BLOCK_SIZE: int, TOTAL_CP_WORLD_SIZE: int = 1, TOTAL_CP_RANK: int = 0, CP_KV_CACHE_INTERLEAVE_SIZE: int = 1, @@ -27,14 +25,12 @@ def _compute_slot_mapping_kernel_impl( BLOCK_SIZE: int = 1024, ) -> None: assert TOTAL_CP_WORLD_SIZE == 1, "Context Parallelism is not supported on CPU." - if BLOCKS_PER_KV_BLOCK != 1: - assert block_size * BLOCKS_PER_KV_BLOCK == KV_CACHE_BLOCK_SIZE torch.ops._C.compute_slot_mapping_kernel_impl( query_start_loc, positions, block_table, slot_mapping, - block_size, + KV_CACHE_BLOCK_SIZE, ) From 7b6310c109135b1b4aba42ed363697987380fd10 Mon Sep 17 00:00:00 2001 From: Bowen Wang Date: Sat, 17 Oct 2026 12:22:11 +0800 Subject: [PATCH 41/49] [Perf][Sampler] Speed up the small-batch split top-p pipeline (#61996) Signed-off-by: Bowen Wang Signed-off-by: Nick Hill Co-authored-by: Claude Opus 5.5 Co-authored-by: Nick Hill Co-authored-by: vllm-agent --- tests/v1/sample/test_topk_topp_sampler.py | 54 ++++++++++++++++++- tests/v1/worker/test_gpu_sampler_flags.py | 24 +++++++++ tests/watermarking/test_watermarking.py | 8 ++- vllm/model_executor/models/diffusion_gemma.py | 7 ++- vllm/v1/sample/ops/topk_topp_sampler.py | 14 ++++- vllm/v1/sample/ops/topk_topp_triton.py | 22 +++++--- vllm/v1/watermarking/gpu_sampler.py | 6 ++- vllm/v1/worker/gpu/sample/sampler.py | 10 ++-- vllm/v1/worker/gpu/sample/states.py | 27 +++++++--- 9 files changed, 148 insertions(+), 24 deletions(-) diff --git a/tests/v1/sample/test_topk_topp_sampler.py b/tests/v1/sample/test_topk_topp_sampler.py index f78f48b3e8..0cd44fe213 100644 --- a/tests/v1/sample/test_topk_topp_sampler.py +++ b/tests/v1/sample/test_topk_topp_sampler.py @@ -293,7 +293,7 @@ class TestTritonTopkTopp: self._compare_results(logits, k, p=None) - @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024]) + @pytest.mark.parametrize("batch_size", [1, 8, 32, 48, 128, 512, 1024]) @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256]) def test_topp_only(self, batch_size: int, vocab_size: int): """Test top-p only (k=None).""" @@ -307,7 +307,7 @@ class TestTritonTopkTopp: self._compare_results(logits, k=None, p=p) - @pytest.mark.parametrize("batch_size", [1, 8, 32, 128, 512, 1024]) + @pytest.mark.parametrize("batch_size", [1, 8, 32, 48, 128, 512, 1024]) @pytest.mark.parametrize("vocab_size", [1024, 32000, 128256]) def test_topk_and_topp(self, batch_size: int, vocab_size: int): """Test combined top-k and top-p.""" @@ -328,6 +328,24 @@ class TestTritonTopkTopp: self._compare_results(logits, k, p) + @pytest.mark.parametrize("batch_size", [8, 48, 64]) + @pytest.mark.parametrize("vocab_size", [151936, 200064, 262144]) + def test_topk_and_topp_large_vocab(self, batch_size: int, vocab_size: int): + """Mixed batches at vocab sizes where the split top-p pipeline uses + its largest split counts (up to _SPLIT_MAX_SPLITS).""" + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + k = torch.randint(1, 100, (batch_size,), generator=self.generator) + p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 + # Rows cycle through top-k + top-p, top-p only, top-k only, neither. + k[1::4] = vocab_size + p[2::4] = 1.0 + k[3::4] = vocab_size + p[3::4] = 1.0 + + self._compare_results(logits, k, p) + def test_both_disabled(self): """Test when both k and p are None (should be no-op).""" from vllm.v1.sample.ops.topk_topp_triton import apply_top_k_top_p_triton @@ -788,6 +806,38 @@ class TestTritonTopkTopp: f"{kept_topk_topp.tolist()}" ) + @pytest.mark.parametrize("batch_size", [8, 48]) + def test_skip_split_without_topp_only_rows( + self, batch_size: int, monkeypatch: pytest.MonkeyPatch + ): + """has_topp_only_rows=False skips the split pipeline; the result must + not change when every top-p row also has a top-k, grammar rows + (k >= #finite logits) included.""" + import vllm.v1.sample.ops.topk_topp_triton as tt + + vocab_size = 32000 + logits = torch.randn( + batch_size, vocab_size, generator=self.generator, dtype=torch.float32 + ) + logits[1] = float("-inf") + idx = torch.randperm(vocab_size, generator=self.generator)[:30] + logits[1, idx] = torch.randn(30, generator=self.generator) + k = torch.randint( + 50, 100, (batch_size,), generator=self.generator, dtype=torch.int32 + ) + p = torch.rand(batch_size, generator=self.generator) * 0.9 + 0.1 + k[::4] = vocab_size + p[::4] = 1.0 + + expected = tt.apply_top_k_top_p_triton(logits.clone(), k, p) + monkeypatch.setattr( + tt, "_apply_topp_split", lambda *args: pytest.fail("split launched") + ) + actual = tt.apply_top_k_top_p_triton( + logits.clone(), k, p, has_topp_only_rows=False + ) + assert torch.equal(actual, expected) + def test_topp_tie_break_deterministic(self): """Boundary-duplicate trimming must be deterministic across calls. diff --git a/tests/v1/worker/test_gpu_sampler_flags.py b/tests/v1/worker/test_gpu_sampler_flags.py index 3f05a2cfb5..7982be1938 100644 --- a/tests/v1/worker/test_gpu_sampler_flags.py +++ b/tests/v1/worker/test_gpu_sampler_flags.py @@ -109,6 +109,30 @@ def test_custom_processor_add_request_gates_pipeline_flag(admitted: bool): assert sampler.needs_logits_processing[3] == admitted +def test_has_topp_only_rows(): + """A row is top-p-only when it has top-p and no effective top-k; top_k + of 0 (the default) or at least the logits width disables top-k.""" + sampler = _make_sampler() + sampler.add_request(0, SamplingParams(top_k=10, top_p=0.9)) + sampler.add_request(1, SamplingParams(top_k=10)) + sampler.add_request(2, SamplingParams(top_p=0.9)) + sampler.add_request(3, SamplingParams(top_k=VOCAB_SIZE + 1, top_p=0.9)) + states = sampler.sampling_states + states.apply_staged_writes() + + def has_topp_only_rows(idx: list[int], vocab_size: int = VOCAB_SIZE) -> bool: + idx_np = np.array(idx, dtype=np.int32) + idx_gpu = torch.from_numpy(idx_np).to(DEVICE) + return states.get_top_k_top_p(idx_gpu, idx_np, vocab_size)[2] + + assert not has_topp_only_rows([0, 1]) + assert has_topp_only_rows([0, 2]) + assert has_topp_only_rows([3]) + # The kernel compares top_k against the logits width, so a top_k at or + # above a narrower logits width counts as no top-k. + assert has_topp_only_rows([0], vocab_size=10) + + def test_logits_processing_cache_only_checks_active_requests(): sampler = _make_sampler() sampler.add_request(0, SamplingParams(temperature=0.0)) diff --git a/tests/watermarking/test_watermarking.py b/tests/watermarking/test_watermarking.py index eec43845db..3379e385a0 100644 --- a/tests/watermarking/test_watermarking.py +++ b/tests/watermarking/test_watermarking.py @@ -368,6 +368,7 @@ def test_gpu_sampler_respects_mixed_request_watermarking( None, None, False, + False, ) assert torch.equal(sampled, torch.tensor([3, 7])) @@ -383,7 +384,7 @@ def test_gpu_sampler_filters_top_k_top_p_before_watermarking( # reference path; the assertion is about the call order, not the kernel. monkeypatch.setattr( "vllm.v1.watermarking.gpu_sampler.apply_top_k_top_p", - lambda logits, k, p: apply_top_k_top_p_pytorch( + lambda logits, k, p, **kwargs: apply_top_k_top_p_pytorch( logits, k, p, allow_cpu_sync=True ), ) @@ -415,6 +416,7 @@ def test_gpu_sampler_filters_top_k_top_p_before_watermarking( torch.tensor([2]), torch.tensor([0.8]), False, + False, ) assert watermarker.captured_logits is not None @@ -482,6 +484,7 @@ def test_gpu_sampler_skips_watermarking_for_repeated_contexts(monkeypatch): None, None, False, + False, ) assert torch.equal(sampled, torch.tensor([3, 7])) @@ -517,6 +520,7 @@ def test_gpu_sampler_can_disable_context_deduplication(monkeypatch): None, None, False, + False, ) assert torch.equal(sampled, torch.tensor([7, 7])) @@ -956,6 +960,7 @@ def test_gpu_sampler_uses_fused_gumbel_for_repeated_contexts(): None, None, False, + False, ) assert torch.equal(repeated, torch.tensor([True, False], device=device)) @@ -992,6 +997,7 @@ def test_gpu_sampler_skips_watermarking_for_greedy_batch(monkeypatch): None, None, False, + False, ) assert actual is expected diff --git a/vllm/model_executor/models/diffusion_gemma.py b/vllm/model_executor/models/diffusion_gemma.py index 7f4f9ffba9..541ae4e2f1 100644 --- a/vllm/model_executor/models/diffusion_gemma.py +++ b/vllm/model_executor/models/diffusion_gemma.py @@ -1404,14 +1404,17 @@ class DiffusionSampler: # unaffected; only the canvas exploration is constrained. Applied # before canvas padding so phantom positions stay uniform. if num_decode > 0: - top_k, top_p = self.sampling_states.get_top_k_top_p( + top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p( decode_slots.repeat_interleave( valid_canvas_len, output_size=int(valid_canvas_len_np.sum()) ), decode_slots_np, + logits.shape[-1], ) if top_k is not None or top_p is not None: - logits = apply_top_k_top_p(logits.float(), top_k, top_p) + logits = apply_top_k_top_p( + logits.float(), top_k, top_p, has_topp_only_rows=has_topp_only_rows + ) # Where each decode request's rows start in the flat logits. Tiles # below gather and pad a request to the tile's width. diff --git a/vllm/v1/sample/ops/topk_topp_sampler.py b/vllm/v1/sample/ops/topk_topp_sampler.py index 085f092341..02d62298cd 100644 --- a/vllm/v1/sample/ops/topk_topp_sampler.py +++ b/vllm/v1/sample/ops/topk_topp_sampler.py @@ -484,13 +484,23 @@ class TopKTopPSampler(nn.Module): def apply_top_k_top_p( - logits: torch.Tensor, k: torch.Tensor | None, p: torch.Tensor | None + logits: torch.Tensor, + k: torch.Tensor | None, + p: torch.Tensor | None, + has_topp_only_rows: bool = True, ) -> torch.Tensor: + """Apply top-k and/or top-p masking to logits. + + has_topp_only_rows must be True if any row has p < 1 and + k >= logits.shape[-1] (no active top-k). + """ if p is None and k is None: return logits if HAS_TRITON: - return apply_top_k_top_p_triton(logits, k, p) + return apply_top_k_top_p_triton( + logits, k, p, has_topp_only_rows=has_topp_only_rows + ) is_cpu = current_platform.is_cpu() return apply_top_k_top_p_pytorch(logits, k, p, allow_cpu_sync=is_cpu) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index d5849a8d67..6945adc662 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -987,13 +987,19 @@ _SPLIT_MAX_BATCH = 64 _SPLIT_MAX_SPLITS = 32 # power of two _SPLIT_FANOUT = 8 _SPLIT_ROUNDS = 5 +# The step kernels are latency-bound at one program per SM, so keep splitting +# past that while slices span at least this many entries; shorter slices no +# longer amortize each program's row stats and combine replay. +_SPLIT_MIN_SLICE = 7168 _TRITON_SPLIT_CACHE: dict[torch.device, dict[str, torch.Tensor]] = {} -def _topp_split_count(batch_size: int, num_sm: int) -> int: +def _topp_split_count(batch_size: int, vocab_size: int, num_sm: int) -> int: splits = 1 - while splits < _SPLIT_MAX_SPLITS and splits * 2 * batch_size <= num_sm: + while splits < _SPLIT_MAX_SPLITS and ( + splits * 2 * batch_size <= num_sm or splits * 2 * _SPLIT_MIN_SLICE <= vocab_size + ): splits *= 2 return splits @@ -1079,7 +1085,7 @@ def _topp_split_stats( num_sm: int, ) -> DispatchSpec: batch_size, vocab_size = logits.shape - splits = _topp_split_count(batch_size, num_sm) + splits = _topp_split_count(batch_size, vocab_size, num_sm) return (logits.shape[0] * splits,), dict( K=k if k is not None else logits, HAS_K=k is not None, @@ -1373,7 +1379,7 @@ def _topp_split_step( num_sm: int, ) -> DispatchSpec: batch_size, vocab_size = logits.shape - splits = _topp_split_count(batch_size, num_sm) + splits = _topp_split_count(batch_size, vocab_size, num_sm) return (logits.shape[0] * splits,), dict( K=k if k is not None else logits, HAS_K=k is not None, @@ -1555,7 +1561,7 @@ def _topp_split_mask( num_sm: int, ) -> DispatchSpec: batch_size, vocab_size = logits.shape - splits = _topp_split_count(batch_size, num_sm) + splits = _topp_split_count(batch_size, vocab_size, num_sm) return (logits.shape[0] * splits,), dict( K=k if k is not None else logits, HAS_K=k is not None, @@ -1619,6 +1625,7 @@ def apply_top_k_top_p_triton( k: torch.Tensor | None, p: torch.Tensor | None, mask_value: float = float("-inf"), + has_topp_only_rows: bool = True, ) -> torch.Tensor: """Apply combined top-k and top-p masking using Triton. @@ -1633,6 +1640,9 @@ def apply_top_k_top_p_triton( p: [batch_size] float32 tensor of top-p values per row (0 to 1), or None to disable top-p mask_value: Value for masked positions (default: -inf) + has_topp_only_rows: Whether any row has p < 1 and k >= vocab_size. + False lets small batches skip the split top-p pipeline; pass it + only if no such row exists, as such a row would get no top-p. Returns: The masked logits tensor. It may or may not be modified in-place. @@ -1715,7 +1725,7 @@ def apply_top_k_top_p_triton( mask_value, num_sm, ) - if use_split: + if use_split and has_topp_only_rows: _apply_topp_split(logits, k_ptr, p_ptr, mask_value, num_sm) return logits diff --git a/vllm/v1/watermarking/gpu_sampler.py b/vllm/v1/watermarking/gpu_sampler.py index 776dece208..48547057fa 100644 --- a/vllm/v1/watermarking/gpu_sampler.py +++ b/vllm/v1/watermarking/gpu_sampler.py @@ -58,6 +58,7 @@ class GPUWatermarkSampler(Sampler): pos: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, + has_topp_only_rows: bool, use_fused_sampler: bool, ) -> tuple[torch.Tensor, torch.Tensor]: enabled = self.watermarking.np[idx_mapping_np] & ( @@ -71,10 +72,13 @@ class GPUWatermarkSampler(Sampler): pos, top_k, top_p, + has_topp_only_rows, use_fused_sampler, ) - processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) + processed_logits = apply_top_k_top_p( + processed_logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows + ) contexts = self._get_contexts(expanded_idx_mapping) repeated_contexts = None if self.deduplicate_contexts != "none": diff --git a/vllm/v1/worker/gpu/sample/sampler.py b/vllm/v1/worker/gpu/sample/sampler.py index 7ffc5bf493..4ad56292e4 100644 --- a/vllm/v1/worker/gpu/sample/sampler.py +++ b/vllm/v1/worker/gpu/sample/sampler.py @@ -298,8 +298,8 @@ class Sampler: seq_lens_upper_bound_np, skip_top_k_top_p=True, ) - top_k, top_p = self.sampling_states.get_top_k_top_p( - expanded_idx_mapping, idx_mapping_np + top_k, top_p, has_topp_only_rows = self.sampling_states.get_top_k_top_p( + expanded_idx_mapping, idx_mapping_np, logits.shape[-1] ) # Don't use a fused sampler if no requests use top_k/top_p, if there are # any greedy requests or per-request seeds, or if post-processed @@ -321,6 +321,7 @@ class Sampler: pos, top_k, top_p, + has_topp_only_rows, use_fused_sampler, ) @@ -332,6 +333,7 @@ class Sampler: pos: torch.Tensor, top_k: torch.Tensor | None, top_p: torch.Tensor | None, + has_topp_only_rows: bool, use_fused_sampler: bool, ) -> tuple[torch.Tensor, torch.Tensor]: if use_fused_sampler: @@ -342,7 +344,9 @@ class Sampler: else: # Use XPU sampler sampled, _ = xpu_sample(processed_logits, top_k, top_p) else: - processed_logits = apply_top_k_top_p(processed_logits, top_k, top_p) + processed_logits = apply_top_k_top_p( + processed_logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows + ) sampled = gumbel_sample( processed_logits, expanded_idx_mapping, diff --git a/vllm/v1/worker/gpu/sample/states.py b/vllm/v1/worker/gpu/sample/states.py index 47cf5de856..9e598b5a91 100644 --- a/vllm/v1/worker/gpu/sample/states.py +++ b/vllm/v1/worker/gpu/sample/states.py @@ -103,13 +103,22 @@ class SamplingStates: apply_min_p(logits, expanded_idx_mapping, self.min_p.gpu) def get_top_k_top_p( - self, expanded_idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray - ) -> tuple[torch.Tensor | None, torch.Tensor | None]: - do_top_k = np.any(self.top_k.np[idx_mapping_np] != self.vocab_size) - do_top_p = np.any(self.top_p.np[idx_mapping_np] != 1.0) + self, + expanded_idx_mapping: torch.Tensor, + idx_mapping_np: np.ndarray, + vocab_size: int, + ) -> tuple[torch.Tensor | None, torch.Tensor | None, bool]: + """Returns (top_k, top_p, has_topp_only_rows).""" + top_k_np = self.top_k.np[idx_mapping_np] + top_p_np = self.top_p.np[idx_mapping_np] + do_top_k = np.any(top_k_np != self.vocab_size) + do_top_p = np.any(top_p_np != 1.0) top_k = self.top_k.gpu[expanded_idx_mapping] if do_top_k else None top_p = self.top_p.gpu[expanded_idx_mapping] if do_top_p else None - return top_k, top_p + has_topp_only_rows = bool( + do_top_p and np.any((top_k_np >= vocab_size) & (top_p_np < 1.0)) + ) + return top_k, top_p, has_topp_only_rows def apply_top_k_top_p( self, @@ -117,10 +126,14 @@ class SamplingStates: expanded_idx_mapping: torch.Tensor, idx_mapping_np: np.ndarray, ) -> torch.Tensor: - top_k, top_p = self.get_top_k_top_p(expanded_idx_mapping, idx_mapping_np) + top_k, top_p, has_topp_only_rows = self.get_top_k_top_p( + expanded_idx_mapping, idx_mapping_np, logits.shape[-1] + ) if top_k is None and top_p is None: return logits - return apply_top_k_top_p(logits, top_k, top_p) + return apply_top_k_top_p( + logits, top_k, top_p, has_topp_only_rows=has_topp_only_rows + ) def any_greedy(self, idx_mapping_np: np.ndarray) -> bool: return bool(np.any(self.temperature.np[idx_mapping_np] == 0.0)) From a0e7137c4bb94dcb057247334978a17364a0a127 Mon Sep 17 00:00:00 2001 From: Micah Williamson Date: Fri, 16 Oct 2026 23:25:39 -0500 Subject: [PATCH 42/49] [CI][ROCm][Bugfix] DeepSeek V4.1: don't pass expert intermediate padding to AITER (#61960) Signed-off-by: Micah Williamson Co-authored-by: Cursor --- vllm/models/deepseek_v4/amd/model.py | 4 ++++ vllm/models/deepseek_v41/amd/model.py | 4 ++++ 2 files changed, 8 insertions(+) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index f933cffcda..d1a8ca1d9d 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -578,6 +578,9 @@ def _validate_aiter_mega_moe_config(vllm_config: VllmConfig) -> None: class DeepseekV4MoE(nn.Module): + # None: use the padding vLLM added. + experts_intermediate_pad: int | None = None + def __init__( self, vllm_config: VllmConfig, @@ -726,6 +729,7 @@ class DeepseekV4MoE(nn.Module): top_k=config.num_experts_per_tok, hidden_size=config.hidden_size, intermediate_size=config.moe_intermediate_size, + intermediate_pad=self.experts_intermediate_pad, renormalize=config.norm_topk_prob, quant_config=quant_config, prefix=f"{prefix}.experts", diff --git a/vllm/models/deepseek_v41/amd/model.py b/vllm/models/deepseek_v41/amd/model.py index 57366f2453..9b8a00388a 100644 --- a/vllm/models/deepseek_v41/amd/model.py +++ b/vllm/models/deepseek_v41/amd/model.py @@ -77,6 +77,10 @@ logger = init_logger(__name__) class DeepseekV4MoE(DeepseekV4MoEBase): + # TODO: revert once https://github.com/ROCm/aiter/pull/6598 is in the + # pinned AITER; with the padding set, its FlyDSL MoE can output NaN. + experts_intermediate_pad = 0 + def __init__( self, vllm_config: VllmConfig, From 41b8c158097c57ccf61626680d0fd13790bd120b Mon Sep 17 00:00:00 2001 From: YukioZzz Date: Sat, 17 Oct 2026 12:57:58 +0800 Subject: [PATCH 43/49] [Feature][MoRIIO] Support hybrid WRITE transfers with DSpark (#60544) Signed-off-by: Yichao Zhu Signed-off-by: tjtanaa Co-authored-by: OpenAI Codex Co-authored-by: tjtanaa Co-authored-by: Claude Opus 5.5 (1M context) --- .buildkite/test_areas/misc.yaml | 53 ++ docs/features/moriio_connector_usage.md | 8 + .../moriio_toy_proxy_server.py | 4 +- .../moriio_integration/__init__.py | 2 + .../moriio_integration/moriio_wide_dummy.py | 55 ++ .../test_write_dummy_e2e_rocm.py | 736 ++++++++++++++++++ .../unit/test_moriio_connector.py | 6 + .../unit/test_moriio_hma_scheduler.py | 321 +++++++- .../unit/test_moriio_kv_layout.py | 94 ++- .../unit/test_moriio_transfer_completion.py | 61 ++ .../unit/test_moriio_write_done_routing.py | 84 +- .../kv_connector/v1/moriio/moriio_common.py | 7 +- .../v1/moriio/moriio_connector.py | 145 ++-- .../kv_connector/v1/moriio/moriio_engine.py | 101 ++- 14 files changed, 1589 insertions(+), 88 deletions(-) create mode 100644 tests/v1/kv_connector/moriio_integration/__init__.py create mode 100644 tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py create mode 100644 tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py diff --git a/.buildkite/test_areas/misc.yaml b/.buildkite/test_areas/misc.yaml index c79b451749..53ce97f6be 100644 --- a/.buildkite/test_areas/misc.yaml +++ b/.buildkite/test_areas/misc.yaml @@ -164,6 +164,59 @@ steps: depends_on: - image-build-amd +# The shared KV Connectors steps above exclude vllm/distributed/kv_transfer/ from +# their triggers, and CUDA skips MoRIIO's mori/aiter-gated tests. Run the whole +# MoRIIO suite on ROCm whenever the connector or its tests change. +- label: ":amd: (MI355 DPX) MoRIIO Connector Unit Tests" + key: amd-moriio-connector-unit-mi355 + dind: false + device: mi355_dpx + num_devices: 1 + timeout_in_minutes: 30 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/unit/test_moriio_ + - tests/v1/kv_connector/unit/utils.py + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/unit/test_moriio_*.py + +# Single-node MoRIIO WRITE P/D on dummy weights: prefill and decode on separate +# GPUs over xGMI, compared exactly against a direct serve. +- label: ":amd: (MI355) MoRIIO WRITE P/D Dummy-Weight E2E Shard %N" + key: amd-moriio-write-dummy-e2e-mi355 + dind: false + device: mi355_2 + num_devices: 2 + # ~3.5-5 min per case; pytest-shard hashes test ids, so with 20 cases the + # largest of 6 shards holds 5 (~22 min on MI355). Every case picks free ports + # for its servers and proxy and its own compile-cache dir. + parallelism: 6 + timeout_in_minutes: 45 + working_dir: /vllm-workspace/tests + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/kv_transfer/kv_connector/v1/moriio/ + - vllm/distributed/kv_transfer/kv_connector/v1/base.py + - vllm/distributed/kv_transfer/kv_connector/v1/ssm_conv_transfer_utils.py + - vllm/model_executor/model_loader/dummy_loader.py + - examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py + - tests/v1/kv_connector/moriio_integration/ + - requirements/kv_connectors_rocm.txt + commands: + - bash /vllm-workspace/.buildkite/scripts/install-kv-connectors.sh + - export VLLM_WORKER_MULTIPROC_METHOD=spawn + - pytest -v -s -rs v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py --shard-id=$$BUILDKITE_PARALLEL_JOB --num-shards=$$BUILDKITE_PARALLEL_JOB_COUNT + - label: ":nvidia: (H200 MIG 35GB) V1 Metrics + LM Eval" device: h200_35gb key: v1-metrics-lmeval diff --git a/docs/features/moriio_connector_usage.md b/docs/features/moriio_connector_usage.md index 6c906e7db1..1c67908a58 100644 --- a/docs/features/moriio_connector_usage.md +++ b/docs/features/moriio_connector_usage.md @@ -136,6 +136,14 @@ WRITE mode is used by default. READ mode can be configured by setting `--kv-tran !!! note `notify_port` is used as a *base* port: each (DP rank, TP rank) pair within an instance uses `notify_port + offset` where the offset is based on the rank. Make sure the range starting at `notify_port` is free on the host. +### Hybrid attention and recurrent state + +WRITE mode supports hybrid attention plus Mamba/KDA models, including DSpark, with equal producer and consumer TP sizes. Use matching cache layouts on both instances, `--mamba-cache-mode align`, and `read_mode: false` in each instance's `kv_connector_extra_config`. Both peers must support hybrid WRITE. Sliding-window WRITE and hybrid speculative methods other than DSpark remain unsupported. + +For DSpark, configure the same draft model on both instances through `--speculative-config`. Target and draft attention groups retain their own block IDs and token spans; only the running recurrent state is transferred, not speculative scratch slots. The producer transfers the state after the first N-1 prompt tokens, and the consumer recomputes the last prompt token. Convolution and SSM transfers must both finish before the layer is complete. + +Hybrid READ keeps its existing single-transferable-attention-group requirement. Plain-attention WRITE retains its existing notification format. + ### Transport configuration MoRI has two transport backends: RDMA and xGMI. You can select backend using `--kv-transfer-config.kv_connector_extra_config.backend $BACKEND`, with `$BACKEND` being `rdma` or `xgmi`. RDMA is the default backend and should be used in multi-node deployments. diff --git a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py index ce4d558fb9..e611b2e690 100644 --- a/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py +++ b/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py @@ -473,9 +473,11 @@ async def stop_profile(): if __name__ == "__main__": parser = argparse.ArgumentParser() parser.add_argument("--port", type=int, default=10001) + # Port the prefill/decode servers register on (their proxy_ping_port). + parser.add_argument("--service-discovery-port", type=int, default=36367) args = parser.parse_args() - t = start_service_discovery("0.0.0.0", 36367) + t = start_service_discovery("0.0.0.0", args.service_discovery_port) # High-concurrency hardening. Quart's app.run() uses a shallow listen # backlog (100) and, with app.debug=True, adds per-request overhead that # starves the single accept loop. Under a burst of ~512 simultaneous client diff --git a/tests/v1/kv_connector/moriio_integration/__init__.py b/tests/v1/kv_connector/moriio_integration/__init__.py new file mode 100644 index 0000000000..208f01a7cb --- /dev/null +++ b/tests/v1/kv_connector/moriio_integration/__init__.py @@ -0,0 +1,2 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project diff --git a/tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py b/tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py new file mode 100644 index 0000000000..ac80a6bc9b --- /dev/null +++ b/tests/v1/kv_connector/moriio_integration/moriio_wide_dummy.py @@ -0,0 +1,55 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Widen the dummy weight init for the MoRIIO dummy-weight P/D test. + +The dummy loader fills float weights in +-1e-3, which keeps every activation so +small that outputs barely depend on the input, so a wrong KV prefix would go +unnoticed. The P/D test names ``WideDummyWeights`` as each server's +``--worker-extension-cls``. Workers import it before they load the model, which +re-registers the ``dummy`` load format in every worker process with a loader +that fills 16/32-bit float weights in +-``MORIIO_DUMMY_WEIGHT_RANGE``; 8-bit +weights and scales keep the default init. The format keeps the name ``dummy`` +so every config path that special-cases dummy weights still applies. + +The test process itself never imports this module. +""" + +import os + +import torch +from torch import nn + +from vllm.config import ModelConfig +from vllm.model_executor.model_loader import register_model_loader +from vllm.model_executor.model_loader.dummy_loader import DummyModelLoader +from vllm.model_executor.model_loader.reload.layerwise import get_layerwise_info +from vllm.model_executor.model_loader.reload.utils import get_layer_tensors +from vllm.model_executor.model_loader.weight_utils import ( + initialize_single_dummy_weight, +) + +_RANGE = float(os.environ["MORIIO_DUMMY_WEIGHT_RANGE"]) + + +@register_model_loader("dummy") +class WideDummyModelLoader(DummyModelLoader): + def load_weights(self, model: nn.Module, model_config: ModelConfig) -> None: + # DummyModelLoader.load_weights with a wider init for 16/32-bit floats. + for layer in model.modules(): + info = get_layerwise_info(layer) + if info.can_load(): + self._process_online_quant_layer(layer, info) + continue + for name, tensor in get_layer_tensors(layer).items(): + if name in layer._non_persistent_buffers_set: + continue + if tensor.is_floating_point() and torch.finfo(tensor.dtype).bits >= 16: + initialize_single_dummy_weight( + tensor, low=-_RANGE, high=_RANGE, name=name + ) + else: + initialize_single_dummy_weight(tensor, name=name) + + +class WideDummyWeights: + """Worker extension that only exists to import this module in workers.""" diff --git a/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py new file mode 100644 index 0000000000..74c9a95698 --- /dev/null +++ b/tests/v1/kv_connector/moriio_integration/test_write_dummy_e2e_rocm.py @@ -0,0 +1,736 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Single-node MoRIIO WRITE P/D checks on dummy weights. + +Each case serves a real model architecture with ``--load-format dummy``, widened +by ``moriio_wide_dummy`` (see ``DUMMY_WEIGHT_RANGE``): the +config keeps the real hidden size, head geometry, MLA ranks and recurrent-state +shapes, so the connector registers and transfers exactly the production cache +layout, but the layer stack is cut to a few layers that still contain every +cache-group kind and the MoE is cut to a handful of experts. Dummy weights are +seeded, so the prefill, decode and reference servers hold identical weights. +Cases marked ``HMA_off`` run every server with +``--disable-hybrid-kv-cache-manager``, so all layers, sliding-window ones +included, share a single cache group. + +Prefill runs on one GPU, decode on another, and MoRI moves the KV over xGMI. +Either way decode computes the last prompt token itself. For hybrid models +prefill computes only the first ``n - 1`` tokens and transfers the KV and the +recurrent state after them. For attention-only models prefill computes and +transfers all ``n``, and decode's scheduler, seeing the whole prompt cached, +recomputes the last token. A direct serve splits the same prompt the same way +when the first ``n - 1`` tokens end on a chunk boundary that both servers share, +so the prompts are chosen to make that happen: one is a single chunk plus one +token, the others are a whole number of hybrid cache blocks plus one token +(prefix caching ends chunks on block boundaries) or, for attention-only models, +a whole number of chunks plus one token. With the split matched, a correct +transfer reproduces the direct serve's first-token log-probabilities exactly, +so the comparison needs no noise allowance. The +servers compile with Inductor's deterministic mode, since otherwise each one +benchmarks its own reduction launch configs and separately started servers +can drift apart by a few bf16 logit steps. Guards: + +* decode must report an external prefix-cache hit for every transferred token + of every prompt, so the match is not a local recompute; +* changing the first half of a prompt must move the direct serve's log-probs by + at least ``MIN_CONTEXT_EFFECT``, so a wrong prefix would be visible; +* ``test_corrupted_prefill_state_is_detected`` makes prefill compute its KV with + a different norm epsilon (same shapes, different values) and requires every + prompt to fail the comparison. +""" + +import atexit +import contextlib +import json +import os +import random +import shutil +import signal +import subprocess +import sys +import tempfile +import time +from collections.abc import Callable, Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from pathlib import Path +from typing import Any + +import pytest +import requests + +from tests.utils import RemoteOpenAIServer +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port + +_ROOT = Path(__file__).resolve().parents[4] +_PROXY = ( + _ROOT / "examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py" +) + +# Max allowed difference between centered first-token log-probabilities of the +# P/D and direct outputs, over the direct serve's top-k tokens. The split is +# matched, so correct transfers produce identical logits; this only absorbs the +# fp32 rounding of the log-softmax (~1e-7), far below one bf16 logit step. +TOLERANCE = float(os.getenv("MORIIO_DUMMY_E2E_TOL", "1e-5")) +# Several bf16 steps of the logits at DUMMY_WEIGHT_RANGE. +MIN_CONTEXT_EFFECT = 0.05 +# Half-width of the uniform dummy init for 16/32-bit float weights, applied in +# every server worker by moriio_wide_dummy.py. The loader default (1e-3) shrinks +# every activation so far that a context change moves the first-token log-probs +# by ~1e-4; at 0.1 it moves them by >0.1 and a corrupted prefill changes the +# top-k set entirely. +DUMMY_WEIGHT_RANGE = 0.1 +_WIDE_DUMMY_DIR = Path(__file__).resolve().parent +TOP_K = 20 +MAX_TOKENS = 4 +MAX_MODEL_LEN = 4096 + +# Prefill token budget of every server. Prompts longer than it are chunked, so +# the final-chunk save path runs; the short prompt fills exactly one chunk. +PREFILL_CHUNK = 299 +SHORT_PROMPT_LEN = PREFILL_CHUNK + 1 +# Hybrid blocks per aligned prompt, before the trailing token. +ALIGNED_BLOCKS = (1, 2) +# Prefill chunks per long attention-only prompt, before the trailing token. +ALIGNED_CHUNKS = (2, 4) + + +def _mori_available() -> bool: + try: + import mori.io # noqa: F401 + except ImportError: + return False + return True + + +pytestmark = [ + pytest.mark.skipif( + not current_platform.is_rocm() or not _mori_available(), + reason="MoRIIO P/D needs ROCm and MoRI", + ), + pytest.mark.skipif( + current_platform.is_rocm() and current_platform.device_count() < 2, + reason="needs one GPU for prefill and one for decode", + ), +] + + +# --------------------------------------------------------------------------- +# Model architectures + + +def _kimi_linear_overrides(cfg) -> dict[str, Any]: + # 3 KDA layers then 1 MLA layer; layer 0 stays dense (first_k_dense_replace). + linear = dict(cfg.linear_attn_config) + linear["kda_layers"] = [1, 2, 3] + linear["full_attn_layers"] = [4] + groups = max(1, cfg.num_expert_group) + return { + "num_hidden_layers": 4, + "linear_attn_config": linear, + "num_experts": 2 * groups, + "num_experts_per_token": 2, + "topk_group": min(cfg.topk_group, groups), + } + + +def _qwen3_5_moe_overrides(cfg) -> dict[str, Any]: + # 3 Gated DeltaNet layers then 1 full-attention layer. + return { + "text_config": { + "num_hidden_layers": 4, + "layer_types": ["linear_attention"] * 3 + ["full_attention"], + "num_experts": 4, + "num_experts_per_tok": 2, + } + } + + +def _qwen3_dense_overrides(cfg) -> dict[str, Any]: + # Attention-only: decode gets the KV for the whole prompt. + return {"num_hidden_layers": 4} + + +def _deepseek_v3_overrides(cfg) -> dict[str, Any]: + # MLA attention-only: layer 0 dense, layer 1 MoE, no MTP layer. + return { + "num_hidden_layers": 2, + "first_k_dense_replace": 1, + "n_routed_experts": 8, + "num_experts_per_tok": 2, + "n_group": 1, + "topk_group": 1, + "num_nextn_predict_layers": 0, + } + + +def _gemma3_overrides(cfg) -> dict[str, Any]: + # Five sliding-window layers then one global layer, as in the real + # 5:1 pattern. Without the hybrid KV cache manager vLLM gives the sliding + # layers full-attention cache specs, so all six share one group. + return { + "num_hidden_layers": 6, + "layer_types": ["sliding_attention"] * 5 + ["full_attention"], + } + + +@dataclass(frozen=True) +class ArchCase: + model: str + overrides: Callable[[Any], dict[str, Any]] + # Norm-epsilon override that perturbs every computed KV/state value + # without changing a single shape. Used by the corruption control. + corrupt_overrides: dict[str, Any] + env: dict[str, str] = field(default_factory=dict) + args: tuple[str, ...] = () + # Has recurrent-state (mamba-style) layers, so prefill stops at n - 1. + hybrid: bool = True + + +ARCHS: dict[str, ArchCase] = { + "kimi_linear": ArchCase( + model="moonshotai/Kimi-Linear-48B-A3B-Instruct", + overrides=_kimi_linear_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "1"}, + args=("--trust-remote-code", "--mamba-ssm-cache-dtype", "float32"), + ), + "qwen3_5_moe": ArchCase( + model="Qwen/Qwen3.5-35B-A3B", + overrides=_qwen3_5_moe_overrides, + corrupt_overrides={"text_config": {"rms_norm_eps": 1.0}}, + # AITER's GDN conv-update kernel rejects the DS conv-state layout that + # MoRIIO's recurrent-state transfer requires. + env={"VLLM_ROCM_USE_AITER": "0"}, + args=("--language-model-only",), + ), + "qwen3_dense": ArchCase( + model="Qwen/Qwen3-32B", + overrides=_qwen3_dense_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "0"}, + hybrid=False, + ), + "deepseek_v3": ArchCase( + model="deepseek-ai/DeepSeek-V3", + overrides=_deepseek_v3_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "1"}, + hybrid=False, + ), + # Text-only Gemma 3. The multimodal Gemma3ForConditionalGeneration takes + # the inputs_embeds path on KV consumers (enable_mm_embeds is forced on + # there) and computes differently from a direct serve without any transfer. + "gemma3": ArchCase( + model="google/gemma-3-1b-it", + overrides=_gemma3_overrides, + corrupt_overrides={"rms_norm_eps": 1.0}, + env={"VLLM_ROCM_USE_AITER": "0"}, + hybrid=False, + ), +} + + +def _arch_overrides(arch: ArchCase) -> dict[str, Any]: + from vllm.transformers_utils.config import get_config + + return arch.overrides(get_config(arch.model, trust_remote_code=True)) + + +def _merge(base: dict[str, Any], extra: dict[str, Any]) -> dict[str, Any]: + merged = dict(base) + for key, value in extra.items(): + if isinstance(value, dict) and isinstance(merged.get(key), dict): + merged[key] = _merge(merged[key], value) + else: + merged[key] = value + return merged + + +# --------------------------------------------------------------------------- +# Serving + + +@dataclass(frozen=True) +class Case: + arch: str + attention_backend: str + prefill_graph: str + decode_graph: str + # Run every server of the case without the hybrid KV cache manager. + hma_off: bool = False + + @property + def extra_args(self) -> tuple[str, ...]: + return ("--disable-hybrid-kv-cache-manager",) if self.hma_off else () + + @property + def id(self) -> str: + return ( + f"{self.arch}-{self.attention_backend}-" + f"P_{self.prefill_graph}-D_{self.decode_graph}" + + ("-HMA_off" if self.hma_off else "") + ) + + +_DEFAULT = "FULL_AND_PIECEWISE" +CASES = [ + # Kimi-Linear: KDA + MLA groups. + Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT), + Case("kimi_linear", "ROCM_AITER_MLA", "NONE", "NONE"), + Case("kimi_linear", "ROCM_AITER_MLA", "PIECEWISE", "PIECEWISE"), + Case("kimi_linear", "ROCM_AITER_MLA", _DEFAULT, "FULL_DECODE_ONLY"), + Case("kimi_linear", "TRITON_MLA", _DEFAULT, _DEFAULT), + # Qwen3.5: Gated DeltaNet + full-attention groups. + Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, _DEFAULT), + Case("qwen3_5_moe", "TRITON_ATTN", "NONE", "NONE"), + Case("qwen3_5_moe", "TRITON_ATTN", "PIECEWISE", "PIECEWISE"), + Case("qwen3_5_moe", "TRITON_ATTN", _DEFAULT, "FULL_DECODE_ONLY"), + # ROCM_ATTN keeps a (2, num_blocks, ...) KV layout and opts out of KV + # connectors, so the second ROCm attention backend is the unified one. + Case("qwen3_5_moe", "ROCM_AITER_UNIFIED_ATTN", _DEFAULT, _DEFAULT), + # Attention-only: prefill transfers the whole prompt. + Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT), + Case("qwen3_dense", "ROCM_AITER_UNIFIED_ATTN", _DEFAULT, _DEFAULT), + # Without the hybrid KV cache manager: one cache group per model. + Case("qwen3_dense", "TRITON_ATTN", _DEFAULT, _DEFAULT, hma_off=True), + Case("deepseek_v3", "ROCM_AITER_MLA", _DEFAULT, _DEFAULT, hma_off=True), + Case("gemma3", "TRITON_ATTN", _DEFAULT, _DEFAULT, hma_off=True), +] + + +def _graph_args(mode: str) -> list[str]: + if mode == "NONE": + return ["--enforce-eager"] + # Inductor otherwise picks reduction launch configs by timing them, so each + # separately compiled server may sum in a different order and drift by a + # few bf16 logit steps; deterministic mode fixes the choice. + config = { + "cudagraph_mode": mode, + "inductor_compile_config": {"deterministic": True}, + } + return ["--compilation-config", json.dumps(config)] + + +def _visible_devices() -> tuple[str, list[str]]: + selectors = ("HIP_VISIBLE_DEVICES", "ROCR_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES") + selector = next((k for k in selectors if k in os.environ), selectors[0]) + visible = os.getenv(selector) + devices = ( + visible.split(",") + if visible + else [str(i) for i in range(current_platform.device_count())] + ) + return selector, devices + + +def _device_env(selector: str, device: str) -> dict[str, str]: + # vLLM mirrors HIP_VISIBLE_DEVICES into CUDA_VISIBLE_DEVICES at import, so + # the test process carries both; pin every selector or the server refuses + # the mismatch. ROCR narrows the set first, so HIP/CUDA then index into it. + env = {selector: device} + for key in ("HIP_VISIBLE_DEVICES", "CUDA_VISIBLE_DEVICES"): + if key != selector: + env[key] = "0" if selector == "ROCR_VISIBLE_DEVICES" else device + return env + + +# KV consumers force enable_mm_embeds on, which moves multimodal-capable models +# (Qwen3.5) from input_ids to inputs_embeds, but the AOT compile-cache key does +# not cover it, so a consumer can load a producer's graph and crash in +# profile_run. Keep each role's compile cache apart until +# vllm-project/vllm#54642 lands (vllm-project/vllm#50891). +_CACHE_ROOT = Path(tempfile.mkdtemp(prefix="moriio_dummy_e2e_cache_")) +atexit.register(shutil.rmtree, _CACHE_ROOT, ignore_errors=True) + + +def _server( + arch: ArchCase, + overrides: dict[str, Any], + backend: str, + graph: str, + device: str, + selector: str, + *, + port: int | None = None, + extra_args: tuple[str, ...] = (), + kv_transfer_config: dict[str, Any] | None = None, +) -> RemoteOpenAIServer: + port = port or get_open_port() + args = [ + *arch.args, + *extra_args, + "--port", + str(port), + "--load-format", + "dummy", + "--worker-extension-cls", + "moriio_wide_dummy.WideDummyWeights", + "--hf-overrides", + json.dumps(overrides), + "--dtype", + "bfloat16", + "--attention-backend", + backend, + "--max-model-len", + str(MAX_MODEL_LEN), + "--max-num-batched-tokens", + str(PREFILL_CHUNK), + "--max-num-seqs", + "16", + "--kv-cache-memory-bytes", + str(2 * 1024**3), + "--gpu-memory-utilization", + "0.5", + "--enable-prefix-caching", + *_graph_args(graph), + ] + role = "direct" + if kv_transfer_config is not None: + args += ["--kv-transfer-config", json.dumps(kv_transfer_config)] + role = kv_transfer_config["kv_role"] + env = { + "PATH": str(Path(sys.executable).parent) + os.pathsep + os.environ["PATH"], + **_device_env(selector, device), + "VLLM_HOST_IP": "127.0.0.1", + "VLLM_PORT": str(get_open_port()), + "VLLM_WORKER_MULTIPROC_METHOD": "spawn", + "VLLM_SSM_CONV_STATE_LAYOUT": "DS", + "VLLM_CACHE_ROOT": str(_CACHE_ROOT / role), + # Only this directory, so the source tree cannot shadow installed vllm. + "PYTHONPATH": os.pathsep.join( + filter(None, [str(_WIDE_DUMMY_DIR), os.environ.get("PYTHONPATH")]) + ), + "MORIIO_DUMMY_WEIGHT_RANGE": str(DUMMY_WEIGHT_RANGE), + **arch.env, + } + return RemoteOpenAIServer( + arch.model, args, env_dict=env, auto_port=False, max_wait_seconds=900 + ) + + +def _moriio_config(role: str, http_port: int, ping_port: int) -> dict[str, Any]: + return { + "kv_connector": "MoRIIOConnector", + "kv_role": role, + "kv_connector_extra_config": { + "backend": "xgmi", + "read_mode": False, + "host_ip": "127.0.0.1", + "proxy_ip": "127.0.0.1", + "proxy_ping_port": str(ping_port), + "http_port": str(http_port), + "handshake_port": str(get_open_port()), + "notify_port": str(get_open_port()), + }, + } + + +@contextlib.contextmanager +def _proxy(ping_port: int) -> Iterator[str]: + # Free ports for both sockets, so cases from parallel shards on one host + # do not collide. + port = get_open_port() + process = subprocess.Popen( + [ + sys.executable, + str(_PROXY), + "--port", + str(port), + "--service-discovery-port", + str(ping_port), + ], + start_new_session=True, + ) + try: + deadline = time.monotonic() + 60 + while time.monotonic() < deadline: + assert process.poll() is None, "MoRIIO proxy exited during startup" + try: + requests.get(f"http://127.0.0.1:{port}/", timeout=1) + break + except requests.ConnectionError: + time.sleep(0.2) + else: + raise TimeoutError("MoRIIO proxy did not start") + yield f"http://127.0.0.1:{port}" + finally: + if process.poll() is None: + os.killpg(process.pid, signal.SIGTERM) + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + os.killpg(process.pid, signal.SIGKILL) + process.wait(timeout=10) + + +@contextlib.contextmanager +def _pd_pair( + arch: ArchCase, + prefill_overrides: dict[str, Any], + decode_overrides: dict[str, Any], + case: Case, +) -> Iterator[tuple[str, RemoteOpenAIServer]]: + selector, devices = _visible_devices() + assert len(devices) >= 2, "needs one GPU for prefill and one for decode" + prefill_port, decode_port, ping_port = (get_open_port() for _ in range(3)) + with _proxy(ping_port) as proxy_url, contextlib.ExitStack() as stack: + servers: list[RemoteOpenAIServer] = [] + stack.callback(RemoteOpenAIServer.shutdown_many, servers) + with ThreadPoolExecutor(2) as pool: + futures = [ + pool.submit( + _server, + arch, + prefill_overrides, + case.attention_backend, + case.prefill_graph, + devices[0], + selector, + port=prefill_port, + extra_args=case.extra_args, + kv_transfer_config=_moriio_config( + "kv_producer", prefill_port, ping_port + ), + ), + pool.submit( + _server, + arch, + decode_overrides, + case.attention_backend, + case.decode_graph, + devices[1], + selector, + port=decode_port, + extra_args=case.extra_args, + kv_transfer_config=_moriio_config( + "kv_consumer", decode_port, ping_port + ), + ), + ] + for future in futures: + servers.append(future.result()) + _wait_for_registration(proxy_url, arch.model) + yield proxy_url, servers[1] + + +def _wait_for_registration(proxy_url: str, model: str) -> None: + deadline = time.monotonic() + 120 + while time.monotonic() < deadline: + # The proxy hands decode max_tokens - 1 (prefill emits the first + # token), so the probe needs at least 2 for decode to accept it. + try: + response = requests.post( + f"{proxy_url}/v1/completions", + json={"model": model, "prompt": [1000, 1001], "max_tokens": 2}, + timeout=300, + ) + # The proxy streams the decode body, so a decode-side error can + # surface as a truncated response rather than a status code. + if response.status_code == 200 and response.json()["choices"]: + return + except requests.RequestException: + pass + time.sleep(1) + raise TimeoutError("prefill/decode never registered with the MoRIIO proxy") + + +# --------------------------------------------------------------------------- +# Measurement + + +def _metric(server: RemoteOpenAIServer, name: str) -> list[Any]: + from prometheus_client.parser import text_string_to_metric_families + + response = requests.get(server.url_for("metrics"), timeout=10) + response.raise_for_status() + return [ + sample + for family in text_string_to_metric_families(response.text) + for sample in family.samples + if sample.name == name + ] + + +def _block_size(server: RemoteOpenAIServer) -> int: + # Hybrid models pad the attention block to the recurrent-state page, so the + # effective size is only known once the server has started. + sizes = { + int(sample.labels["block_size"]) + for sample in _metric(server, "vllm:cache_config_info") + } + assert len(sizes) == 1, f"expected one cache block size, got {sizes}" + return sizes.pop() + + +def _prompt_lengths(arch: ArchCase, server: RemoteOpenAIServer) -> list[int]: + # Hybrid chunks end on cache-block boundaries; attention-only chunks end on + # the prefill budget, so the last token is computed on its own either way. + if arch.hybrid: + unit, counts = _block_size(server), ALIGNED_BLOCKS + else: + unit, counts = PREFILL_CHUNK, ALIGNED_CHUNKS + aligned = [ + k * unit + 1 for k in counts if k * unit + 1 + MAX_TOKENS <= MAX_MODEL_LEN + ] + assert aligned, f"alignment unit {unit} leaves no aligned prompt" + return [SHORT_PROMPT_LEN, *aligned] + + +def _prompts(lengths: list[int]) -> list[list[int]]: + rng = random.Random(0) + return [[rng.randrange(1000, 30000) for _ in range(n)] for n in lengths] + + +def _first_step(url: str, model: str, prompt: list[int]) -> dict[str, float]: + response = requests.post( + f"{url}/v1/completions", + json={ + "model": model, + "prompt": prompt, + "max_tokens": MAX_TOKENS, + "temperature": 0, + "logprobs": TOP_K, + "return_tokens_as_token_ids": True, + }, + timeout=600, + ) + response.raise_for_status() + top = response.json()["choices"][0]["logprobs"]["top_logprobs"] + assert top, "server returned no logprobs" + return top[0] + + +def _distance(ref: dict[str, float], got: dict[str, float]) -> float: + """Max difference of centered log-probs over the shared top-k tokens. + + Centering removes the log-sum-exp offset, leaving the logit differences the + KV actually determines. + """ + common = sorted(ref.keys() & got.keys()) + if len(common) < len(ref) // 2: + return float("inf") + ref_mean = sum(ref[t] for t in common) / len(common) + got_mean = sum(got[t] for t in common) / len(common) + return max(abs((ref[t] - ref_mean) - (got[t] - got_mean)) for t in common) + + +def _external_hits(server: RemoteOpenAIServer) -> float: + samples = _metric(server, "vllm:external_prefix_cache_hits_total") + return sum(sample.value for sample in samples) + + +@dataclass +class Reference: + prompts: list[list[int]] + logprobs: list[dict[str, float]] + hybrid: bool + + @property + def expected_external_hits(self) -> int: + # Hybrid prefill stops one token short; attention-only prefill hands + # decode the whole prompt. + return sum(len(p) - int(self.hybrid) for p in self.prompts) + + +def _reference(arch: ArchCase, overrides: dict[str, Any], case: Case) -> Reference: + """Direct serve; also checks that the comparison can see a wrong prefix.""" + selector, devices = _visible_devices() + # Prefill graphs run the chunks and the decode-only last step runs the same + # captured graph in every decode mode used here, so the prefill mode + # reproduces both halves of the P/D computation. + with _server( + arch, + overrides, + case.attention_backend, + case.prefill_graph, + devices[0], + selector, + extra_args=case.extra_args, + ) as server: + url = server.url_root + lengths = _prompt_lengths(arch, server) + prompts = _prompts(lengths) + refs = [_first_step(url, arch.model, p) for p in prompts] + # Same tail, different early context: the comparison has to see it. + for prompt, ref in zip(prompts, refs): + altered = [t + 1 for t in prompt[: len(prompt) // 2]] + altered += prompt[len(prompt) // 2 :] + moved = _distance(ref, _first_step(url, arch.model, altered)) + print(f"[sensitivity] len={len(prompt)} context-change distance={moved}") + assert moved >= MIN_CONTEXT_EFFECT, ( + f"first-token log-probs barely depend on the context (moved " + f"{moved:.4g} for len {len(prompt)}); a wrong prefix would go " + "unnoticed" + ) + # Determinism of the reference itself (served from its prefix cache). + again = _first_step(url, arch.model, prompts[-1]) + print(f"[noise] repeat distance={_distance(refs[-1], again)}") + return Reference(prompts, refs, arch.hybrid) + + +def _pd_distances( + arch: ArchCase, + prefill_overrides: dict[str, Any], + decode_overrides: dict[str, Any], + case: Case, + ref: Reference, +) -> tuple[list[float], float]: + with _pd_pair(arch, prefill_overrides, decode_overrides, case) as ( + proxy_url, + decode, + ): + hits_before = _external_hits(decode) + distances = [ + _distance(expected, _first_step(proxy_url, arch.model, prompt)) + for prompt, expected in zip(ref.prompts, ref.logprobs) + ] + hits = _external_hits(decode) - hits_before + lengths = [len(p) for p in ref.prompts] + print(f"[pd] {case.id} lens={lengths} distances={distances} external_hits={hits}") + return distances, hits + + +# --------------------------------------------------------------------------- +# Tests + + +@pytest.mark.parametrize("case", CASES, ids=lambda c: c.id) +def test_write_matches_direct_serve(case: Case): + arch = ARCHS[case.arch] + overrides = _arch_overrides(arch) + ref = _reference(arch, overrides, case) + distances, hits = _pd_distances(arch, overrides, overrides, case, ref) + + # Every transferred token of every prompt must come from prefill. + assert hits == ref.expected_external_hits, ( + f"decode reported {hits} external prefix-cache hits, expected " + f"{ref.expected_external_hits}; it did not use all the transferred KV" + ) + lengths = [len(p) for p in ref.prompts] + bad = {n: d for n, d in zip(lengths, distances) if not d <= TOLERANCE} + assert not bad, f"P/D first-token log-probs differ from direct serve: {bad}" + + +@pytest.mark.parametrize("arch_name", sorted(ARCHS)) +def test_corrupted_prefill_state_is_detected(arch_name: str): + """Negative control: KV computed with different values must fail.""" + arch = ARCHS[arch_name] + default = next(c for c in CASES if c.arch == arch_name) + case = Case( + arch_name, default.attention_backend, _DEFAULT, _DEFAULT, default.hma_off + ) + overrides = _arch_overrides(arch) + ref = _reference(arch, overrides, case) + corrupted = _merge(overrides, arch.corrupt_overrides) + distances, hits = _pd_distances(arch, corrupted, overrides, case, ref) + + assert hits == ref.expected_external_hits, ( + "decode did not use the transferred KV; the control proves nothing" + ) + assert all(d > TOLERANCE for d in distances), ( + f"P/D outputs with wrongly computed prefill KV still matched the " + f"reference (distances {distances}); the comparison cannot detect " + "transfer corruption" + ) diff --git a/tests/v1/kv_connector/unit/test_moriio_connector.py b/tests/v1/kv_connector/unit/test_moriio_connector.py index 73cc5f0583..ffd78969a1 100644 --- a/tests/v1/kv_connector/unit/test_moriio_connector.py +++ b/tests/v1/kv_connector/unit/test_moriio_connector.py @@ -146,6 +146,7 @@ def _write_consumer_scheduler_for_finished_request(tp_size: int = 2): scheduler.is_producer = False scheduler.mode = MoRIIOMode.WRITE scheduler.tp_size = tp_size + scheduler._has_mamba = False scheduler._reqs_need_recv = {} scheduler.unmap_request_id = MagicMock() return scheduler @@ -158,6 +159,10 @@ def _write_producer_scheduler(block_size: int = 1) -> Any: scheduler.mode = MoRIIOMode.WRITE scheduler.is_producer = True scheduler.block_size = block_size + scheduler._has_mamba = False + scheduler.kv_cache_config = SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=object())] + ) scheduler.blocks_per_sw = [0] scheduler._group_block_sizes = [block_size] scheduler.transfer_id_to_request_id = {} @@ -944,6 +949,7 @@ def test_clamp_to_prompt_blocks_clamps_each_group_by_its_own_block_size(): # their own prompt-block count; a sliding-window group is left untouched. scheduler = MoRIIOConnectorScheduler.__new__(MoRIIOConnectorScheduler) scheduler.blocks_per_sw = [0, 0, 3] + scheduler._has_mamba = False scheduler._group_block_sizes = [16, 32, 16] req = SimpleNamespace(num_prompt_tokens=40) diff --git a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py index 6d8dc24ea2..1120062e54 100644 --- a/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py +++ b/tests/v1/kv_connector/unit/test_moriio_hma_scheduler.py @@ -12,12 +12,13 @@ wiring. import importlib from dataclasses import asdict from types import SimpleNamespace +from unittest.mock import Mock import pytest import torch from vllm.v1.core.block_pool import BlockPool -from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.output import KVConnectorBlockState, SchedulerOutput from vllm.v1.core.single_type_kv_cache_manager import FullAttentionManager from vllm.v1.kv_cache_interface import FullAttentionSpec, UniformTypeKVCacheSpecs @@ -52,6 +53,7 @@ class _FakeScheduler(moriio_connector.MoRIIOConnectorScheduler): # type: ignore attributes read by the method under test are set.""" def __init__(self, **attrs): + self.mode = MoRIIOMode.READ self._mamba_group_ids: list[int] = [] self._attn_group_ids: list[int] = [0] self._num_ssm_scratch_blocks = 0 @@ -287,7 +289,8 @@ def test_scheduler_rejects_mamba_group_without_two_states(): moriio_connector.MoRIIOConnectorScheduler(_gate_vllm_config(), "engine", config) -def test_scheduler_rejects_non_dspark_speculative_hybrid_read(): +@pytest.mark.parametrize("read_mode", [True, False]) +def test_scheduler_rejects_non_dspark_speculative_hybrid(read_mode): config = SimpleNamespace( kv_cache_groups=[ SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), @@ -298,7 +301,9 @@ def test_scheduler_rejects_non_dspark_speculative_hybrid_read(): with pytest.raises(moriio_common.MoRIIOError, match="speculative decoding"): moriio_connector.MoRIIOConnectorScheduler( - _gate_vllm_config(speculative_config=_spec_config("ngram")), + _gate_vllm_config( + read_mode=read_mode, speculative_config=_spec_config("ngram") + ), "engine", config, ) @@ -336,14 +341,16 @@ def _dspark_wrapped_attention_config(*, dcp_sharded: bool = True): @pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize("read_mode", [True, False]) @pytest.mark.parametrize( "dcp_size,dcp_sharded,expected_tail", [(1, True, 4), (4, True, 1), (4, False, 4)] ) def test_scheduler_accepts_dspark_with_wrapped_full_attention( - dcp_size, dcp_sharded, expected_tail + dcp_size, dcp_sharded, expected_tail, read_mode ): config = _dspark_wrapped_attention_config(dcp_sharded=dcp_sharded) vllm_config = _gate_vllm_config( + read_mode=read_mode, speculative_config=_spec_config("dspark"), block_size=4, num_lookahead_tokens=7, @@ -384,19 +391,234 @@ def test_scheduler_accepts_dspark_with_wrapped_full_attention( ) -def test_scheduler_rejects_hybrid_write(): - config = SimpleNamespace( - kv_cache_groups=[ - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=object()), - SimpleNamespace(enable_kv_transfer=True, kv_cache_spec=_mamba_spec()), - ], +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_clamps_attention_and_excludes_scratch_slots(): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config ) - config.transfer_groups = tuple(config.kv_cache_groups) + assert sched._clamp_to_prompt_blocks( + SimpleNamespace(num_prompt_tokens=9), [[10, 11, 12, 13], [1, 2, 3, 4]] + ) == [[11], [1, 2]] - with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"): - moriio_connector.MoRIIOConnectorScheduler( - _gate_vllm_config(read_mode=False), "engine", config + +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_resolves_current_state_instead_of_cached_block_ids(): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config + ) + state = KVConnectorBlockState( + req_ids={"req"}, + resolve_block_ids=lambda _: ([0, 51, 52, 53], [31, 32, 33]), + boundary_state_offloads={}, + ) + assert sched._get_write_blocks( + SimpleNamespace(request_id="req", num_prompt_tokens=9), + [[10, 11, 12], [1]], + SimpleNamespace(kv_connector_block_state=state), + ) == [[51], [31, 32]] + + +@pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize("external_tokens", [0, 9]) +def test_hybrid_write_notifies_grouped_slots_or_releases_a_complete_local_hit( + external_tokens, +): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False), "engine", config + ) + sched.send_notify_block = Mock() + sched._send_transfer_release = Mock() + req = _mk_request( + list(range(10)), + params={ + "do_remote_prefill": True, + "transfer_id": "tx", + "remote_host": "127.0.0.2", + "remote_notify_port": 7000, + }, + ) + req.request_id = "req" + sched.update_state_after_alloc( + req, _FakeBlocks(([10, 11, 12], [1, 2, 3])), external_tokens + ) + if external_tokens: + sched.send_notify_block.assert_called_once_with( + req_id="req", + transfer_id="tx", + block_notify_list=[[10], [1, 2, 3]], + host="127.0.0.2", + port=7000, ) + sched._send_transfer_release.assert_not_called() + else: + sched.send_notify_block.assert_not_called() + sched._send_transfer_release.assert_called_once_with("tx", "127.0.0.2", 7000) + assert not req.kv_transfer_params["do_remote_prefill"] + + +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_keeps_separate_draft_groups_and_their_block_spans(): + config = _dspark_wrapped_attention_config() + # A disabled cache is absent from the wire. Target and draft may have + # different token spans, and neither can borrow the other's block IDs. + draft = SimpleNamespace( + enable_kv_transfer=True, + kv_cache_spec=FullAttentionSpec( + block_size=4, + num_kv_heads=1, + head_size=2, + dtype=torch.float32, + dcp_sharded=False, + ), + ) + disabled = SimpleNamespace(enable_kv_transfer=False, kv_cache_spec=object()) + config.kv_cache_groups.insert(1, disabled) + config.kv_cache_groups.append(draft) + config.transfer_groups = tuple( + g for g in config.kv_cache_groups if g.enable_kv_transfer + ) + config.select_transfer_block_ids = lambda blocks: (blocks[0], blocks[2], blocks[3]) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config( + read_mode=False, dcp_size=2, speculative_config=_spec_config("dspark") + ), + "engine", + config, + ) + req = _mk_request( + list(range(9)), + params={ + "do_remote_decode": True, + "transfer_id": "tx", + }, + ) + req.request_id = "req" + req.status = moriio_connector.RequestStatus.FINISHED_LENGTH_CAPPED + blocks = ([10, 11, 12, 13], [99], [20, 21, 22], [30, 31, 32, 33]) + expected = [[11], [20, 21], [30, 31, 32]] + assert sched._clamp_to_prompt_blocks(req, blocks) == expected + # Exercise the connector entry point too: completion must not flatten + # a multi-group WRITE payload through READ's single-attention helper. + delay_free, params = _FakeConnector(sched).request_finished_all_groups(req, blocks) + assert delay_free + assert params["remote_block_ids"] == expected + + +@pytest.mark.usefixtures("restore_moriio_role") +@pytest.mark.parametrize( + "blocks,delay_free", + [(([0, 51, 52, 53], [31, 32, 33]), True), (([], []), False)], +) +def test_hybrid_write_request_finished_defers_free_only_with_blocks_to_send( + blocks, delay_free +): + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config + ) + sched.is_producer = True + req = _mk_request( + list(range(9)), params={"do_remote_decode": True, "transfer_id": "tx"} + ) + req.request_id = "req" + req.status = moriio_connector.RequestStatus.FINISHED_LENGTH_CAPPED + + assert sched.request_finished(req, blocks)[0] is delay_free + # The WRITE ACK, not the READ watchdog, releases the deferred blocks. + if delay_free: + assert sched._deferred_send_deadlines["req"][1] == "tx" + else: + assert "req" not in sched._deferred_send_deadlines + + +@pytest.mark.usefixtures("restore_moriio_role") +def test_hybrid_write_build_meta_saves_current_state_on_the_final_chunk(): + """Chunked prefill must not ship block ids cached at admission: align-mode + state slots move between chunks without appending a block.""" + config = _dspark_wrapped_attention_config() + config.select_transfer_block_ids = lambda blocks: tuple(blocks) + sched = moriio_connector.MoRIIOConnectorScheduler( + _gate_vllm_config(read_mode=False, dcp_size=2), "engine", config + ) + sched.is_producer = True + set_role(moriio_common.ROLE.PRODUCER) + params = { + "do_remote_decode": True, + "transfer_id": "tx", + "remote_engine_id": "decode", + "remote_block_ids": [], + "remote_host": "127.0.0.1", + "remote_handshake_port": 6000, + "remote_notify_port": 6001, + } + req = _mk_request(list(range(9)), params=params) + req.request_id = "req" + req.num_computed_tokens = 0 + sched._reqs_need_save["req"] = (req, [[10, 11, 12], [1]]) + sched._req_kv_params["req"] = dict(params) + + def _output(cached_req_ids, new_block_ids, num_scheduled): + return SimpleNamespace( + scheduled_cached_reqs=SimpleNamespace( + req_ids=cached_req_ids, new_block_ids=new_block_ids + ), + num_scheduled_tokens={"req": num_scheduled}, + new_block_ids_to_zero=None, + kv_connector_block_state=KVConnectorBlockState( + req_ids={"req"}, + resolve_block_ids=lambda _: ([0, 51, 52, 53], [31, 32, 33]), + boundary_state_offloads={}, + ), + ) + + # First chunk: not final, so nothing is saved yet. + meta = sched.build_connector_meta(_output([], [], 4)) + assert not meta.reqs_to_save + assert "req" in sched._reqs_need_pending_save + + # Final chunk: stale cached ids plus new ids are ignored in favour of + # the scheduler's current table, clamped to the prompt. + req.num_computed_tokens = 4 + meta = sched.build_connector_meta(_output(["req"], [([77], [78])], 5)) + assert meta.reqs_to_save["req"].local_block_ids == [[51], [31, 32]] + assert "req" not in sched._reqs_need_pending_save + + +@pytest.mark.parametrize( + "mode,deadlines,expected", + [ + (MoRIIOMode.WRITE, {"req": (0.0, "tx")}, True), + (MoRIIOMode.WRITE, {}, False), + (MoRIIOMode.READ, {"req": (0.0, "tx")}, False), + ], +) +def test_has_pending_push_work_tracks_write_deferred_frees(mode, deadlines, expected): + connector = _FakeConnector(_FakeScheduler(_deferred_send_deadlines=deadlines)) + connector.mode = mode + assert connector.has_pending_push_work() is expected + + +def test_has_pending_push_work_is_false_on_worker_role(): + connector = _FakeConnector(None) + connector.mode = MoRIIOMode.WRITE + assert connector.has_pending_push_work() is False + + +def test_loaded_kv_cache_groups_are_every_transfer_group(): + # The connector restores the request's own blocks, not a hash-addressed + # prefix, so a transfer-enabled Mamba group counts as loaded too. + connector = _FakeConnector(None) + connector._kv_cache_config = SimpleNamespace( + transfer_group_ids=(0, 2), prefix_cacheable_group_ids=(0,) + ) + assert connector.get_loaded_kv_cache_group_ids(SimpleNamespace()) == (0, 2) def test_split_block_groups_rejects_non_hybrid_use(): @@ -728,11 +950,15 @@ def test_session_build_rejects_per_layer_region_count_mismatch(): worker._get_built_session("prefill") -def test_register_kv_caches_rejects_hybrid_write_before_registration(): - worker = _FakeWorker(mode=MoRIIOMode.WRITE, _transfer_layer_names={"kda.0"}) +def test_register_kv_caches_rejects_invalid_hybrid_write_state_layout(): + worker = _FakeWorker( + mode=MoRIIOMode.WRITE, + _transfer_layer_names={"kda.0"}, + layer_to_spec={"kda.0": _mamba_spec(1)}, + ) worker._is_mamba_layer = lambda _layer_name: True - with pytest.raises(moriio_common.MoRIIOError, match="READ mode only"): + with pytest.raises(moriio_common.MoRIIOError, match="exactly two Mamba states"): worker.register_kv_caches({"kda.0": object()}) @@ -933,6 +1159,40 @@ def test_read_matches_only_requests_with_pending_remote_prefill(params, expected assert scheduler.get_num_new_matched_tokens(request, 0) == expected +@pytest.mark.parametrize( + "has_mamba,expected", [(False, (32, True)), (True, (31, True))] +) +def test_write_consumer_matches_full_prompt_unless_hybrid(has_mamba, expected): + """Without a mamba group (attention-only, or HMA disabled) prefill pushed + KV for all N tokens; with one it pushed h(N-1).""" + scheduler = _FakeScheduler( + is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=has_mamba + ) + request = SimpleNamespace( + num_prompt_tokens=32, kv_transfer_params={"do_remote_prefill": True} + ) + + assert scheduler.get_num_new_matched_tokens(request, 0) == expected + + +@pytest.mark.parametrize( + "is_producer,has_mamba,expected_prompt_len", + [(True, True, 4), (True, False, 5), (False, True, 5)], +) +def test_on_new_request_truncates_only_hybrid_producers( + is_producer, has_mamba, expected_prompt_len +): + """Attention-only prefill (including models run with HMA disabled) + computes all N prompt tokens; only a hybrid producer drops the last.""" + scheduler = _FakeScheduler(is_producer=is_producer, _has_mamba=has_mamba) + request = _mk_request([10, 11, 12, 13, 14], params={"do_remote_decode": True}) + + scheduler.on_new_request(request) + + assert request.num_prompt_tokens == expected_prompt_len + assert len(request.prompt_token_ids) == expected_prompt_len + + def test_read_rematch_after_alloc_keeps_queued_transfer(): """Clearing the flag blocks another external hit, not the queued READ.""" sched = _make_read_scheduler() @@ -981,12 +1241,37 @@ def test_get_num_new_matched_tokens_write_plain_keeps_all_tokens(): assert is_async is True +@pytest.mark.parametrize("params", [None, {}, {"do_remote_prefill": False}]) +def test_write_does_not_wait_again_after_receive_or_for_local_requests(params): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True) + req = _mk_request(list(range(10))) + req.kv_transfer_params = params + assert sched.get_num_new_matched_tokens(req, 0) == (0, False) + + +def test_write_never_declares_an_empty_async_load(): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=False) + req = _mk_request(list(range(10)), params={"do_remote_prefill": True}) + assert sched.get_num_new_matched_tokens(req, 10) == (0, False) + + +@pytest.mark.parametrize( + "prompt_len,computed,expected", [(16, 0, 15), (17, 16, 0), (1, 0, 0)] +) +def test_hybrid_write_advertises_only_the_state_computed_by_prefill( + prompt_len, computed, expected +): + sched = _FakeScheduler(is_producer=False, mode=MoRIIOMode.WRITE, _has_mamba=True) + req = _mk_request(list(range(prompt_len)), params={"do_remote_prefill": True}) + assert sched.get_num_new_matched_tokens(req, computed) == (expected, expected > 0) + + @pytest.mark.parametrize( ("mode", "num_computed_tokens", "expected", "is_async"), [ (MoRIIOMode.READ, 0, 9, False), (MoRIIOMode.READ, 10, 0, False), - (MoRIIOMode.WRITE, 2, 8, True), + (MoRIIOMode.WRITE, 2, 7, True), ], ) def test_get_num_new_matched_tokens_supports_embeds_only_prompts( diff --git a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py index 5ec6fb5755..f0559f259d 100644 --- a/tests/v1/kv_connector/unit/test_moriio_kv_layout.py +++ b/tests/v1/kv_connector/unit/test_moriio_kv_layout.py @@ -7,6 +7,7 @@ import threading from collections import OrderedDict from types import SimpleNamespace from typing import Any +from unittest.mock import Mock import pytest import torch @@ -106,7 +107,7 @@ def _write_task(layer_name: str, transfer_id: str = "xfer") -> Any: request_id="req", transfer_id=transfer_id, dst_engine_id="remote-engine", - local_block_ids=[1, 3], + local_block_ids=[[1, 3]], remote_block_ids_hint=None, layer_name=layer_name, event=None, @@ -345,6 +346,8 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry(): calls: list[str] = [] class FakeWorker: + _has_mamba: bool + layer_to_group: dict[str, int] kv_caches: dict[str, torch.Tensor] layer_name_to_local_kv_cache_metadata: dict[str, list[Any]] @@ -364,11 +367,13 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry(): return ([call_id], [call_id + 10], [call_id + 20]) fake_worker = FakeWorker() + fake_worker._has_mamba = False + fake_worker.layer_to_group = dict.fromkeys(kv_caches, 0) fake_worker.kv_caches = kv_caches fake_worker.layer_name_to_local_kv_cache_metadata = {name: [] for name in kv_caches} writer = MoRIIOWriter.__new__(MoRIIOWriter) writer._worker_ref = lambda: fake_worker - request_info = RemoteAllocInfo(block_ids=[4, 5]) + request_info = RemoteAllocInfo(block_ids=[[4, 5]]) remote_meta = _remote_meta() dense0_plan = writer._prepare_transfer_plan( @@ -395,7 +400,7 @@ def test_write_transfer_plan_caches_offsets_per_layer_geometry(): def test_write_scheduler_deduplicates_layers_and_seals_expected_count(): - request_info = RemoteAllocInfo(block_ids=[4, 5]) + request_info = RemoteAllocInfo(block_ids=[[4, 5]]) wrapper = _wrapper_for_messages() wrapper.done_remote_allocate_req_dict["xfer"] = request_info writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) @@ -411,6 +416,80 @@ def test_write_scheduler_deduplicates_layers_and_seals_expected_count(): assert writer._sealed_writes["xfer"] == 2 +def test_hybrid_write_attention_uses_its_own_group_not_the_first_group(): + worker = SimpleNamespace( + _has_mamba=True, + kv_caches={"draft": torch.empty((8, 1, 4, 3))}, + layer_to_group={"draft": 2}, + _region_session_indices=lambda _: [5], + _compute_block_transfer_offsets=Mock(return_value=([1], [2], [3])), + ) + writer = make_moriio_writer(worker) + task = _write_task("draft") + task.local_block_ids = [[10], [20], [30, 31]] + remote = RemoteAllocInfo(block_ids=[[40], [50], [60, 61, 62]]) + meta = _remote_meta() + plan = writer._prepare_transfer_plan(task, remote, meta) + worker._compute_block_transfer_offsets.assert_called_once_with( + "draft", [30, 31], [60, 61], meta, remote_engine_id=task.dst_engine_id + ) + assert plan.sess_idx == 5 + + +def test_hybrid_write_posts_both_regions_before_layer_completion(): + from unittest.mock import Mock + + from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_connector import ( + MoRIIOConnectorWorker, + ) + + wrapper = _wrapper_for_messages() + info = RemoteAllocInfo(block_ids=[[], [7], [9]]) + wrapper.done_remote_allocate_req_dict["xfer"] = info + wrapper.write_remote_data = Mock(side_effect=["conv-status", "ssm-status"]) + worker = SimpleNamespace( + moriio_wrapper=wrapper, + tp_rank=0, + get_engine_name_with_dp=lambda name, rank: name, + _get_built_session=lambda name: (["conv-session", "ssm-session"], None), + _is_mamba_layer=lambda name: True, + _mamba_payload_index_by_layer={"kda": 2}, + _region_session_indices=lambda name: [0, 1], + _compute_mamba_transfer_offsets=Mock( + return_value=([10, 20, 30, 40], [50, 60, 70, 80], [2, 4, 8, 16], 3) + ), + ) + worker._mamba_blocks_for_layer = lambda name, blocks: ( + MoRIIOConnectorWorker._mamba_blocks_for_layer(worker, name, blocks) + ) + writer = make_moriio_writer(worker) + task = _write_task("kda") + task.local_block_ids = [[], [1], [3]] + task.remote_tp_size = 2 + task.event = Mock() + writer._execute_write_task(task) + + task.event.synchronize.assert_called_once() + worker._compute_mamba_transfer_offsets.assert_called_once_with( + "kda", [3], [9], remote_tp_size=2 + ) + assert wrapper.write_remote_data.call_args_list[0].args == ( + [2, 4, 8], + [10, 20, 30], + [50, 60, 70], + "conv-session", + ) + assert wrapper.write_remote_data.call_args_list[1].args == ( + [16], + [40], + [80], + "ssm-session", + ) + assert info.transfer_statuses == ["conv-status", "ssm-status"] + assert info.writes_done == 1 + assert not info.completion_notified + + def test_write_completion_notifies_once_after_all_sealed_writes_finish(): class FakeWrapper: def __init__(self): @@ -445,7 +524,7 @@ def test_write_completion_notifies_once_after_all_sealed_writes_finish(): ) wrapper = FakeWrapper() - request_info = RemoteAllocInfo(block_ids=[4, 5], writes_expected=2) + request_info = RemoteAllocInfo(block_ids=[[4, 5]], writes_expected=2) request_info.transfer_statuses.extend(["status-a", "status-b"]) request_info.completion_request_id = "req" # Final notify port is resolved in _execute_write_task; finalize reads it @@ -487,6 +566,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain(): wrapper = MoRIIOWrapper.__new__(MoRIIOWrapper) wrapper.lock = threading.Lock() wrapper._transfer_timeout = 1 + wrapper.moriio_engine = object() global_status = FakeStatus() scoped_status = FakeStatus() wrapper.transfer_status = [global_status] @@ -500,7 +580,7 @@ def test_moriio_wrapper_waits_scoped_statuses_without_global_drain(): def test_write_failure_marks_terminal_and_clears_scheduled_state(): wrapper = _wrapper_for_messages() - wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5]) + wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[[4, 5]]) writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) writer._scheduled_writes["xfer"] = 2 writer._scheduled_layers["xfer"] = {"dense0", "indexer"} @@ -518,7 +598,7 @@ def test_write_failure_marks_terminal_and_clears_scheduled_state(): def test_schedule_write_rejects_terminal_transfer_without_recreating_state(): wrapper = _wrapper_for_messages() - wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[4, 5]) + wrapper.done_remote_allocate_req_dict["xfer"] = RemoteAllocInfo(block_ids=[[4, 5]]) writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) writer._scheduled_writes["xfer"] = 1 writer._scheduled_layers["xfer"] = {"dense0"} @@ -610,7 +690,7 @@ def test_moriio_wrapper_routes_valid_messages(role, payload, expected): if expected == "remote_blocks": request_info = wrapper.done_remote_allocate_req_dict["xfer"] - assert request_info.block_ids == [4, 5] + assert request_info.block_ids == [[4, 5]] assert request_info.decode_dp_rank == 3 elif expected == "write_done": assert wrapper.done_write_cache_req_ids == ["xfer"] diff --git a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py index 2a3092ad5a..8fa59e5fb5 100644 --- a/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py +++ b/tests/v1/kv_connector/unit/test_moriio_transfer_completion.py @@ -11,6 +11,8 @@ import importlib import importlib.util import threading import time +from types import SimpleNamespace +from unittest.mock import Mock import pytest @@ -127,3 +129,62 @@ def test_wait_blocks_until_terminal(): assert time.monotonic() - start >= 0.15 finally: timer.cancel() + + +@pytest.mark.parametrize("timeout,expected_ms", [(0.05, 50), (0.0001, 1)]) +def test_native_wait_receives_all_statuses_and_a_positive_deadline( + timeout, expected_ms +): + statuses = [_pending(), _ok()] + engine = SimpleNamespace( + wait_all=Mock(return_value=moriio_engine.StatusCode.SUCCESS) + ) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=timeout) + wrapper.waiting_for_transfer_complete(statuses) + engine.wait_all.assert_called_once_with(statuses, timeout_ms=expected_ms) + + +@pytest.mark.parametrize( + "result,statuses,message", + [ + ("ERR_RDMA_OP", [_pending(), _bad()], "SQ full"), + ("IN_PROGRESS", [_pending(), _ok()], "timed out"), + ("ERR_INVALID_ARGS", [_pending()], "ERR_INVALID_ARGS"), + ], +) +def test_native_wait_reports_failure_or_timeout(result, statuses, message): + engine = SimpleNamespace( + wait_all=Mock(return_value=getattr(moriio_engine.StatusCode, result)) + ) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + with pytest.raises(TransferError, match=message): + wrapper.waiting_for_transfer_complete(statuses) + + +def test_wait_falls_back_without_native_status_enum(monkeypatch): + monkeypatch.setattr(moriio_engine, "StatusCode", None) + engine = SimpleNamespace(wait_all=Mock()) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + with pytest.raises(TransferError, match="SQ full"): + wrapper.waiting_for_transfer_complete([_bad(), _ok()]) + engine.wait_all.assert_not_called() + + +def test_nonblocking_poll_does_not_call_native_wait(): + engine = SimpleNamespace(wait_all=Mock()) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + assert wrapper.poll_transfer_batch([_pending()]) is TransferBatchState.PENDING + wrapper.waiting_for_transfer_complete([]) + engine.wait_all.assert_not_called() + + +def test_native_wait_drains_only_the_unscoped_status_snapshot(): + statuses = [_pending(), _ok()] + engine = SimpleNamespace( + wait_all=Mock(return_value=moriio_engine.StatusCode.SUCCESS) + ) + wrapper = MoRIIOWrapper(moriio_engine=engine, transfer_timeout=0.05) + wrapper.transfer_status.extend(statuses) + wrapper.waiting_for_transfer_complete() + engine.wait_all.assert_called_once_with(statuses, timeout_ms=50) + assert wrapper.transfer_status == [] diff --git a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py index 57b6893869..ee04c2cabe 100644 --- a/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py +++ b/tests/v1/kv_connector/unit/test_moriio_write_done_routing.py @@ -14,6 +14,8 @@ import pytest from vllm.distributed.kv_transfer.kv_connector.v1.moriio import moriio_engine from vllm.distributed.kv_transfer.kv_connector.v1.moriio.moriio_common import ( ROLE, + LayerTransferPlan, + MoRIIOError, MoRIIOMode, RemoteAllocInfo, ReqMeta, @@ -96,7 +98,7 @@ def _write_task( request_id=request_id, transfer_id=transfer_id, dst_engine_id="remote-engine", - local_block_ids=[1, 3], + local_block_ids=[[1, 3]], remote_block_ids_hint=None, layer_name="layer0", event=None, @@ -249,6 +251,7 @@ def test_finalize_single_pod_port_uses_global_rank(): def test_write_blocks_for_req_forwards_routing_to_schedule_write(): worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = False captured: dict[str, Any] = {} def _capture_schedule_write_blocks(**kwargs: Any) -> None: @@ -266,13 +269,14 @@ def test_write_blocks_for_req_forwards_routing_to_schedule_write(): assert captured["multi_pod_hosts"] == ["10.0.0.1", "10.0.0.2"] assert captured["remote_dp_size_local"] == 8 assert captured["remote_ip"] == "10.0.0.1" - assert captured["local_block_ids"] == [1] - assert captured["remote_block_ids"] == [2] + assert captured["local_block_ids"] == [[1]] + assert captured["remote_block_ids"] == [[2]] assert not hasattr(worker, "multi_pod_hosts") def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty(): worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = False captured: dict[str, Any] = {} def _capture_schedule_write_blocks(**kwargs: Any) -> None: @@ -290,8 +294,8 @@ def test_write_blocks_for_req_falls_back_when_multi_pod_hosts_empty(): assert captured["multi_pod_hosts"] == ["10.0.0.9"] assert captured["remote_dp_size_local"] == 4 - assert captured["local_block_ids"] == [1] - assert captured["remote_block_ids"] == [2] + assert captured["local_block_ids"] == [[1]] + assert captured["remote_block_ids"] == [[2]] def test_deferred_write_timeout_fails_transfer_on_every_decode_rank(): @@ -373,3 +377,73 @@ def test_write_failed_is_reported_as_failed_recving_on_owning_rank(): assert results.failed_recving == {"req-1"} assert wrapper.pop_finished_write_req_ids() == set() assert list(worker._unmatched_write_failures) == ["tOtherRank"] + + +def test_hybrid_write_rejects_unequal_tp_before_scheduling_any_layer(): + from unittest.mock import Mock + + worker = MoRIIOConnectorWorker.__new__(MoRIIOConnectorWorker) + worker._has_mamba = True + worker.world_size = 2 + worker.schedule_write_blocks = Mock() + + with pytest.raises(NotImplementedError, match="heterogeneous-TP"): + worker._write_blocks_for_req("req", _req_meta(), "attention", kv_layer=None) + + worker.schedule_write_blocks.assert_not_called() + + +@pytest.mark.parametrize( + "notify,expected", + [ + # Plain WRITE peers still send one flat attention list. + ([7, 8], [[7, 8]]), + # Hybrid WRITE sends one list per transfer group, kept as-is even + # when the first group is empty (full attention hit, state only). + ([[7], [9]], [[7], [9]]), + ([[], [9]], [[], [9]]), + ], +) +def test_remote_blocks_message_normalizes_to_grouped_block_ids(notify, expected): + wrapper = _real_wrapper() + with patch.object(moriio_engine, "get_role", return_value=ROLE.PRODUCER): + wrapper._handle_remote_blocks_message( + {"transfer_id": "tA", "block_notify_list": notify, "decode_rank": 2} + ) + info = wrapper.done_remote_allocate_req_dict["tA"] + assert info.block_ids == expected + assert info.decode_dp_rank == 2 + + +def test_empty_layer_plan_posts_no_rdma_write(): + from unittest.mock import Mock + + wrapper = SimpleNamespace(write_remote_data=Mock(), write_remote_data_single=Mock()) + writer = make_moriio_writer(SimpleNamespace(moriio_wrapper=wrapper)) + plan = LayerTransferPlan( + request_id="req", + transfer_id="tA", + layer_name="layer0", + sess_idx=0, + transfer_local_offsets=[], + transfer_remote_offsets=[], + transfer_sizes=[], + ) + + assert writer._do_layer_write(plan, ["session"]) == [] + wrapper.write_remote_data.assert_not_called() + wrapper.write_remote_data_single.assert_not_called() + + +@pytest.mark.parametrize("local_state", [[], [1, 2]]) +def test_hybrid_write_requires_exactly_one_running_state_slot(local_state): + worker = SimpleNamespace( + _mamba_blocks_for_layer=lambda layer_name, block_ids: block_ids[1], + ) + writer = make_moriio_writer(worker) + task = _write_task("tA", "10.0.0.1") + task.layer_name = "kda" + task.local_block_ids = [[], local_state] + + with pytest.raises(MoRIIOError, match="one running-state slot"): + writer._prepare_mamba_transfer_plans(task, RemoteAllocInfo(block_ids=[[], [9]])) diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py index e82d066818..04fd38be2f 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py @@ -63,12 +63,13 @@ class WriteTask: request_id: ReqId transfer_id: TransferId dst_engine_id: str - local_block_ids: list[int] - remote_block_ids_hint: list[int] | None + local_block_ids: BlockIds + remote_block_ids_hint: BlockIds | None layer_name: str event: torch.cuda.Event remote_notify_port: int remote_ip: str + remote_tp_size: int | None = None multi_pod_hosts: list[str] = field(default_factory=list) remote_dp_size_local: int = 0 remote_dp_size: int = 1 @@ -94,7 +95,7 @@ class LayerTransferPlan: class RemoteAllocInfo: """Information about remote block allocation.""" - block_ids: list[int] + block_ids: BlockIds | None writes_done: int = 0 writes_expected: int | None = None decode_dp_rank: int = 0 diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py index fa75eaee3d..9760d8bef1 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_connector.py @@ -338,7 +338,10 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): block_ids: BlockIds, ) -> tuple[bool, dict[str, Any] | None]: assert self.connector_scheduler is not None - if self.connector_scheduler._has_mamba: + if ( + self.connector_scheduler._has_mamba + and self.connector_scheduler.mode == MoRIIOMode.READ + ): attn_block_ids, mamba_block_groups = ( self.connector_scheduler.split_block_groups(block_ids) ) @@ -420,6 +423,16 @@ class MoRIIOConnector(KVConnectorBase_V1, SupportsHMA): ) self.connector_worker.wait_for_save(self._connector_metadata) + def has_pending_push_work(self) -> bool: + return ( + self.mode == MoRIIOMode.WRITE + and self.connector_scheduler is not None + and bool(self.connector_scheduler._deferred_send_deadlines) + ) + + def get_loaded_kv_cache_group_ids(self, request: "Request") -> tuple[int, ...]: + return self._kv_cache_config.transfer_group_ids + def shutdown(self): if self.connector_worker is not None: self.connector_worker.shutdown() @@ -463,11 +476,10 @@ def _validate_hybrid_speculation(vllm_config: VllmConfig) -> None: speculative_config = vllm_config.speculative_config if speculative_config is None: return - # DSpark is allowlisted because it is the only hybrid speculative method - # validated end to end with MoRIIO READ. + # Other hybrid speculative methods need their own transfer validation. if not speculative_config.use_dspark(): raise MoRIIOError( - "MoRIIO hybrid READ supports DSpark speculative decoding only, got " + "MoRIIO hybrid transfer supports DSpark speculative decoding only, got " f"method={speculative_config.method!r}" ) @@ -480,14 +492,14 @@ def _validate_mamba_specs(specs: Collection[MambaSpec]) -> MambaSpec | None: for spec in specs: if len(spec.shapes) != 2 or len(spec.dtypes) != 2: raise MoRIIOError( - "MoRIIO hybrid READ supports exactly two Mamba states " + "MoRIIO hybrid transfer supports exactly two Mamba states " "(conv and SSM), got " f"{len(spec.shapes)} shapes and {len(spec.dtypes)} dtypes" ) reference = specs[0] if any(spec != reference for spec in specs[1:]): raise MoRIIOError( - "MoRIIO hybrid READ requires all transferable Mamba groups " + "MoRIIO hybrid transfer requires all transferable Mamba groups " "to use the same cache spec" ) return reference @@ -515,7 +527,9 @@ class MoRIIOConnectorScheduler: self._has_mamba = bool(self._mamba_group_ids) mamba_spec = None if self._has_mamba: - if len(self._attn_group_ids) != 1: + if not self._attn_group_ids: + raise MoRIIOError("MoRIIO hybrid transfer requires an attention group") + if self.mode == MoRIIOMode.READ and len(self._attn_group_ids) != 1: raise MoRIIOError( "MoRIIO hybrid READ requires exactly one transferable " "attention group, got " @@ -529,11 +543,6 @@ class MoRIIOConnectorScheduler: ] mamba_spec = _validate_mamba_specs(mamba_specs) _validate_hybrid_speculation(vllm_config) - if self.mode != MoRIIOMode.READ: - raise MoRIIOError( - "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " - "mode only; set kv_connector_extra_config.read_mode=true" - ) self._num_ssm_scratch_blocks = ( mamba_spec.num_speculative_blocks if mamba_spec is not None else 0 ) @@ -581,11 +590,13 @@ class MoRIIOConnectorScheduler: f"cache group spec(s): {sorted(set(unsupported))}. Pass " "--disable-hybrid-kv-cache-manager to run without HMA." ) - if self.mode == MoRIIOMode.WRITE: - # TODO(simondanielsson): support HMA in WRITE mode + if self.mode == MoRIIOMode.WRITE and any( + isinstance(g.kv_cache_spec, SlidingWindowSpec) + for g in kv_cache_config.transfer_groups + ): raise NotImplementedError( - "MoRIIO WRITE mode does not support hybrid KV cache groups " - "(sliding-window attention). Use READ mode, or pass " + "MoRIIO WRITE mode does not support sliding-window " + "cache groups. Use READ mode, or pass " "--disable-hybrid-kv-cache-manager." ) @@ -784,10 +795,11 @@ class MoRIIOConnectorScheduler: if not params.get("do_remote_prefill"): return 0, False if self.mode == MoRIIOMode.WRITE: - # MoriiO in write mode, no remote prefill. Hybrid models never get - # here: register_kv_caches refuses WRITE mode for them, so there is - # no recurrent-state accounting to do. - return num_external_tokens, True + if self._has_mamba: + # P pushes h(N-1). Do not publish it under the N-token hash + # when the scheduler caches a completed asynchronous load. + num_external_tokens = max(num_external_tokens - 1, 0) + return num_external_tokens, num_external_tokens > 0 # READ mode always recomputes the last token locally on the decoder. # @@ -1178,6 +1190,16 @@ class MoRIIOConnectorScheduler: f"remote_notify_port={remote_notify_port!r})" ) + block_notify_list: list[int] | list[list[int]] + if self._has_mamba: + block_notify_list = list( + self._clamp_to_prompt_blocks( + request, blocks.get_block_ids() + ) + ) + else: + block_notify_list = blocks.get_block_ids()[0] + # Wide-EP multi-pod: a pod binds notify sockets only for # its LOCAL ranks, so the port offset must use the per-pod # local rank (% dp_local), not the global rank. Single-pod @@ -1204,7 +1226,7 @@ class MoRIIOConnectorScheduler: self.send_notify_block( req_id=request.request_id, transfer_id=request.kv_transfer_params["transfer_id"], - block_notify_list=blocks.get_block_ids()[0], + block_notify_list=block_notify_list, host=_notify_host, port=target_port, ) @@ -1227,6 +1249,26 @@ class MoRIIOConnectorScheduler: sliding-window groups are left unchanged. """ + if self._has_mamba: + # WRITE preserves transfer-group order, including a separate drafter + # attention group. READ retains its attention-first wire format. + transfer_blocks = self.kv_cache_config.select_transfer_block_ids(block_ids) + dcp_size = self.vllm_config.parallel_config.decode_context_parallel_size + clamped_hybrid: list[list[int]] = [] + for group, blocks in zip( + self.kv_cache_config.transfer_groups, transfer_blocks, strict=True + ): + spec = group.kv_cache_spec + if isinstance(spec, MambaSpec): + clamped_hybrid.append( + clip_ssm_state_blocks(list(blocks), spec.num_speculative_blocks) + ) + else: + block_size = spec.block_size * (dcp_size if spec.dcp_sharded else 1) + clamped_hybrid.append( + list(blocks[: cdiv(req.num_prompt_tokens, block_size)]) + ) + return clamped_hybrid clamped = list(block_ids) changed = False for g, group_blocks in enumerate(block_ids): @@ -1252,6 +1294,19 @@ class MoRIIOConnectorScheduler: num_scheduled = scheduler_output.num_scheduled_tokens.get(req.request_id, 0) return req.num_computed_tokens + num_scheduled >= req.num_prompt_tokens + def _get_write_blocks( + self, req: "Request", block_ids: BlockIds, scheduler_output: SchedulerOutput + ) -> BlockIds: + if self._has_mamba: + # Align-mode state slots can move without appending a new block. + # Resolve the authoritative table at the final prefill step. + block_state = scheduler_output.kv_connector_block_state + assert block_state is not None + current_blocks = block_state.get_block_ids(req.request_id) + assert current_blocks is not None + block_ids = current_blocks + return self._clamp_to_prompt_blocks(req, block_ids) + def build_connector_meta( self, scheduler_output: SchedulerOutput, @@ -1274,10 +1329,10 @@ class MoRIIOConnectorScheduler: # reserved speculative-lookahead slack), so new_block_ids is None; # still fall through to the final-chunk check below. new_block_ids = scheduler_output.scheduled_cached_reqs.new_block_ids[i] - if new_block_ids is not None: + if new_block_ids is not None and not self._has_mamba: existing_blocks = [ - existing_blocks[g] + new_block_ids[g] - for g in range(len(new_block_ids)) + old + new + for old, new in zip(existing_blocks, new_block_ids, strict=True) ] self._reqs_need_pending_save[req_id] = (req, existing_blocks) if not self._is_final_prefill_chunk(req, scheduler_output): @@ -1287,7 +1342,9 @@ class MoRIIOConnectorScheduler: kv_params = self._req_kv_params.pop( req_id, req.kv_transfer_params or {} ) - save_block_ids = self._clamp_to_prompt_blocks(req, existing_blocks) + save_block_ids = self._get_write_blocks( + req, existing_blocks, scheduler_output + ) meta.add_new_req( request_id=req_id, local_block_ids=save_block_ids, @@ -1328,7 +1385,7 @@ class MoRIIOConnectorScheduler: if not self._is_final_prefill_chunk(req, scheduler_output): self._reqs_need_pending_save[req_id] = (req, block_ids) continue - save_block_ids = self._clamp_to_prompt_blocks(req, block_ids) + save_block_ids = self._get_write_blocks(req, block_ids, scheduler_output) meta.add_new_req( request_id=req_id, local_block_ids=save_block_ids, @@ -1425,7 +1482,7 @@ class MoRIIOConnectorScheduler: """Once a request is finished, determine whether request blocks should be freed now or will be sent asynchronously and freed later. - ``block_ids`` normally contains all cache groups. Hybrid callers from + ``block_ids`` normally contains all cache groups. Hybrid READ callers from ``request_finished_all_groups`` pass a flat attention list and the recurrent-state groups separately in ``mamba_block_groups``. """ @@ -1494,6 +1551,10 @@ class MoRIIOConnectorScheduler: list(cast(list[int], block_ids)), *[list(group) for group in mamba_block_groups], ] + elif self._has_mamba and self.mode == MoRIIOMode.WRITE: + computed_block_ids = self._clamp_to_prompt_blocks( + request, cast(BlockIds, block_ids) + ) elif self._has_mamba: raise MoRIIOError("Hybrid request completion requires all KV cache groups") else: @@ -1660,7 +1721,7 @@ class MoRIIOConnectorWorker: _, mamba_transfer_group_ids = _split_kv_cache_group_kinds(kv_cache_config) self._num_mamba_transfer_groups = len(mamba_transfer_group_ids) self._mamba_payload_index_by_layer = { - layer_name: payload_index + layer_name: group_id if self.mode == MoRIIOMode.WRITE else payload_index for payload_index, group_id in enumerate(mamba_transfer_group_ids, start=1) for layer_name in kv_cache_config.transfer_groups[group_id].layer_names } @@ -1884,8 +1945,8 @@ class MoRIIOConnectorWorker: request_id: ReqId, transfer_id: TransferId, dst_engine_id: str, - local_block_ids: list[int], - remote_block_ids: list[int] | None, + local_block_ids: BlockIds, + remote_block_ids: BlockIds | None, layer_name: str, kv_layer: torch.Tensor, remote_notify_port: int, @@ -1893,6 +1954,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts: list[str], remote_dp_size_local: int, remote_dp_size: int, + remote_tp_size: int | None = None, ) -> None: """Schedule a block write operation. @@ -1909,6 +1971,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts: List of pod IPs for multi-pod Wide-EP remote_dp_size_local: Per-pod DP size for multi-pod remote_dp_size: Global DP size of the decode instance + remote_tp_size: Tensor parallel size advertised by the peer """ # synchronization to prevent dirty reads between @@ -1933,6 +1996,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts=multi_pod_hosts, remote_dp_size_local=remote_dp_size_local, remote_dp_size=remote_dp_size, + remote_tp_size=remote_tp_size, ) self._writer.schedule_write(task) @@ -2359,15 +2423,6 @@ class MoRIIOConnectorWorker: self._has_mamba = any( self._is_mamba_layer(layer_name) for layer_name in kv_caches ) - if self._has_mamba and self.mode == MoRIIOMode.WRITE: - # WRITE needs a distinct conv/ssm transfer plan. Reject before - # allocating or registering any regions rather than silently - # serving a zero recurrent state. - raise MoRIIOError( - "MoRIIO hybrid (mamba/KDA) transfer is implemented for READ " - "mode only; set kv_connector_extra_config.read_mode=true " - "for hybrid models" - ) if self._has_mamba: mamba_specs = [ cast(MambaSpec, self.layer_to_spec[layer_name]) @@ -3292,6 +3347,9 @@ class MoRIIOConnectorWorker: ) def _write_blocks_for_req(self, req_id: ReqId, meta: ReqMeta, layer_name, kv_layer): + if self._has_mamba: + # Reject before even an attention-layer write can be queued. + self._mamba_tp_ratio(int(meta.tp_size)) # Compute per-request values to pass through the task (no shared state). # This avoids race conditions when concurrent requests target different # decode pods - each WriteTask carries its own routing info. @@ -3303,16 +3361,12 @@ class MoRIIOConnectorWorker: if meta.remote_dp_size_local else int(meta.remote_dp_size) ) - # WRITE does not support HMA, so unwrap blocks into flat lists. - # remote_block_ids can itself be empty so need to be careful when unwrapping. - local_block_ids = meta.local_block_ids[0] - remote_block_ids = meta.remote_block_ids[0] if meta.remote_block_ids else [] self.schedule_write_blocks( request_id=req_id, transfer_id=meta.transfer_id, dst_engine_id=meta.remote_engine_id, - local_block_ids=local_block_ids, - remote_block_ids=remote_block_ids, + local_block_ids=meta.local_block_ids, + remote_block_ids=meta.remote_block_ids, layer_name=layer_name, kv_layer=kv_layer, remote_notify_port=meta.remote_notify_port, @@ -3320,6 +3374,7 @@ class MoRIIOConnectorWorker: multi_pod_hosts=hosts, remote_dp_size_local=dp_local, remote_dp_size=int(meta.remote_dp_size), + remote_tp_size=int(meta.tp_size), ) def merge_contiguous_blocks( diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py index e33a1524e1..3ffd6d10f9 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_engine.py @@ -3,6 +3,7 @@ import threading import time from collections import OrderedDict, defaultdict +from math import ceil from queue import Empty, Queue from typing import TYPE_CHECKING, Any from weakref import ref as weakref_ref @@ -61,6 +62,11 @@ try: except ImportError: logger.error("MoRIIO is not available") +try: + from mori.io import StatusCode +except ImportError: + StatusCode = None + """Write task execution logic for MoRIIO connector.""" @@ -372,11 +378,15 @@ class MoRIIOWriter: task.dst_engine_id ) - # Prepare transfer plan - plan = self._prepare_transfer_plan(task, request_info, remote_moriio_meta) - - # Execute transfer - transfer_statuses = self._do_layer_write(plan, sessions) + if self.worker._is_mamba_layer(task.layer_name): + plans = self._prepare_mamba_transfer_plans(task, request_info) + else: + plans = [ + self._prepare_transfer_plan(task, request_info, remote_moriio_meta) + ] + transfer_statuses = [] + for plan in plans: + transfer_statuses.extend(self._do_layer_write(plan, sessions)) with self._write_state_lock: request_info.transfer_statuses.extend(transfer_statuses) @@ -404,10 +414,17 @@ class MoRIIOWriter: key = (task.layer_name, *_get_write_geometry_key(layer_cache)) offsets = request_info.transfer_offsets.get(key) if offsets is None: + group = self.worker.layer_to_group[task.layer_name] + assert request_info.block_ids is not None + local_blocks = task.local_block_ids[group] + remote_blocks = request_info.block_ids[group] + if self.worker._has_mamba: + # D can reserve a tail block for the token recomputed locally. + remote_blocks = remote_blocks[: len(local_blocks)] offsets = self.worker._compute_block_transfer_offsets( task.layer_name, - task.local_block_ids, - request_info.block_ids, + local_blocks, + remote_blocks, remote_moriio_meta, remote_engine_id=task.dst_engine_id, ) @@ -430,6 +447,44 @@ class MoRIIOWriter: use_batch=True, ) + def _prepare_mamba_transfer_plans( + self, task: WriteTask, request_info: RemoteAllocInfo + ) -> list[LayerTransferPlan]: + """Plan both registered regions as one scheduled layer write.""" + assert request_info.block_ids is not None + local_slots = self.worker._mamba_blocks_for_layer( + task.layer_name, task.local_block_ids + ) + remote_slots = self.worker._mamba_blocks_for_layer( + task.layer_name, request_info.block_ids + ) + if len(local_slots) != 1 or len(remote_slots) != 1: + raise MoRIIOError( + "Hybrid WRITE requires one running-state slot per Mamba group" + ) + local, remote, sizes, n_conv = self.worker._compute_mamba_transfer_offsets( + task.layer_name, + local_slots, + remote_slots, + remote_tp_size=task.remote_tp_size, + ) + conv_session, ssm_session = self.worker._region_session_indices(task.layer_name) + return [ + LayerTransferPlan( + request_id=task.request_id, + transfer_id=task.transfer_id, + layer_name=task.layer_name, + sess_idx=session, + transfer_local_offsets=local[region], + transfer_remote_offsets=remote[region], + transfer_sizes=sizes[region], + ) + for session, region in ( + (conv_session, slice(0, n_conv)), + (ssm_session, slice(n_conv, None)), + ) + ] + def _do_layer_write(self, plan: LayerTransferPlan, sessions: list) -> list[Any]: """Perform the actual layer write. @@ -438,6 +493,8 @@ class MoRIIOWriter: sessions: List of transfer sessions """ + if not plan.transfer_sizes: + return [] if plan.use_batch: return [ self.worker.moriio_wrapper.write_remote_data( @@ -721,7 +778,28 @@ class MoRIIOWrapper: if not transfers_to_wait: return - self._poll_transfers_until_done(transfers_to_wait) + wait_all = getattr(self.moriio_engine, "wait_all", None) + if wait_all is None or StatusCode is None: + self._poll_transfers_until_done(transfers_to_wait) + return + + # Zero means a nonblocking progress call in Mori, not an immediate + # deadline. Round up so even sub-millisecond waits remain blocking. + timeout_ms = ceil(self._transfer_timeout * 1000) + result = wait_all(transfers_to_wait, timeout_ms=timeout_ms) + if result != StatusCode.SUCCESS: + details = [ + f"{status.Message()} (code={status.Code()})" + for status in transfers_to_wait + if status.Failed() + ] + if not details and result in (StatusCode.INIT, StatusCode.IN_PROGRESS): + details.append( + f"RDMA transfer timed out after {self._transfer_timeout:g}s" + ) + raise TransferError( + f"MoRIIO batch wait failed ({result}): " + "; ".join(details) + ) def _poll_transfers_until_done(self, transfers_to_wait: list[Any]) -> None: """Fallback for mori builds without the batched wait. @@ -858,7 +936,12 @@ class MoRIIOWrapper: ) return self.done_remote_allocate_req_dict[transfer_id] = RemoteAllocInfo( - block_ids=block_notify_list, decode_dp_rank=decode_dp_rank + block_ids=( + [block_notify_list] + if isinstance(block_notify_list[0], int) + else block_notify_list + ), + decode_dp_rank=decode_dp_rank, ) def _handle_write_done_message(self, data: dict): From 39983a5e7a5a575a7829f5d5ada50a33d9337ec1 Mon Sep 17 00:00:00 2001 From: Jee Jee Li Date: Sat, 17 Oct 2026 13:17:55 +0800 Subject: [PATCH 44/49] [Kernel] Add FP32 router GEMM support for Step-5 Preview and Muse Spark Open (#62126) Signed-off-by: Jee Jee Li --- csrc/fp32_router_gemm.cu | 29 ++++++++++++++++++- csrc/fp32_router_gemm_entry.cu | 12 ++++++-- tests/kernels/test_fp32_router_gemm.py | 23 +++++++++------ .../layers/fused_moe/router/gate_linear.py | 8 ++++- 4 files changed, 59 insertions(+), 13 deletions(-) diff --git a/csrc/fp32_router_gemm.cu b/csrc/fp32_router_gemm.cu index 832874f914..d5784ff0e6 100644 --- a/csrc/fp32_router_gemm.cu +++ b/csrc/fp32_router_gemm.cu @@ -279,6 +279,27 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, launchFp32RouterGemm( output, mat_a, mat_b, stream); } + } else if constexpr (std::is_same_v && + kNumExperts == 256 && kHiddenDim == 8192) { + // MuseSpark1x: the 8MB weight makes this shape FFMA-bound past M~10, so + // the wider block wins at every M and token groups clear the M >= 28 + // accumulator cliff. Thresholds are from a B300 sweep. BS=192/384 are not + // instantiable: VPT=8 requires kBlockSize to divide kHiddenDim / 8 = 1024. + if (!isBlackwellFamily()) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + return; + } + if constexpr (kNumTokens >= 10 && kNumTokens % 2 == 0) { + launchFp32RouterGemm(output, mat_a, mat_b, stream); + } else if constexpr (kNumTokens >= 13) { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } else { + launchFp32RouterGemm( + output, mat_a, mat_b, stream); + } } else { launchFp32RouterGemm( output, mat_a, mat_b, stream); @@ -287,7 +308,8 @@ void invokeFp32RouterGemm(float* output, InputT const* mat_a, // --------------------------------------------------------------------------- // Explicit instantiations for both input types. -// Supported (E, H) pairs: (256, 3072), (128, 6144), and (256, 5120). +// Supported (E, H) pairs: (256, 3072), (128, 6144), (256, 5120), +// (352, 4096) [Step-5], and (256, 8192) [Muse Spark]. // --------------------------------------------------------------------------- #define INSTANTIATE(T, M, E, H) \ @@ -336,6 +358,11 @@ INSTANTIATE_ALL(__nv_bfloat16, 128, 6144) INSTANTIATE_ALL(float, 256, 5120) INSTANTIATE_ALL(__nv_bfloat16, 256, 5120) +INSTANTIATE_ALL(float, 352, 4096) +INSTANTIATE_ALL(__nv_bfloat16, 352, 4096) +INSTANTIATE_ALL(float, 256, 8192) +INSTANTIATE_ALL(__nv_bfloat16, 256, 8192) + #define INSTANTIATE_LARGE(T, E, H) \ INSTANTIATE(T, 33, E, H) \ INSTANTIATE(T, 34, E, H) \ diff --git a/csrc/fp32_router_gemm_entry.cu b/csrc/fp32_router_gemm_entry.cu index a6efc32831..4a83dbd73a 100644 --- a/csrc/fp32_router_gemm_entry.cu +++ b/csrc/fp32_router_gemm_entry.cu @@ -29,7 +29,9 @@ static constexpr int FP32_MAX_TOKENS = 32; static inline bool fp32_router_gemm_supported(int hidden_dim, int num_experts) { return (hidden_dim == 3072 && num_experts == 256) || (hidden_dim == 6144 && num_experts == 128) || - (hidden_dim == 5120 && num_experts == 256); + (hidden_dim == 5120 && num_experts == 256) || + (hidden_dim == 4096 && num_experts == 352) || + (hidden_dim == 8192 && num_experts == 256); } // Forward declarations — 4 template params must match fp32_router_gemm.cu @@ -80,6 +82,12 @@ void dispatchFp32RouterGemm(int num_experts, int hidden_dim, int num_tokens, } else if (num_experts == 256 && hidden_dim == 5120) { Fp32LoopUnroller::unroll(num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 352 && hidden_dim == 4096) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); + } else if (num_experts == 256 && hidden_dim == 8192) { + Fp32LoopUnroller::unroll( + num_tokens, output, mat_a, mat_b, stream); } else { throw std::invalid_argument( "fp32_router_gemm: unsupported (hidden_dim, num_experts) pair"); @@ -114,7 +122,7 @@ void fp32_router_gemm( STD_TORCH_CHECK( fp32_router_gemm_supported(hidden_dim, num_experts), "fp32_router_gemm: supported (hidden_dim, num_experts) pairs are " - "(3072, 256), (6144, 128), and (5120, 256)"); + "(3072, 256), (6144, 128), (5120, 256), (4096, 352), and (8192, 256)"); const int max_tokens = hidden_dim == 5120 && num_experts == 256 ? 48 : FP32_MAX_TOKENS; STD_TORCH_CHECK(num_tokens <= max_tokens, diff --git a/tests/kernels/test_fp32_router_gemm.py b/tests/kernels/test_fp32_router_gemm.py index 0f55d860a6..54331dc3d3 100644 --- a/tests/kernels/test_fp32_router_gemm.py +++ b/tests/kernels/test_fp32_router_gemm.py @@ -3,11 +3,11 @@ """Tests for fp32_router_gemm kernel: activation×weight→fp32. Supported (hidden_size, num_experts) pairs: - (3072, 256), (6144, 128), (5120, 256) + (3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256) Correctness baseline: F.linear in float32. Every supported M is covered so all tuned geometries (wide-block, experts-per-block, token-group; boundaries -at M=4/5, odd/even, M=15/16) are exercised on Blackwell, and the legacy +at M=4/5, M=9/10, M=12/13, odd/even, M=15/16) are exercised on Blackwell, and the legacy 128/1 geometry everywhere else. """ @@ -18,7 +18,7 @@ from vllm._custom_ops import fp32_router_gemm from vllm.platforms import current_platform # (hidden_size, num_experts) -SHAPES = [(3072, 256), (6144, 128), (5120, 256)] +SHAPES = [(3072, 256), (6144, 128), (5120, 256), (4096, 352), (8192, 256)] CASES = [ (m, hidden, experts) for hidden, experts in SHAPES @@ -126,11 +126,11 @@ def test_topk_routing_consistency(num_tokens: int, hidden_dim: int, num_experts: ) -def test_zero_tokens_returns_empty(): +@pytest.mark.parametrize("hidden_dim,num_experts", SHAPES) +def test_zero_tokens_returns_empty(hidden_dim, num_experts): """M=0 is a graceful no-op returning an empty [0, E] tensor.""" _requires_sm90() device = torch.device("cuda") - hidden_dim, num_experts = SHAPES[0] mat_a = torch.empty(0, hidden_dim, dtype=torch.float32, device=device) mat_b = torch.randn(num_experts, hidden_dim, dtype=torch.float32, device=device) out = fp32_router_gemm(mat_a, mat_b) @@ -183,8 +183,13 @@ def test_rejects_invalid_inputs(hidden_dim, num_experts): @pytest.mark.parametrize( "tokens", [*range(1, 17), 17, 24, 32, 33, 40, 48, 56, 64, 65, 128] ) +@pytest.mark.parametrize( + "hidden_dim,num_experts", [(5120, 256), (4096, 352), (8192, 256)] +) @torch.inference_mode() -def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, tokens): +def test_router_retains_fp32_weights_and_logits( + default_vllm_config, dist_init, tokens, hidden_dim, num_experts +): """Low-latency routing must retain FP32 weight precision and logits.""" from vllm.model_executor.layers.fused_moe.router.gate_linear import GateLinear from vllm.platforms import current_platform @@ -193,16 +198,16 @@ def test_router_retains_fp32_weights_and_logits(default_vllm_config, dist_init, pytest.skip("SM100 required") torch.manual_seed(42) router = GateLinear( - 5120, 256, params_dtype=torch.float32, out_dtype=torch.float32 + hidden_dim, num_experts, params_dtype=torch.float32, out_dtype=torch.float32 ).cuda() router.weight.normal_(std=0.01) - x = torch.randn(tokens, 5120, device="cuda", dtype=torch.bfloat16) + x = torch.randn(tokens, hidden_dim, device="cuda", dtype=torch.bfloat16) expected = (x.double() @ router.weight.double().T).float() actual, bias = router(x) assert bias is None and actual.dtype == torch.float32 torch.testing.assert_close(actual, expected, rtol=2e-5, atol=2e-5) torch.testing.assert_close(actual.topk(8).indices, expected.topk(8).indices) - if tokens in (1, 16, 48, 56): + if tokens in (1, 16, 32, 33, 48, 56): graph = torch.cuda.CUDAGraph() with torch.cuda.graph(graph): captured, _ = router(x) diff --git a/vllm/model_executor/layers/fused_moe/router/gate_linear.py b/vllm/model_executor/layers/fused_moe/router/gate_linear.py index 7d14f105c9..3237783ece 100644 --- a/vllm/model_executor/layers/fused_moe/router/gate_linear.py +++ b/vllm/model_executor/layers/fused_moe/router/gate_linear.py @@ -47,7 +47,13 @@ class GateLinear(ReplicatedLinear): """ # (hidden_size, num_experts) pairs with an instantiated fp32 kernel. - FP32_SUPPORTED_SHAPES = {(3072, 256), (6144, 128), (5120, 256)} + FP32_SUPPORTED_SHAPES = { + (3072, 256), + (6144, 128), + (5120, 256), + (4096, 352), + (8192, 256), + } FP32_MAX_TOKENS = 32 # Largest batch for which tier 1 beats the cuBLAS epilogue. Subclasses # whose router shape stays ahead further up the range may raise it. From d74f4180cdff64d410084173d7f9bf73af70f615 Mon Sep 17 00:00:00 2001 From: Cy Cheng Date: Fri, 16 Oct 2026 22:20:38 -0700 Subject: [PATCH 45/49] [Bugfix][Frontend] Detokenize echoed prompts after truncation (#61876) Signed-off-by: Cy <15118312+cycancode@users.noreply.github.com> Co-authored-by: Cy <15118312+cycancode@users.noreply.github.com> --- tests/renderers/test_completions.py | 27 +++++++++++++++++++++++++++ vllm/renderers/base.py | 12 ++++++++++-- 2 files changed, 37 insertions(+), 2 deletions(-) diff --git a/tests/renderers/test_completions.py b/tests/renderers/test_completions.py index 9d37ff78d1..f7e53b1389 100644 --- a/tests/renderers/test_completions.py +++ b/tests/renderers/test_completions.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import asyncio import io from collections.abc import Sequence from dataclasses import dataclass @@ -376,6 +377,32 @@ class TestRenderPrompt: assert results[0]["prompt_token_ids"] == tokens assert results[0]["prompt"] == "[1, 2, 3, 4]" + @pytest.mark.parametrize( + ("truncation_side", "expected"), + [("left", [106, 107, 108, 109]), ("right", [100, 101, 102, 103])], + ) + def test_detokenization_matches_truncated_tokens(self, truncation_side, expected): + """Echoed prompt text must match the truncated tokens sent to the engine.""" + renderer = _build_renderer(MockModelConfig()) + params = TokenizeParams( + max_total_tokens=100, + truncate_prompt_tokens=4, + truncation_side=truncation_side, + needs_detokenization=True, + ) + + def render(): + return renderer.render_prompts( + _preprocess_prompt(renderer.model_config, list(range(100, 110))) + ) + + results = renderer.tokenize_prompts(render(), params) + async_results = asyncio.run(renderer.tokenize_prompts_async(render(), params)) + + for result in (*results, *async_results): + assert result["prompt_token_ids"] == expected + assert result["prompt"] == str(expected) + def test_explicit_side_tokenizer_unbounded(self): renderer = _build_renderer(MockModelConfig()) diff --git a/vllm/renderers/base.py b/vllm/renderers/base.py index 7b0f4debaa..f22b4215bd 100644 --- a/vllm/renderers/base.py +++ b/vllm/renderers/base.py @@ -687,13 +687,17 @@ class BaseRenderer(ABC, Generic[_T]): prompt = self._tokenize_prompt(prompt, params) prompt = self._apply_prompt_char_offset(prompt, prompt_char_offset) + prompt = params.apply_post_tokenization(self.tokenizer, prompt) # type: ignore[arg-type] + + # Detokenize after truncation so the text matches the tokens sent + # to the engine. if params.needs_detokenization and "prompt" not in prompt: if "prompt_token_ids" not in prompt: raise RuntimeError("Cannot run detokenization on embeddings") prompt = self._detokenize_prompt(prompt) # type: ignore[arg-type] - return params.apply_post_tokenization(self.tokenizer, prompt) # type: ignore[arg-type] + return prompt @overload async def _tokenize_singleton_prompt_async( @@ -727,13 +731,17 @@ class BaseRenderer(ABC, Generic[_T]): prompt = await self._tokenize_prompt_async(prompt, params) prompt = self._apply_prompt_char_offset(prompt, prompt_char_offset) + prompt = params.apply_post_tokenization(self.tokenizer, prompt) # type: ignore[arg-type] + + # Detokenize after truncation so the text matches the tokens sent + # to the engine. if params.needs_detokenization and "prompt" not in prompt: if "prompt_token_ids" not in prompt: raise RuntimeError("Cannot run detokenization on embeddings") prompt = await self._detokenize_prompt_async(prompt) # type: ignore[arg-type] - return params.apply_post_tokenization(self.tokenizer, prompt) # type: ignore[arg-type] + return prompt def _tokenize_enc_dec_prompt( self, From 90d6d871d6700dd229a76b788bf0dca22bab6c70 Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sat, 17 Oct 2026 01:21:55 -0400 Subject: [PATCH 46/49] [Attention] Add CuTe DSL backend selection for MLA context prefill (#62069) Signed-off-by: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Signed-off-by: wzhao18 --- tests/v1/attention/test_mla_backends.py | 38 +++++++++++++++++++ vllm/config/attention.py | 5 +++ .../layers/attention/mla_attention.py | 20 ++++++++++ .../backends/mla/prefill/trtllm_ragged.py | 23 +++++++++++ 4 files changed, 86 insertions(+) diff --git a/tests/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 9d4be102d9..09a8158ff1 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -10,6 +10,7 @@ Known Issues: import sys from types import MethodType, SimpleNamespace +from typing import Literal import pytest import torch @@ -47,6 +48,9 @@ from vllm.v1.attention.backends.mla.prefill.base import MLADimensions from vllm.v1.attention.backends.mla.prefill.selector import ( MLAPrefillSelectorConfig, ) +from vllm.v1.attention.backends.mla.prefill.trtllm_ragged import ( + is_cute_dsl_mla_prefill_supported, +) from vllm.v1.attention.backends.registry import AttentionBackendEnum from vllm.v1.attention.ops.flashmla import is_flashmla_dense_supported from vllm.v1.kv_cache_interface import ( @@ -1483,6 +1487,7 @@ def run_attention_backend( # Force the prefill backend selection (None means auto-select). vllm_config.attention_config.mla_prefill_backend = prefill_backend + mla_attention_module.backend_supports_prefill_query_quantization.cache_clear() # Set the current vllm config so that get_current_vllm_config() works # in the backend implementations @@ -1641,6 +1646,8 @@ def _run_backend_correctness( v_head_dim: int, chunked_prefill_workspace_size: int | None = None, adaptive_verification: bool = False, + trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"] = "trtllm-gen", + use_prefill_query_quantization: bool = False, ): """Test that all backends produce similar outputs to a reference implementation using torch.nn.functional.scaled_dot_product_attention. @@ -1720,6 +1727,10 @@ def _run_backend_correctness( hf_config_override["num_attention_heads"] ) vllm_config.cache_config.cache_dtype = kv_cache_dtype + vllm_config.attention_config.trtllm_ragged_backend = trtllm_ragged_backend + vllm_config.attention_config.use_prefill_query_quantization = ( + use_prefill_query_quantization + ) # For spec decode tests, add a speculative_config to set the reorder_batch_threshold if is_spec_decode_test: @@ -1910,6 +1921,11 @@ def _run_backend_correctness( q_mha = torch.cat([q_nope, q_pe], dim=-1) # [q_len, num_heads, total_dim] k_pe_full_expanded = k_pe_prefill.expand(-1, num_q_heads, -1) k_full = torch.cat([k_nope_full, k_pe_full_expanded], dim=-1) + if use_prefill_query_quantization: + fp8_dtype = current_platform.fp8_dtype() + q_mha = q_mha.to(fp8_dtype).to(dtype) + k_full = k_full.to(fp8_dtype).to(dtype) + v_full = v_full.to(fp8_dtype).to(dtype) # Create custom attention mask: # - Query tokens can attend to all context tokens @@ -2184,6 +2200,8 @@ def test_backend_correctness( _prefill_backend_dimension_params(), ) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize("trtllm_ragged_backend", ["trtllm-gen", "cute-dsl"]) +@pytest.mark.parametrize("use_prefill_query_quantization", [False, True]) def test_chunked_context_backend_correctness( default_vllm_config, dist_init, @@ -2192,8 +2210,26 @@ def test_chunked_context_backend_correctness( qk_nope_head_dim: int, v_head_dim: int, kv_cache_dtype: str, + trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"], + use_prefill_query_quantization: bool, ): """Split, packed, and context-free requests match the SDPA reference.""" + if prefill_backend != MLAPrefillBackendEnum.TRTLLM_RAGGED and ( + trtllm_ragged_backend == "cute-dsl" or use_prefill_query_quantization + ): + pytest.skip("TRTLLM ragged-specific test configuration") + if use_prefill_query_quantization and kv_cache_dtype == "auto": + pytest.skip("FP8 prefill queries require an FP8 KV cache") + if trtllm_ragged_backend == "cute-dsl": + query_dtype = ( + current_platform.fp8_dtype() + if use_prefill_query_quantization + else torch.bfloat16 + ) + if not is_cute_dsl_mla_prefill_supported( + qk_nope_head_dim + 64, v_head_dim, query_dtype + ): + pytest.skip("Unsupported CuTe DSL prefill dimensions or query dtype") _run_backend_correctness( default_vllm_config, dist_init, @@ -2208,6 +2244,8 @@ def test_chunked_context_backend_correctness( qk_nope_head_dim=qk_nope_head_dim, v_head_dim=v_head_dim, chunked_prefill_workspace_size=1024, + trtllm_ragged_backend=trtllm_ragged_backend, + use_prefill_query_quantization=use_prefill_query_quantization, ) diff --git a/vllm/config/attention.py b/vllm/config/attention.py index fc3ecb0dfc..1b7c996a64 100644 --- a/vllm/config/attention.py +++ b/vllm/config/attention.py @@ -85,6 +85,11 @@ class AttentionConfig: use_prefill_query_quantization: bool = False """If set, quantize query for attention in prefill.""" + trtllm_ragged_backend: Literal["trtllm-gen", "cute-dsl"] = "trtllm-gen" + """Kernel backend for cached-context TRTLLM_RAGGED MLA prefill. + CuTe DSL requires FP8 queries for head dimension 192 and does not support + head dimension 256. New-token prefill always uses TRTLLM.""" + indexer_kv_dtype: IndexerKVDType = "auto" """Data type for the sparse-attention indexer K cache. "auto" picks the model's default (bf16 for MiniMax M3, fp8 for the DeepSeek sparse diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index caf60b009f..865913b782 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2024,6 +2024,7 @@ def build_mla_chunked_context_metadata( dcp_local_block_size: int, dcp_virtual_block_size: int, dcp_manager: MLADCPManager | None = None, + align_indptr: bool = False, ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None": """Build chunked-context metadata for an MLA prefill. @@ -2043,6 +2044,7 @@ def build_mla_chunked_context_metadata( dcp_local_block_size: Per-rank interleave block size for DCP. dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``. dcp_manager: Shared MLA DCP collective manager. + align_indptr: Align packed int32 indptr slices to 16 bytes for CuTe DSL. Returns: The chunked-context metadata, or None when no prefill has any context. @@ -2123,6 +2125,15 @@ def build_mla_chunked_context_metadata( num_tokens = sum(plan.seq_lens) num_local_tokens = num_tokens + if align_indptr: + # CuTe DSL requires 16-byte-aligned int32 indptrs. + padding = -boundary_offset % 4 + cu_seq_lens_flat.extend([0] * padding) + cu_seqlens_q_flat.extend([0] * padding) + if use_dcp: + padded_local_cu_seq_lens_flat.extend([0] * padding) + boundary_offset += padding + seq_lens_flat.extend(plan.seq_lens) cu_seq_lens_flat.extend(itertools.accumulate(plan.seq_lens, initial=0)) query_base = query_start_loc[plan.request_start] @@ -2616,6 +2627,14 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): query_start_loc_cpu[reqs_start:] - query_start_loc_cpu[reqs_start] ) + align_indptr = False + if ( + self._prefill_backend.get_name() == "TRTLLM_RAGGED" + and self.vllm_config.attention_config.trtllm_ragged_backend + == "cute-dsl" + ): + align_indptr = True + chunked_context_metadata = build_mla_chunked_context_metadata( context_lens_cpu=context_lens_cpu, prefill_query_start_loc_cpu=prefill_query_start_loc_cpu, @@ -2628,6 +2647,7 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): dcp_local_block_size=self.dcp_local_block_size, dcp_virtual_block_size=self.dcp_virtual_block_size, dcp_manager=self.dcp_manager, + align_indptr=align_indptr, ) prefill_metadata = MLACommonPrefillMetadata( diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index 49c9dfa21c..a1f21cb0d1 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -22,6 +22,17 @@ if TYPE_CHECKING: from vllm.platforms.interface import DeviceCapability +def is_cute_dsl_mla_prefill_supported( + qk_head_dim: int, v_head_dim: int, query_dtype: torch.dtype +) -> bool: + # FlashInfer has no public compatibility API; replace these checks when available. + return ( + query_dtype in (torch.float16, torch.bfloat16, torch.float8_e4m3fn) + and (qk_head_dim, v_head_dim) in ((128, 128), (192, 128)) + and (qk_head_dim != 192 or query_dtype == torch.float8_e4m3fn) + ) + + class TrtllmRaggedPrefillBackend(MLAPrefillBackend): """TRT-LLM Ragged backend for MLA prefill.""" @@ -79,6 +90,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): v_head_dim=v_head_dim, vllm_config=vllm_config, ) + self._context_backend = vllm_config.attention_config.trtllm_ragged_backend self._workspace_spec = ( (envs.VLLM_FLASHINFER_WORKSPACE_BUFFER_SIZE,), torch.uint8, @@ -198,6 +210,16 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): "TRTLLM ragged context prefill contains an empty query or KV row" ) + if self._context_backend == "cute-dsl": + if not is_cute_dsl_mla_prefill_supported(q.shape[-1], v.shape[-1], q.dtype): + raise ValueError( + "CuTe DSL context prefill does not support " + f"Q/K head dimension {q.shape[-1]}, " + f"V head dimension {v.shape[-1]}, and query dtype {q.dtype}." + ) + # Generic MLA supplies V as a projection slice; CuTe requires packed V. + v = v.contiguous() + if out is None: out = torch.empty( q.shape[0], @@ -225,6 +247,7 @@ class TrtllmRaggedPrefillBackend(MLAPrefillBackend): cum_seq_lens_kv=chunk.cu_seq_lens, enable_pdl=False, is_causal=False, + backend=self._context_backend, return_lse=True, out=out, skip_all_rows_active_check=True, From 5be4266f5c7c10c69e4be5f3ec01d6a1d8bae9a1 Mon Sep 17 00:00:00 2001 From: Jovan <101700786+KuanJia@users.noreply.github.com> Date: Sat, 17 Oct 2026 14:10:16 +0800 Subject: [PATCH 47/49] [Bugfix][Frontend] Preserve Cohere and Anthropic streaming errors (#61947) Signed-off-by: Kevin Co-authored-by: OpenAI Codex --- .../test_anthropic_messages_conversion.py | 85 ++-- .../cohere/test_serving_streaming.py | 135 +++--- .../chat_completion/test_serving_chat.py | 146 +++++- vllm/entrypoints/anthropic/serving.py | 435 +++++++++--------- vllm/entrypoints/cohere/serving.py | 40 +- .../openai/chat_completion/serving.py | 65 ++- 6 files changed, 531 insertions(+), 375 deletions(-) diff --git a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py index 399c898d0b..5e82072d18 100644 --- a/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py +++ b/tests/entrypoints/anthropic/test_anthropic_messages_conversion.py @@ -49,6 +49,9 @@ from vllm.entrypoints.openai.chat_completion.protocol import ( ChatMessage, ) from vllm.entrypoints.serve.engine.protocol import PromptTokenUsageInfo, UsageInfo +from vllm.entrypoints.serve.exception_handling.error_response import ( + create_error_response, +) from vllm.entrypoints.serve.exception_handling.handlers.validation import ( validation_exception_handler, ) @@ -1034,7 +1037,7 @@ def _make_stream_chunk( stop_reason: int | str | None = None, choices: list[ChatCompletionResponseStreamChoice] | None = None, usage: UsageInfo | None = None, -) -> str: +) -> ChatCompletionStreamResponse: if choices is None: choices = [ ChatCompletionResponseStreamChoice( @@ -1051,7 +1054,7 @@ def _make_stream_chunk( choices=choices, usage=usage, ) - return f"data: {chunk.model_dump_json()}" + return chunk def _tc(*, args, id=None, name=None): @@ -1062,6 +1065,39 @@ def _tc(*, args, id=None, name=None): ) +@pytest.mark.asyncio +@pytest.mark.parametrize("block_type", [None, "text", "thinking", "tool_use"]) +async def test_stream_error_preserves_message_and_closes_block(block_type): + async def source(): + if block_type is not None: + yield _make_stream_chunk(delta=DeltaMessage(role="assistant")) + if block_type == "text": + delta = DeltaMessage(content="Hello") + elif block_type == "thinking": + delta = DeltaMessage(reasoning="Thinking") + else: + delta = DeltaMessage( + tool_calls=[_tc(id="call_1", name="search", args='{"q":')] + ) + yield _make_stream_chunk(delta=delta) + yield create_error_response(ValueError("backend failed")) + + converter = _make_stream_converter() + events = _parse_sse_events( + [event async for event in converter.message_stream_converter(source())] + ) + assert events[-1] == ( + "error", + { + "type": "error", + "error": {"type": "internal_error", "message": "backend failed"}, + }, + ) + assert not any(event_type == "message_stop" for event_type, _ in events) + if block_type is not None: + assert events[-2][0] == "content_block_stop" + + class TestMessageStreamConverterToolUseContentBuffering: """Regression test for tool_use arguments being silently dropped. @@ -1074,7 +1110,7 @@ class TestMessageStreamConverterToolUseContentBuffering: async def test_tool_use_args_not_dropped_when_content_in_same_chunk( self, ): - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(role="assistant"), usage=UsageInfo(prompt_tokens=10, total_tokens=10), @@ -1111,11 +1147,10 @@ class TestMessageStreamConverterToolUseContentBuffering: completion_tokens=20, ), ) - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1156,7 +1191,7 @@ class TestMessageStreamConverterToolUseContentBuffering: stream jumps straight from finish_reason to [DONE], skipping the empty-choices usage chunk.""" - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(role="assistant"), usage=UsageInfo(prompt_tokens=10, total_tokens=10), @@ -1178,11 +1213,10 @@ class TestMessageStreamConverterToolUseContentBuffering: finish_reason="tool_calls", ) # No empty-choices usage chunk — go straight to [DONE]. - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1216,7 +1250,7 @@ class TestMessageStartIncludesTypeAndRole: @pytest.mark.asyncio async def test_message_start_contains_message_type_and_role(self): - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(content="Hello"), usage=UsageInfo( @@ -1226,11 +1260,10 @@ class TestMessageStartIncludesTypeAndRole: ), ) yield _make_stream_chunk(finish_reason="stop") - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1260,7 +1293,7 @@ class TestStreamingCacheUsageSemantics: message_start omits cache fields. The final chunk carries prompt_tokens_details, so message_delta carries resolved values.""" - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(role="assistant", content="hi"), usage=UsageInfo(prompt_tokens=100, total_tokens=100), @@ -1277,11 +1310,10 @@ class TestStreamingCacheUsageSemantics: ), ), ) - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1305,7 +1337,7 @@ class TestStreamingCacheUsageSemantics: """When the final chunk reports cached_tokens=0, message_delta carries cache fields = 0 (cache miss); message_start still omits them.""" - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(role="assistant"), usage=UsageInfo(prompt_tokens=50, total_tokens=50), @@ -1322,11 +1354,10 @@ class TestStreamingCacheUsageSemantics: ), ), ) - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1346,7 +1377,7 @@ class TestStreamingCacheUsageSemantics: """If --enable-prompt-tokens-details is off, no chunk carries cache info; both message_start and message_delta omit cache fields.""" - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(role="assistant"), usage=UsageInfo(prompt_tokens=30, total_tokens=30), @@ -1356,11 +1387,10 @@ class TestStreamingCacheUsageSemantics: choices=[], usage=UsageInfo(prompt_tokens=30, completion_tokens=2, total_tokens=32), ) - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1628,7 +1658,7 @@ class TestStopSequenceReason: @pytest.mark.asyncio async def test_streaming_stop_string_maps_to_stop_sequence(self): - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk( delta=DeltaMessage(content="hi"), usage=UsageInfo(prompt_tokens=5, total_tokens=5, completion_tokens=0), @@ -1638,11 +1668,10 @@ class TestStopSequenceReason: choices=[], usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), ) - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1654,7 +1683,7 @@ class TestStopSequenceReason: async def test_streaming_no_stop_string_emits_explicit_null_stop_sequence(self): """exclude_unset=True drops stop_sequence unless it is set explicitly.""" - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk(delta=DeltaMessage(role="assistant")) yield _make_stream_chunk(delta=DeltaMessage(content="hi")) yield _make_stream_chunk(finish_reason="stop") @@ -1662,11 +1691,10 @@ class TestStopSequenceReason: choices=[], usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), ) - yield "data: [DONE]" converter = _make_stream_converter() output = [] - async for event in converter.message_stream_converter(sse_input()): + async for event in converter.message_stream_converter(chunk_input()): output.append(event) events = _parse_sse_events(output) @@ -1724,7 +1752,7 @@ class TestToolUseStopReason: ], ) async def test_streaming(self, args, finish_reason, stop_reason, expected): - async def sse_input(): + async def chunk_input(): yield _make_stream_chunk(delta=DeltaMessage(role="assistant")) yield _make_stream_chunk( delta=DeltaMessage( @@ -1738,11 +1766,10 @@ class TestToolUseStopReason: choices=[], usage=UsageInfo(prompt_tokens=5, total_tokens=8, completion_tokens=3), ) - yield "data: [DONE]" converter = _make_stream_converter() events = _parse_sse_events( - [e async for e in converter.message_stream_converter(sse_input())] + [e async for e in converter.message_stream_converter(chunk_input())] ) msg_deltas = [data for ev_type, data in events if ev_type == "message_delta"] diff --git a/tests/entrypoints/cohere/test_serving_streaming.py b/tests/entrypoints/cohere/test_serving_streaming.py index 8c2e8623ad..c0dc86bf9a 100644 --- a/tests/entrypoints/cohere/test_serving_streaming.py +++ b/tests/entrypoints/cohere/test_serving_streaming.py @@ -5,7 +5,7 @@ The stream-translation entry point is :meth:`CohereServingChatV2._chat_completion_stream_to_v2`, which turns an -async iterable of OpenAI SSE chunks into Cohere's +async iterable of OpenAI chunk objects into Cohere's ``message-start → (content|tool-call|citation)* → message-end → [DONE]`` event stream. @@ -15,7 +15,7 @@ state machine. """ import json -from collections.abc import AsyncGenerator +from collections.abc import AsyncGenerator, Sequence from typing import Any import pytest @@ -26,10 +26,7 @@ from vllm.entrypoints.cohere.cohere_chat_message import ( from vllm.entrypoints.cohere.cohere_chat_message import ( CitationSource, ) -from vllm.entrypoints.cohere.protocol import ( - CohereChatV2Request, - MessageStartEvent, -) +from vllm.entrypoints.cohere.protocol import MessageStartEvent from vllm.entrypoints.cohere.serving import ( _DONE_FRAME, CohereServingChatV2, @@ -41,6 +38,7 @@ from vllm.entrypoints.cohere.serving import ( from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionStreamResponse, ) +from vllm.entrypoints.serve.engine.protocol import ErrorResponse # ---------------------------------------------------------------------- # Helpers @@ -85,9 +83,8 @@ def _make_chunk( usage: dict[str, Any] | None = None, omit_choices: bool = False, citations: list[Any] | None = None, -) -> str: - """Build the ``data: {...}\\n\\n`` SSE frame the production code - consumes.""" +) -> ChatCompletionStreamResponse: + """Build the raw chat chunk the production converter consumes.""" delta: dict[str, Any] = {} if role is not None: delta["role"] = role @@ -119,27 +116,23 @@ def _make_chunk( payload["choices"] = [] if usage is not None: payload["usage"] = usage - # Use ChatCompletionStreamResponse to normalize the payload shape. - chunk = ChatCompletionStreamResponse.model_validate(payload) - return f"data: {chunk.model_dump_json(exclude_none=False)}\n\n" + return ChatCompletionStreamResponse.model_validate(payload) -def _make_done() -> str: - return "data: [DONE]\n\n" - - -async def _stream_from(items: list[str]) -> AsyncGenerator[str, None]: +async def _stream_from( + items: Sequence[ChatCompletionStreamResponse | ErrorResponse], +) -> AsyncGenerator[ChatCompletionStreamResponse | ErrorResponse, None]: for item in items: yield item -async def _drain(serving: CohereServingChatV2, items: list[str]) -> list[str]: +async def _drain( + serving: CohereServingChatV2, + items: Sequence[ChatCompletionStreamResponse | ErrorResponse], +) -> list[str]: """Drive ``_chat_completion_stream_to_v2`` over ``items`` and collect the emitted SSE frames.""" - request = CohereChatV2Request( - model="m", messages=[{"role": "user", "content": "hi"}], stream=True - ) - gen = serving._chat_completion_stream_to_v2(_stream_from(items), request) + gen = serving._chat_completion_stream_to_v2(_stream_from(items)) return [frame async for frame in gen] @@ -793,15 +786,14 @@ class TestChatCompletionStreamToV2: assert "content-end" not in types @pytest.mark.asyncio - async def test_done_marker_in_middle_closes_open_block(self): - # Some upstreams send [DONE] without a trailing usage-only chunk. + async def test_stream_exhaustion_closes_open_block_without_usage(self): + # The raw stream may finish without a trailing usage-only chunk. # The translator must still emit message-end before [DONE] so # Cohere clients don't hang. serving = _serving(is_reasoning_model=True) items = [ _make_chunk(role="assistant"), _make_chunk(content="Hi"), - _make_done(), ] frames = await _drain(serving, items) types = [_parse_event(f)["type"] for f in frames[:-1]] @@ -815,72 +807,65 @@ class TestChatCompletionStreamToV2: assert frames[-1] == _DONE_FRAME @pytest.mark.asyncio - async def test_skips_empty_and_non_data_lines(self): + @pytest.mark.parametrize("block_type", [None, "text", "thinking", "tool_call"]) + async def test_upstream_error_preserves_message(self, block_type): serving = _serving() - items = [ - "\n", - "event: ping\n\n", - _make_chunk(role="assistant"), - "data: \n\n", # empty data - _make_chunk(content="Hi"), - _make_chunk( - omit_choices=True, - usage={ - "prompt_tokens": 1, - "completion_tokens": 1, - "total_tokens": 2, - }, - ), - ] + items: list[ChatCompletionStreamResponse | ErrorResponse] = [] + if block_type is not None: + items.append(_make_chunk(role="assistant")) + if block_type == "text": + items.append(_make_chunk(content="Hi")) + elif block_type == "thinking": + items.append(_make_chunk(reasoning="Thinking")) + else: + items.append( + _make_chunk( + tool_calls=[ + { + "index": 0, + "id": "call_1", + "function": {"name": "search", "arguments": '{"q":'}, + } + ] + ) + ) + items.append(serving.create_error_response(ValueError("backend failed"))) + frames = await _drain(serving, items) assert frames[-1] == _DONE_FRAME - # Should still produce a complete lifecycle. - types = [_parse_event(f)["type"] for f in frames[:-1]] - assert types[0] == "message-start" - assert types[-1] == "message-end" + events = [_parse_event(frame) for frame in frames[:-1]] + error_ends = [event for event in events if event["type"] == "message-end"] + assert len(error_ends) == 1 + assert error_ends[0]["delta"] == { + "finish_reason": "ERROR", + "error": "backend failed", + } + assert events[-1] == error_ends[0] + if block_type is not None: + assert events[-2]["type"] == ( + "tool-call-end" if block_type == "tool_call" else "content-end" + ) @pytest.mark.asyncio - async def test_exception_in_chunk_parsing_emits_error_message_end(self): - # An invalid JSON payload after the first chunk triggers the - # error path: a synthetic message-end with finish_reason=ERROR - # followed by [DONE]. - serving = _serving() - items = [ - _make_chunk(role="assistant"), - "data: {not valid json}\n\n", - ] - frames = await _drain(serving, items) - assert frames[-1] == _DONE_FRAME - # Find the error-shaped message-end. - error_end = _parse_event(frames[-2]) - assert error_end["type"] == "message-end" - assert error_end["delta"]["finish_reason"] == "ERROR" - assert "error" in error_end["delta"] - - @pytest.mark.asyncio - async def test_error_message_end_is_sanitized(self): + @pytest.mark.parametrize("started", [False, True]) + async def test_error_message_end_is_sanitized(self, started): # The client-visible ``error`` field on the terminal message-end # must be routed through ``sanitize_message`` so tracebacks, # host filesystem paths, and Python object memory addresses # can't leak from a mid-stream exception into the SSE payload. serving = _serving() - async def _raises_mid_stream() -> AsyncGenerator[str, None]: - yield _make_chunk(role="assistant") + async def _raises_mid_stream() -> AsyncGenerator[ + ChatCompletionStreamResponse | ErrorResponse, None + ]: + if started: + yield _make_chunk(role="assistant") raise RuntimeError( "boom reading /Users/dev/vllm/x.py" ) - request = CohereChatV2Request( - model="m", - messages=[{"role": "user", "content": "hi"}], - stream=True, - ) frames = [ - f - async for f in serving._chat_completion_stream_to_v2( - _raises_mid_stream(), request - ) + f async for f in serving._chat_completion_stream_to_v2(_raises_mid_stream()) ] assert frames[-1] == _DONE_FRAME error_end = _parse_event(frames[-2]) diff --git a/tests/entrypoints/openai/chat_completion/test_serving_chat.py b/tests/entrypoints/openai/chat_completion/test_serving_chat.py index 85404af06c..0c7526cfc9 100644 --- a/tests/entrypoints/openai/chat_completion/test_serving_chat.py +++ b/tests/entrypoints/openai/chat_completion/test_serving_chat.py @@ -28,6 +28,7 @@ from vllm.entrypoints.generate.base.serving import build_per_request_timing_metr from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionRequest, ChatCompletionResponse, + ChatCompletionStreamResponse, ) from vllm.entrypoints.openai.chat_completion.serving import ( OpenAIServingChat, @@ -43,7 +44,7 @@ from vllm.entrypoints.openai.models.serving import ( from vllm.entrypoints.openai.parser.harmony_utils import get_encoding from vllm.entrypoints.serve.engine.protocol import ErrorResponse from vllm.entrypoints.serve.utils.tool_calls_utils import resolve_finish_reason -from vllm.exceptions import QueueOverflowError, VLLMValidationError +from vllm.exceptions import GenerationError, QueueOverflowError, VLLMValidationError from vllm.inputs import TokensPrompt from vllm.logprobs import Logprob from vllm.multimodal.inputs import PlaceholderRange @@ -716,6 +717,149 @@ async def _collect_metrics_stream_chunks( return chunks +@pytest.mark.asyncio +@pytest.mark.parametrize("raw", [False, True]) +async def test_streaming_chunks_keep_usage_and_metadata(raw): + serving = _build_minimal_metrics_serving_chat(False) + request = ChatCompletionRequest( + model="test-model", + messages=[{"role": "user", "content": "Test"}], + stream=True, + stream_options={"include_usage": True}, + ) + metadata = RequestResponseMetadata(request_id="chatcmpl-test") + method = ( + serving.chat_completion_stream_generator_raw + if raw + else serving.chat_completion_stream_generator + ) + chunks = [ + chunk + async for chunk in method( + request, + _single_request_output(_make_metrics_request_output()), + "chatcmpl-test", + "test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=metadata, + ) + ] + if raw: + assert all(isinstance(chunk, ChatCompletionStreamResponse) for chunk in chunks) + else: + assert chunks[-1] == "data: [DONE]\n\n" + first = json.loads(chunks[0].removeprefix("data: ")) + assert first["choices"][0]["logprobs"] is None + usage = json.loads(chunks[-2].removeprefix("data: ")) + assert "system_fingerprint" not in usage + assert "prompt_tokens_details" not in usage["usage"] + chunks = [ + ChatCompletionStreamResponse.model_validate_json( + chunk.removeprefix("data: ") + ) + for chunk in chunks[:-1] + ] + assert len(chunks) == 3 + assert chunks[0].choices[0].delta.role == "assistant" + assert chunks[1].choices[0].delta.content == "Hello" + assert chunks[1].choices[0].finish_reason == "stop" + assert chunks[2].choices == [] + assert chunks[2].usage.prompt_tokens == 3 + assert chunks[2].usage.completion_tokens == 2 + assert metadata.final_usage_info.total_tokens == 5 + + +@pytest.mark.asyncio +async def test_streaming_serialization_failure_emits_error_and_done(): + serving = _build_minimal_metrics_serving_chat(False) + serving.parser_cls = HarmonyParser + parser = MagicMock() + parser.parse_delta.return_value = DeltaMessage( + content="Hello", extra_object=object() + ) + serving._make_parser = MagicMock(return_value=parser) + request = ChatCompletionRequest( + model="test-model", messages=[{"role": "user", "content": "Test"}], stream=True + ) + frames = [ + frame + async for frame in serving.chat_completion_stream_generator( + request, + _single_request_output(_make_metrics_request_output()), + "chatcmpl-test", + "test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test"), + ) + ] + assert frames[-1] == "data: [DONE]\n\n" + error = ErrorResponse.model_validate_json(frames[-2].removeprefix("data: ")) + assert "serialize" in error.error.message + + +@pytest.mark.asyncio +@pytest.mark.parametrize("raw", [False, True]) +@pytest.mark.parametrize( + "failure", ["parser", "before_first", "after_content", "generation"] +) +async def test_streaming_errors_preserve_envelope(raw, failure): + serving = _build_minimal_metrics_serving_chat(False) + error = ( + GenerationError("backend failed") + if failure == "generation" + else ValueError("backend failed") + ) + if failure == "parser": + serving.parser_cls = HarmonyParser + serving._make_parser = MagicMock(side_effect=error) + + async def results(): + if failure == "after_content": + output = _make_metrics_request_output() + output.outputs[0].finish_reason = None + output.finished = False + yield output + raise error + + request = ChatCompletionRequest( + model="test-model", messages=[{"role": "user", "content": "Test"}], stream=True + ) + method = ( + serving.chat_completion_stream_generator_raw + if raw + else serving.chat_completion_stream_generator + ) + chunks = [ + chunk + async for chunk in method( + request, + results(), + "chatcmpl-test", + "test-model", + conversation=[], + tokenizer=MagicMock(), + request_metadata=RequestResponseMetadata(request_id="chatcmpl-test"), + ) + ] + if raw: + assert all( + isinstance(chunk, (ChatCompletionStreamResponse, ErrorResponse)) + for chunk in chunks + ) + assert isinstance(chunks[-1], ErrorResponse) + error_response = chunks[-1] + else: + assert chunks[-1] == "data: [DONE]\n\n" + assert chunks.count("data: [DONE]\n\n") == 1 + error_response = ErrorResponse.model_validate_json( + chunks[-2].removeprefix("data: ") + ) + assert error_response.error.message == "backend failed" + assert error_response.error.code == (500 if failure == "generation" else 400) + + def test_build_per_request_timing_metrics_valid_timestamps(): metrics = build_per_request_timing_metrics( _PER_REQUEST_STATS, num_generation_tokens=10 diff --git a/vllm/entrypoints/anthropic/serving.py b/vllm/entrypoints/anthropic/serving.py index 4b7f2b95a5..8eb6448550 100644 --- a/vllm/entrypoints/anthropic/serving.py +++ b/vllm/entrypoints/anthropic/serving.py @@ -820,7 +820,14 @@ class AnthropicServingMessages(OpenAIServingChat): elif isinstance(generator, ChatCompletionResponse): return self.messages_full_converter(generator) - return self.message_stream_converter(generator) + return generator + + def chat_completion_stream_generator( + self, *args: Any, **kwargs: Any + ) -> AsyncGenerator[str, None]: + return self.message_stream_converter( + self.chat_completion_stream_generator_raw(*args, **kwargs) + ) def messages_full_converter( self, @@ -881,7 +888,7 @@ class AnthropicServingMessages(OpenAIServingChat): async def message_stream_converter( self, - generator: AsyncGenerator[str, None], + generator: AsyncGenerator[ChatCompletionStreamResponse | ErrorResponse, None], ) -> AsyncGenerator[str, None]: try: @@ -991,238 +998,208 @@ class AnthropicServingMessages(OpenAIServingChat): events.extend(stop_active_block()) return events - async for item in generator: - if item.startswith("data:"): - data_str = item[5:].strip().rstrip("\n") - if data_str == "[DONE]": - for event in stop_and_flush(): - yield event - stop_message = AnthropicStreamEvent( - type="message_stop", - ) - data = stop_message.model_dump_json( - exclude_unset=True, exclude_none=True - ) - yield wrap_data_with_event(data, "message_stop") - else: - origin_chunk = ChatCompletionStreamResponse.model_validate_json( - data_str - ) - - if first_item: - chunk = AnthropicStreamEvent( - type="message_start", - message=AnthropicMessagesResponse( - id=origin_chunk.id, - # Set explicitly: this event is serialized - # with exclude_unset=True, which drops - # default-valued fields, while strict - # Anthropic SDK clients require - # message.type/role (issue #45367). - type="message", - role="assistant", - content=[], - model=origin_chunk.model, - stop_reason=None, - stop_sequence=None, - usage=_build_anthropic_usage(origin_chunk.usage), - ), - ) - first_item = False - data = chunk.model_dump_json(exclude_unset=True) - yield wrap_data_with_event(data, "message_start") - continue - - # last chunk including usage info - if len(origin_chunk.choices) == 0: - for event in stop_and_flush(): - yield event - anthropic_stop_reason, anthropic_stop_sequence = ( - _to_anthropic_stop_reason( - finish_reason or "stop", - stop_sequence, - tool_used=tool_used, - tool_complete=( - tool_used and _is_json("".join(last_tool_args)) - ), - ) - ) - stop_delta = AnthropicDelta( - stop_reason=anthropic_stop_reason, - # Set explicitly so exclude_unset=True keeps it. - stop_sequence=anthropic_stop_sequence, - ) - chunk = AnthropicStreamEvent( - type="message_delta", - delta=stop_delta, - usage=_build_anthropic_usage(origin_chunk.usage), - ) - data = chunk.model_dump_json(exclude_unset=True) - yield wrap_data_with_event(data, "message_delta") - continue - - if origin_chunk.choices[0].finish_reason is not None: - finish_reason = origin_chunk.choices[0].finish_reason - stop_sequence = origin_chunk.choices[0].stop_reason - # continue - - # thinking / text content - reasoning_delta = origin_chunk.choices[0].delta.reasoning - if reasoning_delta is not None: - if reasoning_delta == "": - pass - else: - if state.block_type != "thinking": - for event in stop_and_flush(): - yield event - start_event = start_block( - AnthropicContentBlock( - type="thinking", thinking="" - ) - ) - yield start_event - chunk = AnthropicStreamEvent( - index=( - state.block_index - if state.block_index is not None - else state.content_block_index - ), - type="content_block_delta", - delta=AnthropicDelta( - type="thinking_delta", - thinking=reasoning_delta, - ), - ) - data = chunk.model_dump_json(exclude_unset=True) - yield wrap_data_with_event(data, "content_block_delta") - - if origin_chunk.choices[0].delta.content is not None: - if origin_chunk.choices[0].delta.content == "": - pass - elif state.block_type == "tool_use": - state.pending_content.append( - origin_chunk.choices[0].delta.content - ) - else: - if state.block_type != "text": - for event in stop_and_flush(): - yield event - start_event = start_block( - AnthropicContentBlock(type="text", text="") - ) - yield start_event - chunk = AnthropicStreamEvent( - index=( - state.block_index - if state.block_index is not None - else state.content_block_index - ), - type="content_block_delta", - delta=AnthropicDelta( - type="text_delta", - text=origin_chunk.choices[0].delta.content, - ), - ) - data = chunk.model_dump_json(exclude_unset=True) - yield wrap_data_with_event(data, "content_block_delta") - - # tool calls - process all tool calls in the delta - if len(origin_chunk.choices[0].delta.tool_calls) > 0: - for tool_call in origin_chunk.choices[0].delta.tool_calls: - if tool_call.id is not None: - # Update mapping for incremental updates - tool_index_to_id[tool_call.index] = tool_call.id - # Only create new block if different tool call - # AND has a name - tool_name = ( - tool_call.function.name - if tool_call.function - else None - ) - if ( - state.tool_use_id != tool_call.id - and tool_name is not None - ): - for event in stop_and_flush(): - yield event - tool_used = True - last_tool_args = [] - start_event = start_block( - AnthropicContentBlock( - type="tool_use", - id=tool_call.id, - name=tool_name, - input={}, - ) - ) - yield start_event - # Handle initial arguments if present - if ( - tool_call.function - and tool_call.function.arguments - and state.tool_use_id == tool_call.id - ): - last_tool_args.append( - tool_call.function.arguments - ) - chunk = AnthropicStreamEvent( - index=( - state.block_index - if state.block_index is not None - else state.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: - # Incremental update - use index to find tool_use_id - tool_use_id = tool_index_to_id.get(tool_call.index) - if ( - tool_use_id is not None - and tool_call.function - and tool_call.function.arguments - and state.tool_use_id == tool_use_id - ): - last_tool_args.append( - tool_call.function.arguments - ) - chunk = AnthropicStreamEvent( - index=( - state.block_index - if state.block_index is not None - else state.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" - ) - continue - else: - error_response = AnthropicStreamEvent( - type="error", - error=AnthropicError( - type="internal_error", - message="Invalid data format received", + async for origin_chunk in generator: + if isinstance(origin_chunk, ErrorResponse): + raise RuntimeError(origin_chunk.error.message) + if first_item: + chunk = AnthropicStreamEvent( + type="message_start", + message=AnthropicMessagesResponse( + id=origin_chunk.id, + # Set explicitly: this event is serialized + # with exclude_unset=True, which drops + # default-valued fields, while strict + # Anthropic SDK clients require + # message.type/role (issue #45367). + type="message", + role="assistant", + content=[], + model=origin_chunk.model, + stop_reason=None, + stop_sequence=None, + usage=_build_anthropic_usage(origin_chunk.usage), ), ) - data = error_response.model_dump_json(exclude_unset=True) - yield wrap_data_with_event(data, "error") + first_item = False + data = chunk.model_dump_json(exclude_unset=True) + yield wrap_data_with_event(data, "message_start") + continue + + # last chunk including usage info + if len(origin_chunk.choices) == 0: + for event in stop_and_flush(): + yield event + anthropic_stop_reason, anthropic_stop_sequence = ( + _to_anthropic_stop_reason( + finish_reason or "stop", + stop_sequence, + tool_used=tool_used, + tool_complete=( + tool_used and _is_json("".join(last_tool_args)) + ), + ) + ) + stop_delta = AnthropicDelta( + stop_reason=anthropic_stop_reason, + # Set explicitly so exclude_unset=True keeps it. + stop_sequence=anthropic_stop_sequence, + ) + chunk = AnthropicStreamEvent( + type="message_delta", + delta=stop_delta, + usage=_build_anthropic_usage(origin_chunk.usage), + ) + data = chunk.model_dump_json(exclude_unset=True) + yield wrap_data_with_event(data, "message_delta") + continue + + if origin_chunk.choices[0].finish_reason is not None: + finish_reason = origin_chunk.choices[0].finish_reason + stop_sequence = origin_chunk.choices[0].stop_reason + # continue + + # thinking / text content + reasoning_delta = origin_chunk.choices[0].delta.reasoning + if reasoning_delta is not None: + if reasoning_delta == "": + pass + else: + if state.block_type != "thinking": + for event in stop_and_flush(): + yield event + start_event = start_block( + AnthropicContentBlock(type="thinking", thinking="") + ) + yield start_event + chunk = AnthropicStreamEvent( + index=( + state.block_index + if state.block_index is not None + else state.content_block_index + ), + type="content_block_delta", + delta=AnthropicDelta( + type="thinking_delta", + thinking=reasoning_delta, + ), + ) + data = chunk.model_dump_json(exclude_unset=True) + yield wrap_data_with_event(data, "content_block_delta") + + if origin_chunk.choices[0].delta.content is not None: + if origin_chunk.choices[0].delta.content == "": + pass + elif state.block_type == "tool_use": + state.pending_content.append( + origin_chunk.choices[0].delta.content + ) + else: + if state.block_type != "text": + for event in stop_and_flush(): + yield event + start_event = start_block( + AnthropicContentBlock(type="text", text="") + ) + yield start_event + chunk = AnthropicStreamEvent( + index=( + state.block_index + if state.block_index is not None + else state.content_block_index + ), + type="content_block_delta", + delta=AnthropicDelta( + type="text_delta", + text=origin_chunk.choices[0].delta.content, + ), + ) + data = chunk.model_dump_json(exclude_unset=True) + yield wrap_data_with_event(data, "content_block_delta") + + # tool calls - process all tool calls in the delta + if len(origin_chunk.choices[0].delta.tool_calls) > 0: + for tool_call in origin_chunk.choices[0].delta.tool_calls: + if tool_call.id is not None: + # Update mapping for incremental updates + tool_index_to_id[tool_call.index] = tool_call.id + # Only create new block if different tool call + # AND has a name + tool_name = ( + tool_call.function.name if tool_call.function else None + ) + if ( + state.tool_use_id != tool_call.id + and tool_name is not None + ): + for event in stop_and_flush(): + yield event + tool_used = True + last_tool_args = [] + start_event = start_block( + AnthropicContentBlock( + type="tool_use", + id=tool_call.id, + name=tool_name, + input={}, + ) + ) + yield start_event + # Handle initial arguments if present + if ( + tool_call.function + and tool_call.function.arguments + and state.tool_use_id == tool_call.id + ): + last_tool_args.append(tool_call.function.arguments) + chunk = AnthropicStreamEvent( + index=( + state.block_index + if state.block_index is not None + else state.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: + # Incremental update - use index to find tool_use_id + tool_use_id = tool_index_to_id.get(tool_call.index) + if ( + tool_use_id is not None + and tool_call.function + and tool_call.function.arguments + and state.tool_use_id == tool_use_id + ): + last_tool_args.append(tool_call.function.arguments) + chunk = AnthropicStreamEvent( + index=( + state.block_index + if state.block_index is not None + else state.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") + continue + for event in stop_and_flush(): + yield event + stop_message = AnthropicStreamEvent(type="message_stop") + data = stop_message.model_dump_json(exclude_unset=True, exclude_none=True) + yield wrap_data_with_event(data, "message_stop") except Exception as e: logger.exception("Error in message stream converter.") + for event in stop_active_block(): + yield event error_response = AnthropicStreamEvent( type="error", error=AnthropicError( diff --git a/vllm/entrypoints/cohere/serving.py b/vllm/entrypoints/cohere/serving.py index d4bf81d8c2..73fcda2e2f 100644 --- a/vllm/entrypoints/cohere/serving.py +++ b/vllm/entrypoints/cohere/serving.py @@ -321,7 +321,14 @@ class CohereServingChatV2(OpenAIServingChat): case ChatCompletionResponse(): return self._chat_completion_to_v2(generator, request) case _: - return self._chat_completion_stream_to_v2(generator, request) + return generator + + def chat_completion_stream_generator( + self, *args: Any, **kwargs: Any + ) -> AsyncGenerator[str, None]: + return self._chat_completion_stream_to_v2( + self.chat_completion_stream_generator_raw(*args, **kwargs) + ) def to_chat_completion_request( self, request: CohereChatV2Request @@ -1428,10 +1435,9 @@ class CohereServingChatV2(OpenAIServingChat): async def _chat_completion_stream_to_v2( self, - generator: AsyncGenerator[str, None], - request: CohereChatV2Request, + generator: AsyncGenerator[ChatCompletionStreamResponse | ErrorResponse, None], ) -> AsyncGenerator[str, None]: - """Translate an OpenAI-style chat completion SSE stream into Cohere's + """Translate raw OpenAI chat completion chunks into Cohere's v2 stream-event format. Cohere's v2 stream lifecycle is: @@ -1445,19 +1451,9 @@ class CohereServingChatV2(OpenAIServingChat): state = _StreamState() try: - async for item in generator: - if not item.startswith("data:"): - continue - data_str = item[len("data:") :].strip().rstrip("\n") - if not data_str: - continue - if data_str == "[DONE]": - # OpenAI's stream terminator. Fall through to the - # post-loop cleanup so we always emit ``message-end`` - # even if the usage-only chunk was skipped. - break - - chunk = ChatCompletionStreamResponse.model_validate_json(data_str) + async for chunk in generator: + if isinstance(chunk, ErrorResponse): + raise RuntimeError(chunk.error.message) state.last_chunk_id = chunk.id if not state.started: @@ -1514,7 +1510,9 @@ class CohereServingChatV2(OpenAIServingChat): except Exception as exc: logger.exception("Error converting chat completion stream to v2") - if state.started and not state.ended: + if not state.ended: + for ev in self._close_open_blocks(state): + yield ev yield _sse( json.dumps( { @@ -1530,11 +1528,7 @@ class CohereServingChatV2(OpenAIServingChat): yield _DONE_FRAME return - # Normal completion or ``[DONE]``: ensure ``message-end`` is always - # emitted. Upstream may close the stream without sending the final - # usage-only chunk (e.g. on shutdown, or when ``[DONE]`` is the only - # terminator); without this fallback Cohere clients would hang - # waiting for the closing event. + # Ensure message-end is emitted even without a final usage chunk. if state.started and not state.ended: for ev in self._close_open_blocks(state): yield ev diff --git a/vllm/entrypoints/openai/chat_completion/serving.py b/vllm/entrypoints/openai/chat_completion/serving.py index 7e1782e241..aac1ee77ff 100644 --- a/vllm/entrypoints/openai/chat_completion/serving.py +++ b/vllm/entrypoints/openai/chat_completion/serving.py @@ -478,6 +478,44 @@ class OpenAIServingChat(GenerateBaseServing): chat_template_kwargs: dict[str, Any] | None = None, mm_token_counts: dict[str, int] | None = None, ) -> AsyncGenerator[str, None]: + try: + async for chunk in self.chat_completion_stream_generator_raw( + request, + result_generator, + request_id, + model_name, + conversation, + tokenizer, + request_metadata, + chat_template_kwargs=chat_template_kwargs, + mm_token_counts=mm_token_counts, + ): + if isinstance(chunk, ErrorResponse): + data = chunk.model_dump_json() + else: + data = chunk.model_dump_json( + exclude_unset=True, exclude_none=not chunk.choices + ) + yield f"data: {data}\n\n" + except Exception as e: + logger.exception("Error serializing chat completion stream.") + data = self.create_streaming_error_response(e) + yield f"data: {data}\n\n" + yield "data: [DONE]\n\n" + + async def chat_completion_stream_generator_raw( + self, + request: ChatCompletionRequest, + result_generator: AsyncIterator[RequestOutput], + request_id: str, + model_name: str, + conversation: list[ConversationMessage], + tokenizer: TokenizerLike, + request_metadata: RequestResponseMetadata, + chat_template_kwargs: dict[str, Any] | None = None, + mm_token_counts: dict[str, int] | None = None, + ) -> AsyncGenerator[ChatCompletionStreamResponse | ErrorResponse, None]: + """Yield chat chunks or errors without serializing them into SSE frames.""" created_time = int(time.time()) chunk_object_type: Final = "chat.completion.chunk" first_iteration = True @@ -510,9 +548,7 @@ class OpenAIServingChat(GenerateBaseServing): parsers = [None] * num_choices except Exception as e: logger.exception("Error in parser creation.") - data = self.create_streaming_error_response(e) - yield f"data: {data}\n\n" - yield "data: [DONE]\n\n" + yield self.create_error_response(e) return stream_options = request.stream_options @@ -583,8 +619,7 @@ class OpenAIServingChat(GenerateBaseServing): ), ) - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" + yield chunk # Send response to echo the input portion of the # last message @@ -624,8 +659,7 @@ class OpenAIServingChat(GenerateBaseServing): ), ) - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" + yield chunk first_iteration = False for output in res.outputs: @@ -821,8 +855,7 @@ class OpenAIServingChat(GenerateBaseServing): ), ) - data = chunk.model_dump_json(exclude_unset=True) - yield f"data: {data}\n\n" + yield chunk if self._include_reasoning_tokens_details and not include_continuous_usage: for i, parser in enumerate(parsers): @@ -882,10 +915,7 @@ class OpenAIServingChat(GenerateBaseServing): system_fingerprint=self.system_fingerprint, metrics=stream_per_request_metrics, ) - final_usage_data = final_usage_chunk.model_dump_json( - exclude_unset=True, exclude_none=True - ) - yield f"data: {final_usage_data}\n\n" + yield final_usage_chunk # report to FastAPI middleware aggregate usage across all choices num_completion_tokens = sum(previous_num_tokens) @@ -919,13 +949,12 @@ class OpenAIServingChat(GenerateBaseServing): ) except GenerationError as e: - yield f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n" + yield self.create_error_response( + str(e), err_type="InternalServerError", status_code=e.status_code + ) except Exception as e: logger.exception("Error in chat completion stream generator.") - data = self.create_streaming_error_response(e) - yield f"data: {data}\n\n" - # Send the final done message after all response.n are finished - yield "data: [DONE]\n\n" + yield self.create_error_response(e) async def chat_completion_full_generator( self, From ad2688ab302a24661a3eb4a4754ebd5743cba71a Mon Sep 17 00:00:00 2001 From: znzryb Date: Sat, 17 Oct 2026 15:51:36 +0800 Subject: [PATCH 48/49] [Bugfix] Validate draft-token capacity in forward benchmark (#61809) Signed-off-by: znzryb Co-authored-by: Woosuk Kwon --- .agents/skills/model-microbenchmark/SKILL.md | 3 +++ tests/benchmarks/test_forward_cli.py | 13 +++++++++++++ vllm/benchmarks/forward/main.py | 8 ++++++++ 3 files changed, 24 insertions(+) diff --git a/.agents/skills/model-microbenchmark/SKILL.md b/.agents/skills/model-microbenchmark/SKILL.md index 7e48325dcb..7f18c2bb87 100644 --- a/.agents/skills/model-microbenchmark/SKILL.md +++ b/.agents/skills/model-microbenchmark/SKILL.md @@ -31,6 +31,9 @@ Requests: `[N*]Q[@K][c][:D]`, comma-separated, per DP rank. | `2048@10240` | prefill after an 8192-token prefix-cache hit | | `2048@10240c` | chunked prefill: the prefix came from earlier chunks | +Verify cases require `--speculative-config` with `num_speculative_tokens` +at least as large as the case's draft count; insufficient capacity is rejected. + Hit vs chunk only differs on models with SWA bounded replay (DSV4.1): a hit also recomputes the last window (128 tokens) of the prefix. diff --git a/tests/benchmarks/test_forward_cli.py b/tests/benchmarks/test_forward_cli.py index 34f8e0c355..d44c6b9466 100644 --- a/tests/benchmarks/test_forward_cli.py +++ b/tests/benchmarks/test_forward_cli.py @@ -9,6 +9,7 @@ import numpy as np import pytest from vllm.benchmarks.forward.batch import build_batch +from vllm.benchmarks.forward.main import run_case from vllm.benchmarks.forward.shapes import jitter_kv_lens, parse_case from vllm.platforms import current_platform @@ -84,6 +85,18 @@ def test_jitter_kv_lens(): jitter_kv_lens(case, 0, 1000, 0) +@pytest.mark.parametrize("capacity", [0, 2]) +def test_bench_forward_rejects_oversized_verify(capacity): + """Reject drafts beyond the runner's buffer before launching CUDA kernels.""" + runner = SimpleNamespace(num_speculative_steps=capacity) + case = parse_case("verify=2*4@64:3") + with pytest.raises( + ValueError, + match=rf"Case 'verify' needs 3 draft tokens.*supports {capacity}", + ): + run_case(runner, case, SimpleNamespace()) + + @pytest.mark.benchmark # TODO(woosuk): Support ROCm and re-enable this test there. @pytest.mark.skipif( diff --git a/vllm/benchmarks/forward/main.py b/vllm/benchmarks/forward/main.py index ff709562a8..d7841c641e 100644 --- a/vllm/benchmarks/forward/main.py +++ b/vllm/benchmarks/forward/main.py @@ -120,6 +120,14 @@ def _make_engine_args(args: argparse.Namespace, cases: list[BenchCase]): def run_case(runner, case: BenchCase, args: argparse.Namespace) -> dict[str, Any]: + required_drafts = max((r.num_draft_tokens for r in case.requests), default=0) + if required_drafts > runner.num_speculative_steps: + raise ValueError( + f"Case {case.name!r} needs {required_drafts} draft tokens per request, " + f"but the runner supports {runner.num_speculative_steps}. " + "Set --speculative-config with a sufficient num_speculative_tokens." + ) + from vllm.benchmarks.forward.batch import build_batch from vllm.benchmarks.forward.capture import ForwardBench from vllm.benchmarks.forward.kv_init import touched_kv_views, zero_kv_blocks From 08a65e287e55099255ce178c5b5f58ce42dfd461 Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Sat, 17 Oct 2026 00:52:17 -0700 Subject: [PATCH 49/49] [Bugfix][Kernel] hadacore: don't load past the input on a partial last warp (#61944) Signed-off-by: khluu Co-authored-by: Claude Opus 5.5 Co-authored-by: Jee Jee Li Co-authored-by: vllm-agent --- .../hadamard/hadacore/hadamard_transform_cuda.cu | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu index 2575b20b8e..e86c518204 100644 --- a/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu +++ b/csrc/quantization/hadamard/hadacore/hadamard_transform_cuda.cu @@ -135,6 +135,13 @@ hadamard_transform_kernel(b16* a, b16* out, int total_num_chunks) { #pragma unroll for (int64_t k = 0; k < num_chunks; k++) { + // The last warp may own fewer chunks than num_chunks: loading the + // rest would read past the end of `a`. The cp.async.wait_group counts + // below already assume only real_num_chunks groups were committed. + if constexpr(enable_mask) { + if (k >= real_num_chunks) + break; + } size_t shared_ptr = __cvta_generic_to_shared(b_frag_ptr); #if (__CUDA_ARCH__ >= 900) // SM90 asm volatile(