From 826bb4088f4671fe79e01183e0301f79134662c1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 12 Aug 2026 13:46:21 +0000 Subject: [PATCH 01/52] fix(structured-output): preserve grammar bitmask source widths Record scheduler-side speculative widths in GrammarOutput so worker-side draft trimming cannot shift flattened grammar masks onto later requests. Destination logits continue to use the worker-visible width, while source offsets use the serialized scheduler width. Validated with focused unit coverage and a 160-request concurrent DeepSeek V4 structured-output workload. --- tests/v1/structured_output/test_utils.py | 56 ++++++++++++++++++++++++ vllm/v1/core/sched/output.py | 4 ++ vllm/v1/core/sched/scheduler.py | 10 ++++- vllm/v1/structured_output/utils.py | 18 ++++++-- vllm/v1/worker/gpu/warmup.py | 4 +- 5 files changed, 86 insertions(+), 6 deletions(-) diff --git a/tests/v1/structured_output/test_utils.py b/tests/v1/structured_output/test_utils.py index c026ab0e4e78..42e97e48f1d9 100644 --- a/tests/v1/structured_output/test_utils.py +++ b/tests/v1/structured_output/test_utils.py @@ -1,8 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + +import numpy as np import pytest +import torch +from vllm.v1.core.sched.output import GrammarOutput +from vllm.v1.structured_output import utils from vllm.v1.structured_output.backend_xgrammar import ( has_xgrammar_unsupported_json_features, ) @@ -104,3 +110,53 @@ def test_supported_json_features(supported_schema): assert not has_xgrammar_unsupported_json_features(supported_schema), ( "Schema should be supported" ) + + +def test_apply_grammar_bitmask_preserves_source_offsets_after_draft_trimming( + monkeypatch, +): + """A trimmed request must not shift another request's grammar rows. + + The scheduler serializes masks at its scheduled speculative width. A worker + may trim grammar-invalid drafts before applying those masks, so source and + destination offsets must be advanced with their respective widths. + """ + scheduler_output = SimpleNamespace( + scheduled_spec_decode_tokens={ + "trimmed-request": [1], + "full-request": [2, 3, 4], + } + ) + grammar_output = GrammarOutput( + structured_output_request_ids=["trimmed-request", "full-request"], + grammar_bitmask=np.array( + [[10], [11], [12], [13], [20], [21], [22], [23]], + dtype=np.int32, + ), + num_spec_tokens=[3, 3], + ) + input_batch = SimpleNamespace(req_ids=["trimmed-request", "full-request"]) + logits = torch.zeros((6, 32)) + applied_bitmask = None + + def capture_bitmask(logits, bitmask, indices): + nonlocal applied_bitmask + applied_bitmask = bitmask.clone() + assert indices is None + + monkeypatch.setattr( + utils, + "xgr", + SimpleNamespace(apply_token_bitmask_inplace=capture_bitmask), + ) + monkeypatch.setattr(utils, "PIN_MEMORY", False) + + utils.apply_grammar_bitmask( + scheduler_output, + grammar_output, + input_batch, + logits, + ) + + assert applied_bitmask is not None + assert applied_bitmask[:, 0].tolist() == [10, 13, 20, 21, 22, 23] diff --git a/vllm/v1/core/sched/output.py b/vllm/v1/core/sched/output.py index e0d3c2a1a6cc..da11eb1dddfb 100644 --- a/vllm/v1/core/sched/output.py +++ b/vllm/v1/core/sched/output.py @@ -295,3 +295,7 @@ class GrammarOutput: structured_output_request_ids: list[str] # Bitmask ordered as structured_output_request_ids. grammar_bitmask: "npt.NDArray[np.int32]" + # Number of speculative rows represented in grammar_bitmask for each + # structured output request. Worker-side draft trimming may reduce the + # number of logits without changing this compact source layout. + num_spec_tokens: list[int] diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ae08df9c35b2..5991074d0092 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1691,7 +1691,15 @@ def get_grammar_bitmask( structured_output_request_ids, scheduler_output.scheduled_spec_decode_tokens, ) - return GrammarOutput(structured_output_request_ids, bitmask) + num_spec_tokens = [ + len(scheduler_output.scheduled_spec_decode_tokens.get(req_id, ())) + for req_id in structured_output_request_ids + ] + return GrammarOutput( + structured_output_request_ids, + bitmask, + num_spec_tokens, + ) def update_from_output( self, diff --git a/vllm/v1/structured_output/utils.py b/vllm/v1/structured_output/utils.py index 0629a6d2e0f6..c2b445a52dbb 100644 --- a/vllm/v1/structured_output/utils.py +++ b/vllm/v1/structured_output/utils.py @@ -131,14 +131,24 @@ def apply_grammar_bitmask( ) sorted_bitmask = sorted_bitmask_tensor.numpy() cumulative_index = 0 - for req_id in grammar_output.structured_output_request_ids: - num_spec_tokens = len(spec_tokens.get(req_id, ())) + for req_id, num_grammar_spec_tokens in zip( + grammar_output.structured_output_request_ids, + grammar_output.num_spec_tokens, + strict=True, + ): + num_worker_spec_tokens = len(spec_tokens.get(req_id, ())) + assert num_worker_spec_tokens <= num_grammar_spec_tokens if (logit_idx := struct_out_req_batch_indices.get(req_id)) is not None: - for i in range(1 + num_spec_tokens): + for i in range(num_worker_spec_tokens): bitmask_index = logit_idx + i sorted_bitmask[bitmask_index] = grammar_bitmask[cumulative_index + i] out_indices.append(bitmask_index) - cumulative_index += 1 + num_spec_tokens + bonus_index = logit_idx + num_worker_spec_tokens + sorted_bitmask[bonus_index] = grammar_bitmask[ + cumulative_index + num_grammar_spec_tokens + ] + out_indices.append(bonus_index) + cumulative_index += 1 + num_grammar_spec_tokens # Copy async to device. grammar_bitmask = sorted_bitmask_tensor.to(logits.device, non_blocking=True) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index ab745d70a295..6dba53f2f324 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -367,7 +367,9 @@ def _alloc_blocks(num_blocks: int) -> list[int]: (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32 ) grammar_output = GrammarOutput( - structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask + structured_output_request_ids=req_ids, + grammar_bitmask=grammar_bitmask, + num_spec_tokens=[0] * len(req_ids), ) worker_sample_tokens(grammar_output) From b9eea38f408f88e39d9a145bd5324cdd0b80e234 Mon Sep 17 00:00:00 2001 From: guptaishaan Date: Fri, 31 Jul 2026 14:33:47 -0700 Subject: [PATCH 02/52] [Bugfix] Stream Kimi K3 tool-call arguments incrementally KimiK3ToolParser.extract_tool_calls_streaming matched calls with _call_re, which requires the closing <|close|>call<|sep|> marker. Until that marker arrived nothing was emitted for the call, so a long tool call produced no SSE deltas for the whole generation and then dumped the entire arguments JSON in one delta. Track the call from its <|open|>call ...<|sep|> marker instead. The name goes out immediately, and _partial_arguments serializes the arguments seen so far as a prefix of the final JSON, so each step can stream the difference against what it already sent. String argument bodies are raw text, so they are forwarded as they arrive with a trailing partial close marker held back; other types still need the whole literal to decode and are held until their block closes. The concatenated deltas are byte-identical to the non-streaming extract_tool_calls output. Signed-off-by: guptaishaan --- tests/tool_use/test_kimi_k3_tool_parser.py | 60 ++++++++- vllm/tool_parsers/kimi_k3_tool_parser.py | 135 ++++++++++++++++----- 2 files changed, 159 insertions(+), 36 deletions(-) diff --git a/tests/tool_use/test_kimi_k3_tool_parser.py b/tests/tool_use/test_kimi_k3_tool_parser.py index 118c51db574e..dab4f5227544 100644 --- a/tests/tool_use/test_kimi_k3_tool_parser.py +++ b/tests/tool_use/test_kimi_k3_tool_parser.py @@ -361,9 +361,63 @@ def test_streaming_split_markers_do_not_leak(): assert content == "Hi" assert OPEN not in content assert SEP not in content - assert len(tool_deltas) == 1 - assert tool_deltas[0].function.name == "calc" - assert json.loads(tool_deltas[0].function.arguments) == {"x": 1} + assert [tool_call.function.name for tool_call in tool_deltas if tool_call.id] == [ + "calc" + ] + arguments = "".join(tool_call.function.arguments or "" for tool_call in tool_deltas) + assert json.loads(arguments) == {"x": 1} + + +def test_streaming_emits_argument_text_as_it_arrives(): + """A long string argument must stream, not land in one delta at the close.""" + parser = KimiK3ToolParser(DummyTokenizer()) + request = _request() + value = "word " * 200 + body_chunks = [value[i : i + 5] for i in range(0, len(value), 5)] + chunks = [ + f"{OPEN}tools{SEP}", + f'{OPEN}call tool="write_file" index="1"{SEP}', + f'{OPEN}argument key="content" type="string"{SEP}', + *body_chunks, + f"{CLOSE}argument{SEP}", + f"{CLOSE}call{SEP}", + ] + previous_text = "" + previous_ids: list[int] = [] + messages: list[DeltaMessage] = [] + + for i, chunk in enumerate(chunks, start=1): + current_text = previous_text + chunk + current_ids = previous_ids + [i] + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=chunk, + previous_token_ids=previous_ids, + current_token_ids=current_ids, + delta_token_ids=[i], + request=request, + ) + if delta is not None: + messages.append(delta) + previous_text = current_text + previous_ids = current_ids + + tool_deltas = [ + tool_call for message in messages for tool_call in (message.tool_calls or []) + ] + + # the name is announced once, up front, and every body chunk moves the stream + assert [tool_call.function.name for tool_call in tool_deltas if tool_call.id] == [ + "write_file" + ] + assert len(tool_deltas) >= len(body_chunks) + assert all(tool_call.index == 0 for tool_call in tool_deltas) + + arguments = "".join(tool_call.function.arguments or "" for tool_call in tool_deltas) + assert json.loads(arguments) == {"content": value} + non_streamed = parser.extract_tool_calls(previous_text, request) + assert arguments == non_streamed.tool_calls[0].function.arguments def test_tool_call_ids_are_unique_across_messages(): diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py index 9f688f2270e4..656f0bf18164 100644 --- a/vllm/tool_parsers/kimi_k3_tool_parser.py +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -37,6 +37,7 @@ import regex as re from openai.types.responses import ToolChoiceFunction +from vllm.entrypoints.chat_utils import make_tool_call_id from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, @@ -86,6 +87,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): self.tools_close = "<|close|>tools<|sep|>" self.response_open = "<|open|>response<|sep|>" self.response_close = "<|close|>response<|sep|>" + self.argument_close = "<|close|>argument<|sep|>" # Regexes operate on detokenized text. The XTML markers reach us as the # literal strings <|open|>/<|close|>/<|sep|>. adjust_request keeps them @@ -127,6 +129,15 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): + _S, re.DOTALL, ) + # Half-open variants: streaming needs to see a call/argument as soon as + # its opening marker lands, before the matching close marker exists. + self._call_open_re = re.compile( + _O + r"\s*call\s+(?P" + _TEXT_UNTIL_SEP + r")" + _S + ) + self._call_close_re = re.compile(_C + r"\s*call\s*" + _S) + self._arg_open_re = re.compile( + _O + r"\s*argument\s+(?P" + _TEXT_UNTIL_SEP + r")" + _S + ) # attr segment: key="value" (value already escaped on the encode side) self._attr_re = re.compile(r'(?P\w+)="(?P[^"]*)"') self._response_re = re.compile( @@ -136,7 +147,8 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): # streaming state self._sent_content_idx = 0 - self._sent_tool_call_count = 0 + # arguments already streamed, per tool-call index + self._streamed_args: list[str] = [] if not self.model_tokenizer: raise ValueError( @@ -218,6 +230,17 @@ def _decode_call(self, attrs: str, body: str) -> ToolCall | None: """ call_attrs = self._attrs(attrs) tool_name = call_attrs.get("tool", "") + if not tool_name: + return None + return ToolCall( + type="function", + function=FunctionCall( + name=tool_name, + arguments=json.dumps(self._decode_arguments(body), ensure_ascii=False), + ), + ) + + def _decode_arguments(self, body: str) -> dict: arguments: dict = {} for arg_match in self._arg_re.finditer(body): arg_attrs = self._attrs(arg_match["attrs"]) @@ -231,15 +254,36 @@ def _decode_call(self, attrs: str, body: str) -> ToolCall | None: arguments[key] = json.loads(raw_value) except json.JSONDecodeError: arguments[key] = raw_value - if not tool_name: - return None - return ToolCall( - type="function", - function=FunctionCall( - name=tool_name, - arguments=json.dumps(arguments, ensure_ascii=False), - ), - ) + return arguments + + def _partial_arguments(self, body: str) -> str: + """Serialize the arguments of a ``call`` block that has not closed yet. + + The result is always a prefix of what the finished call serializes to, + so the caller can stream the difference against what it already sent. + String values are emitted raw by the template, so their text is + forwarded as it arrives (minus a trailing partial close marker); other + types need the whole literal to decode and are held until their block + closes. + """ + closed_end = 0 + for arg_match in self._arg_re.finditer(body): + closed_end = arg_match.end() + arguments = self._decode_arguments(body[:closed_end]) + + m_open = self._arg_open_re.search(body, closed_end) + if m_open is None: + # drop the closing "}" of the object + return json.dumps(arguments, ensure_ascii=False)[:-1] + + arg_attrs = self._attrs(m_open["attrs"]) + if arg_attrs.get("type", "string") != "string": + return json.dumps(arguments, ensure_ascii=False)[:-1] + raw_value = body[m_open.end() :] + held = _partial_tag_overlap(raw_value, self.argument_close) + arguments[arg_attrs.get("key", "")] = raw_value[: len(raw_value) - held] + # drop the closing quote of the open value and the "}" of the object + return json.dumps(arguments, ensure_ascii=False)[:-2] def _strip_response_content(self, text: str) -> str | None: """Strip XTML response/message markers from generated response text. @@ -361,35 +405,60 @@ def extract_tool_calls_streaming( delta_token_ids: Sequence[int], request: ChatCompletionRequest, ) -> DeltaMessage | None: - # Conservative streaming: stream unwrapped response-channel text, then - # buffer tool calls and emit each call once its block closes. + # Stream unwrapped response-channel text, then stream each tool call as + # it is generated: the open marker carries the name, and argument text + # is forwarded as it arrives rather than buffered until the call closes. content = self._extract_response_content(current_text) - # tools channel is open: parse fully-closed calls we have not emitted yet m_tools = self._tools_open_re.search(current_text) if m_tools is None: return DeltaMessage(content=content) if content else None section = current_text[m_tools.end() :] - calls = [ - tc - for m in self._call_re.finditer(section) - if (tc := self._decode_call(m["attrs"], m["body"])) is not None - ] - if len(calls) <= self._sent_tool_call_count: + opens = list(self._call_open_re.finditer(section)) + deltas: list[DeltaToolCall] = [] + index = 0 + for i, m_open in enumerate(opens): + end = opens[i + 1].start() if i + 1 < len(opens) else len(section) + body = section[m_open.end() : end] + name = self._attrs(m_open["attrs"]).get("tool", "") + if not name: + # an empty/garbage block is dropped, as in _decode_call + continue + m_close = self._call_close_re.search(body) + if m_close is None: + arguments = self._partial_arguments(body) + else: + arguments = json.dumps( + self._decode_arguments(body[: m_close.start()]), ensure_ascii=False + ) + + if index == len(self._streamed_args): + self._streamed_args.append(arguments) + deltas.append( + DeltaToolCall( + index=index, + id=make_tool_call_id(), + type="function", + function=DeltaFunctionCall( + name=name, arguments=arguments + ).model_dump(exclude_none=True), + ) + ) + else: + sent = self._streamed_args[index] + if arguments != sent and arguments.startswith(sent): + self._streamed_args[index] = arguments + deltas.append( + DeltaToolCall( + index=index, + function=DeltaFunctionCall( + arguments=arguments[len(sent) :] + ).model_dump(exclude_none=True), + ) + ) + index += 1 + + if not deltas: return DeltaMessage(content=content) if content else None - new = calls[self._sent_tool_call_count :] - - deltas = [ - DeltaToolCall( - index=self._sent_tool_call_count + i, - id=tc.id, - type="function", - function=DeltaFunctionCall( - name=tc.function.name, arguments=tc.function.arguments - ).model_dump(exclude_none=True), - ) - for i, tc in enumerate(new) - ] - self._sent_tool_call_count = len(calls) return DeltaMessage(content=content, tool_calls=deltas) From efbaa9581032d01e9ce111849ab2d92a26b16757 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 17 Aug 2026 13:10:32 +0000 Subject: [PATCH 03/52] tools: preserve partial Kimi XTML close markers Withhold whitespace-tolerant argument-close fragments until they form a complete XTML marker. This keeps streamed JSON argument deltas prefix-stable for every marker form accepted by the parser. Co-authored-by: OpenAI Codex --- tests/tool_use/test_kimi_k3_tool_parser.py | 59 ++++++++++++++++++++++ vllm/tool_parsers/kimi_k3_tool_parser.py | 35 ++++++++++--- 2 files changed, 87 insertions(+), 7 deletions(-) diff --git a/tests/tool_use/test_kimi_k3_tool_parser.py b/tests/tool_use/test_kimi_k3_tool_parser.py index dab4f5227544..8c0c1e880344 100644 --- a/tests/tool_use/test_kimi_k3_tool_parser.py +++ b/tests/tool_use/test_kimi_k3_tool_parser.py @@ -385,8 +385,15 @@ def test_streaming_emits_argument_text_as_it_arrives(): previous_text = "" previous_ids: list[int] = [] messages: list[DeltaMessage] = [] + arguments_before_close = "" for i, chunk in enumerate(chunks, start=1): + if chunk == f"{CLOSE}argument{SEP}": + arguments_before_close = "".join( + tool_call.function.arguments or "" + for message in messages + for tool_call in (message.tool_calls or []) + ) current_text = previous_text + chunk current_ids = previous_ids + [i] delta = parser.extract_tool_calls_streaming( @@ -413,6 +420,10 @@ def test_streaming_emits_argument_text_as_it_arrives(): ] assert len(tool_deltas) >= len(body_chunks) assert all(tool_call.index == 0 for tool_call in tool_deltas) + assert ( + arguments_before_close + == json.dumps({"content": value}, ensure_ascii=False)[:-2] + ) arguments = "".join(tool_call.function.arguments or "" for tool_call in tool_deltas) assert json.loads(arguments) == {"content": value} @@ -420,6 +431,54 @@ def test_streaming_emits_argument_text_as_it_arrives(): assert arguments == non_streamed.tool_calls[0].function.arguments +def test_streaming_holds_whitespace_tolerant_argument_close_fragments(): + parser = KimiK3ToolParser(DummyTokenizer()) + request = _request() + chunks = [ + f"{OPEN}tools{SEP}", + f'{OPEN}call tool="write_file" index="1"{SEP}', + f'{OPEN}argument key="content" type="string"{SEP}', + "payload", + f"{CLOSE} arg", + "ument ", + "<|sep", + "|>", + f"{CLOSE}call{SEP}", + ] + previous_text = "" + previous_ids: list[int] = [] + streamed_arguments = "" + partial_close_snapshots: list[str] = [] + + for i, chunk in enumerate(chunks, start=1): + current_text = previous_text + chunk + current_ids = previous_ids + [i] + delta = parser.extract_tool_calls_streaming( + previous_text=previous_text, + current_text=current_text, + delta_text=chunk, + previous_token_ids=previous_ids, + current_token_ids=current_ids, + delta_token_ids=[i], + request=request, + ) + if delta is not None: + streamed_arguments += "".join( + tool_call.function.arguments or "" + for tool_call in (delta.tool_calls or []) + ) + if 5 <= i <= 7: + partial_close_snapshots.append(streamed_arguments) + previous_text = current_text + previous_ids = current_ids + + expected_prefix = json.dumps({"content": "payload"}, ensure_ascii=False)[:-2] + assert partial_close_snapshots == [expected_prefix] * 3 + assert json.loads(streamed_arguments) == {"content": "payload"} + non_streamed = parser.extract_tool_calls(previous_text, request) + assert streamed_arguments == non_streamed.tool_calls[0].function.arguments + + def test_tool_call_ids_are_unique_across_messages(): output = _tools(_call("calc", 1)) diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py index 656f0bf18164..7dea33d589ad 100644 --- a/vllm/tool_parsers/kimi_k3_tool_parser.py +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -70,6 +70,25 @@ def _partial_tag_overlap(text: str, tag: str) -> int: return 0 +def _partial_pattern_overlap(text: str, pattern: re.Pattern) -> int: + """Return a trailing prefix accepted by a partial regular expression. + + Args: + text: Generated text that may end inside a structural marker. + pattern: Complete structural-marker expression. + + Returns: + The number of trailing characters that form an incomplete match. + """ + candidate = text.rfind("<") + while candidate >= 0: + match = pattern.fullmatch(text[candidate:], partial=True) + if match is not None and match.partial: + return len(text) - candidate + candidate = text.rfind("<", 0, candidate) + return 0 + + class KimiK3ToolParser(ToolParser): supports_required_and_named = False # Enables the vLLM-side XTML structural tag builder @@ -138,6 +157,7 @@ def __init__(self, tokenizer: TokenizerLike, tools: list[Tool] | None = None): self._arg_open_re = re.compile( _O + r"\s*argument\s+(?P" + _TEXT_UNTIL_SEP + r")" + _S ) + self._arg_close_re = re.compile(_C + r"\s*argument\s*" + _S) # attr segment: key="value" (value already escaped on the encode side) self._attr_re = re.compile(r'(?P\w+)="(?P[^"]*)"') self._response_re = re.compile( @@ -259,12 +279,13 @@ def _decode_arguments(self, body: str) -> dict: def _partial_arguments(self, body: str) -> str: """Serialize the arguments of a ``call`` block that has not closed yet. - The result is always a prefix of what the finished call serializes to, - so the caller can stream the difference against what it already sent. - String values are emitted raw by the template, so their text is - forwarded as it arrives (minus a trailing partial close marker); other - types need the whole literal to decode and are held until their block - closes. + Args: + body: The incomplete XTML call body. + + Returns: + A prefix of the completed call's JSON arguments. String values can + stream as generated. Other argument types remain buffered until + their complete JSON literal is available. """ closed_end = 0 for arg_match in self._arg_re.finditer(body): @@ -280,7 +301,7 @@ def _partial_arguments(self, body: str) -> str: if arg_attrs.get("type", "string") != "string": return json.dumps(arguments, ensure_ascii=False)[:-1] raw_value = body[m_open.end() :] - held = _partial_tag_overlap(raw_value, self.argument_close) + held = _partial_pattern_overlap(raw_value, self._arg_close_re) arguments[arg_attrs.get("key", "")] = raw_value[: len(raw_value) - held] # drop the closing quote of the open value and the "}" of the object return json.dumps(arguments, ensure_ascii=False)[:-2] From 9972d3c7ad3e90b0652d5a1a0f222a9f706d9fe6 Mon Sep 17 00:00:00 2001 From: myshytf Date: Sun, 16 Aug 2026 02:33:29 +0900 Subject: [PATCH 04/52] fix(cache): preserve Mamba CoW after external hits --- .../test_partial_prefix_cache_hits.py | 72 +++++++++++++++++++ vllm/v1/core/single_type_kv_cache_manager.py | 1 + 2 files changed, 73 insertions(+) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index 5f549ec8df7e..b193f43abd6c 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -333,6 +333,78 @@ def test_hybrid_mamba_partial_tail_owner_uses_cow_on_continue(): assert moved[0].block_hash_num_tokens == 6 +def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): + """An external mid-block hit must become a running request even when its + first continuation does not need another Mamba block.""" + hash_block_size = 2 + mamba_block_size = 4 * hash_block_size + kv_cache_config = KVCacheConfig( + num_blocks=32, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=hash_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + manager = make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + hash_block_size=hash_block_size, + ) + + request = make_request("0", [0] * 15, hash_block_size, sha256) + loaded_blocks = manager.allocate_slots( + request, + num_new_tokens=0, + num_external_computed_tokens=10, + delay_cache_blocks=True, + ) + assert loaded_blocks is not None + + request.num_computed_tokens = 10 + first_step_blocks = manager.allocate_slots(request, num_new_tokens=4) + assert first_step_blocks is not None + + source_block_id = manager.get_blocks("0").get_block_ids()[1][1] + partial_hash = request.block_hashes[14 // hash_block_size - 1] + partial_block = manager.block_pool.get_cached_block( + partial_hash, kv_cache_group_ids=[1] + ) + assert partial_block is not None + assert partial_block[0].block_id == source_block_id + + request.num_computed_tokens = 14 + continuation_blocks = manager.allocate_slots(request, num_new_tokens=1) + assert continuation_blocks is not None + + assert continuation_blocks.get_block_ids()[1] == [] + assert manager.get_blocks("0").get_block_ids()[1][1] == source_block_id + copies, _ = manager.take_kv_cache_block_copies() + cow_copy = next(c for c in copies if c.src_block_id == source_block_id) + assert cow_copy.dst_block_id != source_block_id + + moved = manager.block_pool.get_cached_block(partial_hash, kv_cache_group_ids=[1]) + assert moved is not None + assert moved[0].block_id == cow_copy.dst_block_id + + def test_take_partial_tail_offloads_returns_cow_target(): """The connector offload hand-off exposes the mamba CoW *target* block Y (the durable boundary state), not the overwritten source X, and only at diff --git a/vllm/v1/core/single_type_kv_cache_manager.py b/vllm/v1/core/single_type_kv_cache_manager.py index a334a8dde9da..686cae8257e6 100644 --- a/vllm/v1/core/single_type_kv_cache_manager.py +++ b/vllm/v1/core/single_type_kv_cache_manager.py @@ -1597,6 +1597,7 @@ def allocate_new_blocks( # `num_required_blocks` might be less than `len(req_blocks)` if blocks are # over-allocated at last round. if num_required_blocks <= len(req_blocks) and not has_partial_hit: + self._allocated_block_reqs.add(request_id) return [] else: prev_block_len = len(req_blocks) From 990a63debf7201ae5446436a712609fcc8ffaef7 Mon Sep 17 00:00:00 2001 From: myshytf Date: Sun, 16 Aug 2026 02:40:10 +0900 Subject: [PATCH 05/52] test(cache): assert external-hit continuation precondition --- tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index b193f43abd6c..bc4f978a8ef8 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -381,6 +381,7 @@ def test_external_mamba_hit_same_block_uses_running_cow_on_continue(): request.num_computed_tokens = 10 first_step_blocks = manager.allocate_slots(request, num_new_tokens=4) assert first_step_blocks is not None + assert first_step_blocks.get_block_ids()[1] == [] source_block_id = manager.get_blocks("0").get_block_ids()[1][1] partial_hash = request.block_hashes[14 // hash_block_size - 1] From 3ebc8967dab28a9df103d8c66bff9688525263bc Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 17 Aug 2026 14:02:09 +0000 Subject: [PATCH 06/52] Ignore Kimi tool calls after tools section close --- tests/tool_use/test_kimi_k3_tool_parser.py | 19 +++++++++++++++++++ vllm/tool_parsers/kimi_k3_tool_parser.py | 8 +++++++- 2 files changed, 26 insertions(+), 1 deletion(-) diff --git a/tests/tool_use/test_kimi_k3_tool_parser.py b/tests/tool_use/test_kimi_k3_tool_parser.py index 8c0c1e880344..47bd9b374acf 100644 --- a/tests/tool_use/test_kimi_k3_tool_parser.py +++ b/tests/tool_use/test_kimi_k3_tool_parser.py @@ -479,6 +479,25 @@ def test_streaming_holds_whitespace_tolerant_argument_close_fragments(): assert streamed_arguments == non_streamed.tool_calls[0].function.arguments +def test_streaming_ignores_call_shaped_text_after_tools_close(): + parser = KimiK3ToolParser(DummyTokenizer()) + request = _request() + output = _tools() + _call("calc", 1, _arg("x", "number", "1")) + + delta = parser.extract_tool_calls_streaming( + previous_text="", + current_text=output, + delta_text=output, + previous_token_ids=[], + current_token_ids=[1], + delta_token_ids=[1], + request=request, + ) + + assert delta is None + assert parser.extract_tool_calls(output, request).tools_called is False + + def test_tool_call_ids_are_unique_across_messages(): output = _tools(_call("calc", 1)) diff --git a/vllm/tool_parsers/kimi_k3_tool_parser.py b/vllm/tool_parsers/kimi_k3_tool_parser.py index 7dea33d589ad..e1fc5a156996 100644 --- a/vllm/tool_parsers/kimi_k3_tool_parser.py +++ b/vllm/tool_parsers/kimi_k3_tool_parser.py @@ -435,7 +435,13 @@ def extract_tool_calls_streaming( if m_tools is None: return DeltaMessage(content=content) if content else None - section = current_text[m_tools.end() :] + section_start = m_tools.end() + m_tools_close = self._tools_close_re.search(current_text, section_start) + section = current_text[ + section_start : ( + len(current_text) if m_tools_close is None else m_tools_close.start() + ) + ] opens = list(self._call_open_re.finditer(section)) deltas: list[DeltaToolCall] = [] index = 0 From c805ebd0896ccfbd2569bc0b2a7944d3282106ff Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Mon, 17 Aug 2026 14:11:21 +0000 Subject: [PATCH 07/52] fix(dspark): preserve draft graph capture contract --- .../test_dspark_cudagraph_contract.py | 43 +++++++++++++++++++ .../gpu/spec_decode/dspark/speculator.py | 1 + 2 files changed, 44 insertions(+) create mode 100644 tests/v1/spec_decode/test_dspark_cudagraph_contract.py diff --git a/tests/v1/spec_decode/test_dspark_cudagraph_contract.py b/tests/v1/spec_decode/test_dspark_cudagraph_contract.py new file mode 100644 index 000000000000..c1c1c444e52f --- /dev/null +++ b/tests/v1/spec_decode/test_dspark_cudagraph_contract.py @@ -0,0 +1,43 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace +from unittest.mock import Mock + +import torch + +from vllm.config.compilation import CUDAGraphMode +from vllm.v1.worker.gpu.spec_decode.dspark.speculator import DSparkSpeculator + + +def test_dspark_generate_draft_accepts_dflash_capture_contract(): + head_hidden = torch.randn(10, 8) + speculator = SimpleNamespace( + num_query_per_req=5, + capacity_activation_batch_size=1, + _speculative_steps_for_query_len=Mock(return_value=5), + _run_model=Mock(return_value=head_hidden), + _sample_sequential=Mock(), + ) + + DSparkSpeculator._generate_draft( + speculator, + num_reqs=2, + num_tokens_padded=10, + attn_metadata=None, + slot_mappings=None, + num_tokens_across_dp=None, + cudagraph_runtime_mode=CUDAGraphMode.FULL, + is_profile=True, + num_query_per_req=5, + capture_only=True, + ) + + speculator._sample_sequential.assert_called_once_with( + 2, + head_hidden, + 5, + 5, + is_profile=True, + use_capacity=True, + ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index 0d126f672c89..8ae79d4079e1 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -424,6 +424,7 @@ def _generate_draft( cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE, is_profile: bool = False, num_query_per_req: int | None = None, + capture_only: bool = False, ) -> None: if num_query_per_req is None: num_query_per_req = self.num_query_per_req From 3938209fa7b5573ace70173730390a3d098ffc94 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 14 Aug 2026 13:17:53 +0000 Subject: [PATCH 08/52] spec_decode: preserve DFlash-family checkpoint semantics Co-authored-by: Codex --- tests/v1/spec_decode/test_dflash_causality.py | 104 ++++++++++++++++++ vllm/model_executor/models/qwen3_dflash.py | 30 +++++ .../v1/worker/gpu/spec_decode/dflash/utils.py | 8 +- .../v1/worker/gpu/spec_decode/dspark/utils.py | 4 +- 4 files changed, 144 insertions(+), 2 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_causality.py b/tests/v1/spec_decode/test_dflash_causality.py index 02e2a0bdbec2..015281027894 100644 --- a/tests/v1/spec_decode/test_dflash_causality.py +++ b/tests/v1/spec_decode/test_dflash_causality.py @@ -10,11 +10,13 @@ from types import SimpleNamespace import pytest +import torch.nn as nn from vllm.model_executor.models.qwen3_dflash import ( _dflash_layer_causal, _get_dflash_fc_input_size, dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, ) from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( get_eagle3_aux_layers_from_config, @@ -89,3 +91,105 @@ def test_eagle_aux_layers_preserves_legacy_layer_ids(config_name): assert get_eagle3_aux_layers_from_config(vllm_config.speculative_config) == tuple( layer_ids ) + + +class _TargetRotaryModule(nn.Module): + def __init__(self, is_neox_style: bool): + super().__init__() + self.is_neox_style = is_neox_style + + +class _TargetModel(nn.Module): + def __init__(self, is_neox_style: bool): + super().__init__() + self.rotary = _TargetRotaryModule(is_neox_style) + + +@pytest.mark.parametrize("is_neox_style", [False, True]) +def test_dflash_target_rope_layout_is_discovered(is_neox_style): + target = _TargetModel(is_neox_style) + + assert dflash_target_rope_is_neox_style(target) is is_neox_style + + +def test_dflash_loader_propagates_target_rope_layout(monkeypatch): + """DFlash configures the target rotary layout before draft construction.""" + from vllm.v1.worker.gpu.spec_decode.dflash import utils as loader_module + + draft_hf_config = SimpleNamespace( + num_hidden_layers=1, + layer_types=["sliding_attention"], + dflash_config={"causal": True}, + ) + speculative_config = SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=draft_hf_config), + attention_backend=None, + kv_cache_dtype=None, + draft_load_config=None, + ) + vllm_config = SimpleNamespace( + speculative_config=speculative_config, + attention_config=SimpleNamespace(), + cache_config=SimpleNamespace(), + quant_config=None, + ) + + def fake_replace(obj, **changes): + values = vars(obj).copy() + values.update(changes) + return SimpleNamespace(**values) + + class DraftConstructionObserved(Exception): + pass + + def fake_get_model(**_kwargs): + assert draft_hf_config.is_neox_style is False + raise DraftConstructionObserved + + monkeypatch.setattr(loader_module, "replace", fake_replace) + monkeypatch.setattr(loader_module, "get_model", fake_get_model) + with pytest.raises(DraftConstructionObserved): + loader_module.load_dflash_model(_TargetModel(is_neox_style=False), vllm_config) + + +def test_dspark_loader_preserves_checkpoint_rope_layout(monkeypatch): + """DSpark inference uses the rotary layout encoded by its training model.""" + from vllm.model_executor.models import utils as model_utils + from vllm.v1.worker.gpu.spec_decode.dspark import utils as loader_module + + draft_hf_config = SimpleNamespace( + num_hidden_layers=1, + layer_types=["sliding_attention"], + dflash_config={"causal": True}, + is_neox_style=True, + ) + speculative_config = SimpleNamespace( + draft_model_config=SimpleNamespace(hf_config=draft_hf_config), + attention_backend=None, + kv_cache_dtype=None, + draft_load_config=None, + ) + vllm_config = SimpleNamespace( + speculative_config=speculative_config, + attention_config=SimpleNamespace(), + cache_config=SimpleNamespace(), + quant_config=None, + ) + + class DraftConstructionObserved(Exception): + pass + + def fake_get_model(**_kwargs): + assert draft_hf_config.is_neox_style is True + raise DraftConstructionObserved + + monkeypatch.setattr(loader_module, "get_model", fake_get_model) + monkeypatch.setattr( + loader_module, + "_create_draft_vllm_config", + lambda _config: vllm_config, + ) + monkeypatch.setattr(model_utils, "get_draft_quant_config", lambda _config: None) + + with pytest.raises(DraftConstructionObserved): + loader_module.load_dspark_model(_TargetModel(is_neox_style=False), vllm_config) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 58c59fcbb5cb..40c92b781faa 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -86,6 +86,27 @@ def dflash_has_any_non_causal(config: Qwen3Config) -> bool: ) +def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: + """Return the target model's rotary-embedding layout when it is exposed. + + A DFlash-family draft must rotate query and key tensors with the same + dimension layout as the target model used during hidden-state extraction. + Draft checkpoints do not encode this property. A mismatch changes every + drafted attention result without raising an error and collapses token + acceptance. + """ + language_model = ( + target_model.get_language_model() + if hasattr(target_model, "get_language_model") + else target_model + ) + for module in language_model.modules(): + style = getattr(module, "is_neox_style", None) + if isinstance(style, bool): + return style + return None + + def _get_dflash_fc_input_size(vllm_config: VllmConfig) -> int: spec_config = vllm_config.speculative_config config = spec_config.draft_model_config.hf_config @@ -216,6 +237,7 @@ def __init__( add_swa_attention_sink_bias: bool = False, sliding_window: int | None = None, causal: bool = False, + is_neox_style: bool = True, cache_config: CacheConfig | None = None, quant_config: QuantizationConfig | None = None, prefix: str = "", @@ -259,6 +281,7 @@ def __init__( self.rotary_emb = get_rope( self.head_dim, max_position=max_position, + is_neox_style=is_neox_style, rope_parameters=rope_parameters, ) @@ -342,6 +365,12 @@ def __init__( # non-causal) from the draft config. sliding_window, causal = _resolve_layer_attention(config, layer_idx) + # The loader copies this value from the built target model. Kimi-K3 + # uses interleaved rotary dimensions while Qwen3 defaults to NeoX + # rotary dimensions, so relying on the Qwen3 default is not valid for + # a Kimi-trained DFlash-family checkpoint. + is_neox_style = getattr(config, "is_neox_style", True) + self.self_attn = DFlashQwen3Attention( hidden_size=self.hidden_size, num_heads=config.num_attention_heads, @@ -352,6 +381,7 @@ def __init__( add_swa_attention_sink_bias=add_swa_attention_sink_bias, sliding_window=sliding_window, causal=causal, + is_neox_style=is_neox_style, head_dim=getattr(config, "head_dim", None), cache_config=cache_config, quant_config=quant_config, diff --git a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py index 256f23c42ace..a5e643ed10c7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dflash/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dflash/utils.py @@ -81,11 +81,17 @@ def maybe_load_mask_embedding( def load_dflash_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Module: from vllm.compilation.backends import set_model_tag - from vllm.model_executor.models.qwen3_dflash import dflash_has_any_non_causal + from vllm.model_executor.models.qwen3_dflash import ( + dflash_has_any_non_causal, + dflash_target_rope_is_neox_style, + ) speculative_config = vllm_config.speculative_config assert speculative_config is not None draft_model_config = speculative_config.draft_model_config + is_neox_style = dflash_target_rope_is_neox_style(target_model) + if is_neox_style is not None: + draft_model_config.hf_config.is_neox_style = is_neox_style # Select an attention backend that supports the drafter's attention: mixing # a non-causal layer onto a causal-only backend would fail. draft_vllm_config = replace( diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index cd23ae1bb559..f366a0386b90 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -59,7 +59,9 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo with rope_ownership, set_model_tag("dspark_head"): draft_model = get_model( - vllm_config=draft_vllm_config, model_config=draft_model_config + vllm_config=draft_vllm_config, + model_config=draft_model_config, + load_config=speculative_config.draft_load_config, ) if get_pp_group().world_size != 1: From 8f7f6c66364f2d6d863510aa6ac016582c4a3b02 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 14 Aug 2026 13:17:53 +0000 Subject: [PATCH 09/52] attention: execute replicated KV groups without DCP partitioning Co-authored-by: Codex --- tests/v1/spec_decode/test_dflash_swa.py | 48 ++++++++++++++++++++++++ tests/v1/worker/test_cp_utils.py | 31 +++++++++++++++ vllm/v1/attention/backends/flash_attn.py | 19 ++++++---- vllm/v1/worker/cp_utils.py | 9 +++++ 4 files changed, 100 insertions(+), 7 deletions(-) diff --git a/tests/v1/spec_decode/test_dflash_swa.py b/tests/v1/spec_decode/test_dflash_swa.py index 2a0c9d6d6f7e..092e86492cb5 100644 --- a/tests/v1/spec_decode/test_dflash_swa.py +++ b/tests/v1/spec_decode/test_dflash_swa.py @@ -10,6 +10,7 @@ from vllm.model_executor.models.qwen3_dflash import DFlashAttention from vllm.transformers_utils.configs.speculators import SpeculatorsConfig from vllm.v1.attention.backend import AttentionType, CommonAttentionMetadata +from vllm.v1.attention.backends import flash_attn as flash_attn_backend from vllm.v1.kv_cache_interface import ( FullAttentionSpec, SlidingWindowSpec, @@ -161,6 +162,53 @@ def test_dflash_swa_layers_keep_sliding_window_kv_cache_spec(monkeypatch): assert spec.dcp_replicated is True +def test_flash_attention_metadata_treats_replicated_kv_as_dcp1(monkeypatch): + """Replicated draft cache metadata uses global sequence lengths locally.""" + spec = SlidingWindowSpec( + block_size=16, + num_kv_heads=1, + head_size=128, + dtype=torch.bfloat16, + sliding_window=2048, + dcp_replicated=True, + ) + model_config = SimpleNamespace( + get_num_attention_heads=lambda _parallel_config: 96, + get_num_kv_heads=lambda _parallel_config: 16, + get_head_size=lambda: 128, + rswa_window=None, + is_mm_prefix_lm=False, + ) + vllm_config = SimpleNamespace( + model_config=model_config, + parallel_config=SimpleNamespace(cp_kv_cache_interleave_size=1), + cache_config=SimpleNamespace(cache_dtype="bfloat16"), + compilation_config=SimpleNamespace( + cudagraph_mode=SimpleNamespace(has_full_cudagraphs=lambda: False), + max_cudagraph_capture_size=None, + ), + attention_config=SimpleNamespace( + flash_attn_max_num_splits_for_cuda_graph=0, + ), + scheduler_config=SimpleNamespace(max_num_seqs=1), + ) + + monkeypatch.setattr( + flash_attn_backend, + "get_dcp_group", + lambda: SimpleNamespace(world_size=16, rank_in_group=7), + ) + builder = flash_attn_backend.FlashAttentionMetadataBuilder( + spec, + ["draft.layer"], + vllm_config, + torch.device("cpu"), + ) + + assert builder.dcp_world_size == 1 + assert builder.dcp_rank == 0 + + def test_dflash_swa_layers_use_causal_metadata(): proposer = object.__new__(DFlashProposer) proposer.model = SimpleNamespace(sliding_attention_layer_names={"layer.sw"}) diff --git a/tests/v1/worker/test_cp_utils.py b/tests/v1/worker/test_cp_utils.py index 186d49a038c2..ddcf0c171dca 100644 --- a/tests/v1/worker/test_cp_utils.py +++ b/tests/v1/worker/test_cp_utils.py @@ -94,3 +94,34 @@ def test_check_attention_cp_compatibility_rejects_no_lse_return(monkeypatch): with pytest.raises(AssertionError, match="requires attention implementations"): cp_utils.check_attention_cp_compatibility(_make_config(dcp_size=2)) + + +def test_replicated_kv_group_executes_attention_as_dcp1(monkeypatch): + """A complete per-rank KV copy must not enter DCP attention collectives.""" + impl = SimpleNamespace( + can_return_lse_for_decode=True, + dcp_world_size=16, + dcp_rank=7, + total_cp_world_size=16, + total_cp_rank=7, + need_to_return_lse_for_decode=True, + supports_pcp=False, + ) + layer = SimpleNamespace( + impl=impl, + get_kv_cache_spec=lambda _config: SimpleNamespace(dcp_replicated=True), + ) + + monkeypatch.setattr( + cp_utils, + "get_layers_from_vllm_config", + lambda vllm_config, layer_type: {"draft.layer": layer}, + ) + + cp_utils.check_attention_cp_compatibility(_make_config(dcp_size=16)) + + assert impl.dcp_world_size == 1 + assert impl.dcp_rank == 0 + assert impl.total_cp_world_size == 1 + assert impl.total_cp_rank == 0 + assert impl.need_to_return_lse_for_decode is False diff --git a/vllm/v1/attention/backends/flash_attn.py b/vllm/v1/attention/backends/flash_attn.py index 97f30736fda3..a8444f7bbb7b 100755 --- a/vllm/v1/attention/backends/flash_attn.py +++ b/vllm/v1/attention/backends/flash_attn.py @@ -395,15 +395,20 @@ def __init__( self.max_num_splits = 0 # No upper bound on the number of splits. self.aot_schedule = get_flash_attn_version() == 3 - try: - from vllm.distributed.parallel_state import get_dcp_group - - self.dcp_world_size = get_dcp_group().world_size - self.dcp_rank = get_dcp_group().rank_in_group - except AssertionError: - # DCP might not be initialized in testing + # A DCP-replicated KV group stores the complete sequence on every + # rank. Build ordinary local-attention metadata for that group instead + # of partitioning its sequence lengths for a second time. + if getattr(kv_cache_spec, "dcp_replicated", False): self.dcp_world_size = 1 self.dcp_rank = 0 + else: + try: + self.dcp_world_size = get_dcp_group().world_size + self.dcp_rank = get_dcp_group().rank_in_group + except AssertionError: + # DCP might not be initialized in testing + self.dcp_world_size = 1 + self.dcp_rank = 0 self.cp_kv_cache_interleave_size = ( self.parallel_config.cp_kv_cache_interleave_size diff --git a/vllm/v1/worker/cp_utils.py b/vllm/v1/worker/cp_utils.py index 243d8f19052c..f6725cba470c 100644 --- a/vllm/v1/worker/cp_utils.py +++ b/vllm/v1/worker/cp_utils.py @@ -47,6 +47,15 @@ def check_attention_cp_compatibility(vllm_config: VllmConfig) -> None: except Exception: spec = None if getattr(spec, "dcp_replicated", False): + # A replicated KV group contains the complete sequence on + # every rank. Its attention kernel must therefore execute + # as a local DCP1 operation; applying DCP collectives would + # partition and reduce the same cache a second time. + layer_impl.dcp_world_size = 1 + layer_impl.dcp_rank = 0 + layer_impl.total_cp_world_size = 1 + layer_impl.total_cp_rank = 0 + layer_impl.need_to_return_lse_for_decode = False continue if vllm_config.speculative_config is not None and interleave_size > 1: assert layer_impl.supports_mtp_with_cp_non_trivial_interleave_size, ( From f85d3d87e196546eff162c2a61b49aab13f2e625 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 14 Aug 2026 14:38:11 +0000 Subject: [PATCH 10/52] models: document DFlash rotary layout detection Document the target model input and optional NeoX layout result using the repository's Google-style docstring contract. This is documentation-only and does not change runtime behavior. Co-authored-by: OpenAI Codex --- vllm/model_executor/models/qwen3_dflash.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/vllm/model_executor/models/qwen3_dflash.py b/vllm/model_executor/models/qwen3_dflash.py index 40c92b781faa..929beae39207 100644 --- a/vllm/model_executor/models/qwen3_dflash.py +++ b/vllm/model_executor/models/qwen3_dflash.py @@ -94,6 +94,12 @@ def dflash_target_rope_is_neox_style(target_model: nn.Module) -> bool | None: Draft checkpoints do not encode this property. A mismatch changes every drafted attention result without raising an error and collapses token acceptance. + + Args: + target_model: Target model that can expose rotary layout modules. + + Returns: + The exposed NeoX rotary-layout setting, or ``None`` when unavailable. """ language_model = ( target_model.get_language_model() From 04a6acfe467f4a208c7231a18fc99faf656d016a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 18 Aug 2026 02:11:29 +0000 Subject: [PATCH 11/52] [II] Keep Kimi K3 protocol markers out of streamed content Initialize fresh assistant generations in the reasoning channel when Kimi thinking is enabled, while preserving rendered marker state for continued assistant messages. Filter complete and split XTML control markers at the composed parser boundary so malformed model transitions cannot expose protocol syntax as API content. The thinking-disabled path and continuation semantics remain unchanged. Validation: 72 Kimi K3 reasoning and tool-parser tests; Ruff format and lint; git diff whitespace validation. --- .../test_kimi_k3_reasoning_parser.py | 98 +++++++++++++++++++ tests/tool_use/test_kimi_k3_tool_parser.py | 42 ++++++++ vllm/parser/kimi_k3.py | 65 ++++++++++-- vllm/reasoning/kimi_k3_reasoning_parser.py | 26 +++++ 4 files changed, 222 insertions(+), 9 deletions(-) diff --git a/tests/reasoning/test_kimi_k3_reasoning_parser.py b/tests/reasoning/test_kimi_k3_reasoning_parser.py index 4f831e851e9b..83ee0b5d0951 100644 --- a/tests/reasoning/test_kimi_k3_reasoning_parser.py +++ b/tests/reasoning/test_kimi_k3_reasoning_parser.py @@ -52,6 +52,7 @@ def test_parser_selection_thinking_disabled(): ) assert parser._thinking_enabled is False + assert parser.thinking_enabled is False def test_extract_reasoning_with_xtml_tags(): @@ -119,6 +120,103 @@ def test_is_reasoning_end_ignores_stale_close_from_prior_turn(): assert not parser.is_reasoning_end([*new_open]) +def test_fresh_assistant_prompt_does_not_inherit_closed_reasoning(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": True, + "continue_final_message": False, + }, + ) + + assert not parser.is_reasoning_end_for_prompt(CLOSE_IDS) + + +def test_continued_assistant_prompt_uses_rendered_reasoning_state(): + parser = KimiK3ReasoningParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": False, + "continue_final_message": True, + }, + ) + + assert parser.is_reasoning_end_for_prompt(CLOSE_IDS) + assert not parser.is_reasoning_end_for_prompt(OPEN_IDS) + + +def test_fresh_assistant_stream_classifies_first_tokens_as_reasoning(): + parser = ReasoningOnlyParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": True, + }, + ) + request = ChatCompletionRequest(model="test-model", messages=[]) + + first = parser.parse_delta( + delta_text=".", + delta_token_ids=[9], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=False, + ) + partial_close = parser.parse_delta( + delta_text=f"{CLOSE}think", + delta_token_ids=CLOSE_IDS[:2], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=False, + ) + closed = parser.parse_delta( + delta_text=f"{SEP}{RESPONSE_OPEN}", + delta_token_ids=[CLOSE_IDS[2], 10], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=False, + ) + + assert first is not None + assert first.reasoning == "." + assert first.content is None + assert partial_close is None + assert closed is None + + +def test_content_filter_holds_and_removes_split_protocol_markers(): + parser = ReasoningOnlyParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": False, + "continue_final_message": True, + }, + ) + request = ChatCompletionRequest(model="test-model", messages=[]) + chunks = [".", f"{CLOSE}think", SEP, RESPONSE_OPEN] + messages: list[DeltaMessage] = [] + + for index, chunk in enumerate(chunks): + delta = parser.parse_delta( + delta_text=chunk, + delta_token_ids=[9 + index], + request=request, + prompt_token_ids=CLOSE_IDS, + finished=index == len(chunks) - 1, + ) + if delta is not None: + messages.append(delta) + + content = "".join(message.content or "" for message in messages) + assert content == "." + assert OPEN not in content + assert CLOSE not in content + assert SEP not in content + + def test_streaming_split_open_marker_is_held_back(): parser = KimiK3ReasoningParser(DummyTokenizer()) diff --git a/tests/tool_use/test_kimi_k3_tool_parser.py b/tests/tool_use/test_kimi_k3_tool_parser.py index 118c51db574e..be6546a5ed6d 100644 --- a/tests/tool_use/test_kimi_k3_tool_parser.py +++ b/tests/tool_use/test_kimi_k3_tool_parser.py @@ -171,6 +171,48 @@ def test_delegating_parser_preserves_tool_calls_after_reasoning(): assert json.loads(tool_calls[0].arguments) == {"x": 1} +def test_fresh_tool_stream_does_not_inherit_closed_reasoning(): + parser = KimiK3DelegatingParser( + DummyTokenizer(), + chat_template_kwargs={ + "thinking": True, + "add_generation_prompt": True, + }, + ) + request = _request() + messages: list[DeltaMessage] = [] + chunks = [ + (".", [9]), + (f"{CLOSE}think", [4, 2]), + ( + f"{SEP}{_response('')}{_tools(_call('calc', 1))}", + [3, 10], + ), + ] + + for index, (text, token_ids) in enumerate(chunks): + delta = parser.parse_delta( + delta_text=text, + delta_token_ids=token_ids, + request=request, + prompt_token_ids=[4, 2, 3], + finished=index == len(chunks) - 1, + ) + if delta is not None: + messages.append(delta) + + reasoning = "".join(message.reasoning or "" for message in messages) + content = "".join(message.content or "" for message in messages) + tool_calls = [call for message in messages for call in (message.tool_calls or [])] + assert reasoning == "." + assert content == "" + assert len(tool_calls) == 1 + assert tool_calls[0].function.name == "calc" + assert OPEN not in reasoning + content + assert CLOSE not in reasoning + content + assert SEP not in reasoning + content + + def test_delegating_parser_required_tool_choice_uses_xtml_parser(): parser = KimiK3DelegatingParser(DummyTokenizer()) request = _request().model_copy(update={"tool_choice": "required"}) diff --git a/vllm/parser/kimi_k3.py b/vllm/parser/kimi_k3.py index 380e2cce5f91..e3f58d25f4d7 100644 --- a/vllm/parser/kimi_k3.py +++ b/vllm/parser/kimi_k3.py @@ -12,6 +12,7 @@ ) from vllm.parser.abstract_parser import DelegatingParser from vllm.reasoning.kimi_k3_reasoning_parser import KimiK3ReasoningParser +from vllm.tool_parsers.utils import partial_tag_overlap if TYPE_CHECKING: from vllm.entrypoints.openai.chat_completion.protocol import ( @@ -23,6 +24,45 @@ class KimiK3Parser(DelegatingParser): """Compose the Kimi K3 reasoning and tool parsers for XTML output.""" + _CONTENT_PROTOCOL_MARKERS = ( + "<|open|>think<|sep|>", + "<|close|>think<|sep|>", + "<|open|>response<|sep|>", + "<|close|>response<|sep|>", + "<|close|>message<|sep|>", + ) + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._pending_content_protocol = "" + + def _strip_content_protocol( + self, content: str | None, *, finished: bool + ) -> str | None: + """Remove Kimi XTML control markers from streamed API content. + + Reasoning and tool parsers normally consume these markers before this + method runs. The final filter preserves the API invariant when a + malformed model transition or an inconsistent initial parser phase + routes a marker through the content field. Marker prefixes are held + across stream chunks so partial control tokens are never exposed. + """ + pending = self._pending_content_protocol + (content or "") + for marker in self._CONTENT_PROTOCOL_MARKERS: + pending = pending.replace(marker, "") + + overlap = max( + partial_tag_overlap(pending, marker) + for marker in self._CONTENT_PROTOCOL_MARKERS + ) + if overlap: + emitted = pending[:-overlap] + self._pending_content_protocol = "" if finished else pending[-overlap:] + else: + emitted = pending + self._pending_content_protocol = "" + return emitted or None + # TODO: Switch Kimi K3 to the parser engine once its XTML reasoning/tool # path is covered there. def _extract_tool_calls( @@ -113,18 +153,25 @@ def parse_delta( ) if ( - self._tool_parser is not None - or not isinstance(self._reasoning_parser, KimiK3ReasoningParser) - or not state.reasoning_ended - or delta_message is None + delta_message is not None + and self._tool_parser is None + and isinstance(self._reasoning_parser, KimiK3ReasoningParser) + and state.reasoning_ended ): - return delta_message + stripped = self._reasoning_parser.strip_content_streaming( + previous_text=previous_content, + current_text=state.previous_text, + ) + delta_message.content = stripped.content if stripped is not None else None + + if delta_message is None: + if finished: + self._pending_content_protocol = "" + return None - stripped = self._reasoning_parser.strip_content_streaming( - previous_text=previous_content, - current_text=state.previous_text, + delta_message.content = self._strip_content_protocol( + delta_message.content, finished=finished ) - delta_message.content = stripped.content if stripped is not None else None if ( delta_message.role is None and delta_message.content is None diff --git a/vllm/reasoning/kimi_k3_reasoning_parser.py b/vllm/reasoning/kimi_k3_reasoning_parser.py index f05d5c044805..d170652a20e8 100644 --- a/vllm/reasoning/kimi_k3_reasoning_parser.py +++ b/vllm/reasoning/kimi_k3_reasoning_parser.py @@ -95,6 +95,9 @@ def __init__(self, tokenizer: PreTrainedTokenizerBase, *args, **kwargs): if thinking is None: thinking = chat_kwargs.get("enable_thinking", True) self._thinking_enabled = bool(thinking) + self._starts_new_assistant_message = bool( + chat_kwargs.get("add_generation_prompt", True) + ) and not bool(chat_kwargs.get("continue_final_message", False)) # XTML markers as literal strings (skip_special_tokens=False at serve time) self._think_open = "<|open|>think<|sep|>" @@ -145,6 +148,11 @@ def reasoning_start_str(self) -> str | None: def reasoning_end_str(self) -> str | None: return self._think_close + @property + def thinking_enabled(self) -> bool: + """Whether the request opens a Kimi K3 reasoning channel.""" + return self._thinking_enabled + def adjust_request( self, request: "ChatCompletionRequest | ResponsesRequest", @@ -170,6 +178,24 @@ def is_reasoning_end(self, input_ids: Sequence[int]) -> bool: _newest_marker(input_ids, self._think_close_ids, self._think_open_ids) == 0 ) + def is_reasoning_end_for_prompt(self, input_ids: Sequence[int]) -> bool: + """Return the reasoning phase at the generation boundary. + + A fresh Kimi K3 assistant message starts in the think channel whenever + thinking is enabled. That request contract is authoritative even when + the token list exposed to the output parser does not contain the + generation-prefix marker; closed think channels from historical turns + cannot determine the phase of a fresh assistant message. + + ``continue_final_message`` does not open a fresh channel, so its prompt + markers remain authoritative. + """ + if not self._thinking_enabled: + return True + if self._starts_new_assistant_message: + return False + return self.is_reasoning_end(input_ids) + def is_reasoning_end_streaming( self, input_ids: Sequence[int], delta_ids: Iterable[int] ) -> bool: From c09bed025ff6ce82b358833cac3c87fe51618862 Mon Sep 17 00:00:00 2001 From: jungjiyu Date: Sun, 2 Aug 2026 10:49:02 +0000 Subject: [PATCH 12/52] Fix invalid block handling for hybrid KV cache groups Signed-off-by: jungjiyu Assisted-by: ChatGPT --- .../unit/test_invalid_blocks_correctness.py | 202 ++++++++++++++++++ tests/v1/kv_connector/unit/utils.py | 2 + vllm/v1/core/sched/scheduler.py | 95 ++++---- 3 files changed, 263 insertions(+), 36 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 77d629729776..d2a14b363984 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -15,8 +15,15 @@ from unittest.mock import Mock import pytest +import torch from vllm.v1.core.sched.scheduler import Scheduler +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + SlidingWindowSpec, +) from vllm.v1.request import FinishReason, Request, RequestStatus from .utils import ( @@ -24,6 +31,7 @@ create_request, create_scheduler, create_vllm_config, + make_kv_cache_config, ) pytestmark = pytest.mark.cpu_test @@ -478,3 +486,197 @@ def cache_blocks_spy(req, num_tokens): # request should be in the running queue assert request in recompute_scheduler.running + + +def test_sync_recompute_handles_invalid_block_in_second_kv_cache_group(): + """Invalid blocks in any hybrid KV group must trigger recomputation.""" + block_size = 16 + num_blocks = 128 + num_prompt_blocks = 8 + num_external_computed_blocks = 7 + invalid_block_idx = 3 + + vllm_config = create_vllm_config( + block_size=block_size, + kv_load_failure_policy="recompute", + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=make_kv_cache_config( + block_size=block_size, + swa_enabled=True, + sw_size=num_prompt_blocks * block_size, + num_blocks=num_blocks, + ), + ) + + request = create_request( + num_tokens=num_prompt_blocks * block_size, + block_size=block_size, + ) + scheduler.add_request(request) + + num_external_computed_tokens = num_external_computed_blocks * block_size + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + { + request.request_id: num_external_computed_tokens, + }, + False, + ) + ) + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.request_finished_all_groups.return_value = ( + False, + None, + ) + scheduler.connector.take_events.return_value = () + + scheduler_output = scheduler.schedule() + + assert request.status == RequestStatus.RUNNING + assert len(scheduler_output.scheduled_new_reqs) == 1 + + block_ids_by_group = scheduler_output.scheduled_new_reqs[0].block_ids + assert len(block_ids_by_group) == 2 + assert all( + len(group_block_ids) > invalid_block_idx + for group_block_ids in block_ids_by_group + ) + + # Report a load failure in the second KV cache group, not the first. + invalid_block_id = block_ids_by_group[1][invalid_block_idx] + model_runner_output = create_model_runner_output( + [request], + invalid_block_ids={invalid_block_id}, + use_eos=False, + ) + + scheduler.update_from_output( + scheduler_output, + model_runner_output, + ) + + assert request.status == RequestStatus.RUNNING + assert request.num_computed_tokens == invalid_block_idx * block_size + assert request in scheduler.running + assert request.request_id in scheduler.requests + + +def test_sync_recompute_handles_mixed_kv_group_block_sizes(): + """Recompute from a common boundary for mixed KV block sizes.""" + hash_block_size = 16 + scheduler_block_size = 64 + num_blocks = 512 + num_prompt_tokens = 8 * scheduler_block_size + num_external_computed_tokens = 7 * scheduler_block_size + + # The invalid block begins at token 3 * 32 = 96. + # Scheduling granularity is 64, so recomputation must restart at 64. + invalid_group_idx = 1 + invalid_block_idx = 3 + + vllm_config = create_vllm_config( + block_size=scheduler_block_size, + kv_load_failure_policy="recompute", + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full_16"], + FullAttentionSpec( + block_size=16, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, + ), + ), + KVCacheGroupSpec( + ["window_32"], + SlidingWindowSpec( + block_size=32, + num_kv_heads=4, + head_size=16, + dtype=torch.float16, + sliding_window=num_prompt_tokens, + ), + ), + KVCacheGroupSpec( + ["full_64"], + FullAttentionSpec( + block_size=64, + num_kv_heads=1, + head_size=32, + dtype=torch.float16, + ), + ), + ], + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + hash_block_size=hash_block_size, + ) + + request = create_request( + num_tokens=num_prompt_tokens, + block_size=hash_block_size, + ) + scheduler.add_request(request) + + scheduler.connector = Mock() + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + { + request.request_id: num_external_computed_tokens, + }, + False, + ) + ) + scheduler.connector.request_finished.return_value = ( + False, + None, + ) + scheduler.connector.request_finished_all_groups.return_value = ( + False, + None, + ) + scheduler.connector.take_events.return_value = () + + scheduler_output = scheduler.schedule() + + assert request.status == RequestStatus.RUNNING + assert len(scheduler_output.scheduled_new_reqs) == 1 + + block_ids_by_group = scheduler_output.scheduled_new_reqs[0].block_ids + assert len(block_ids_by_group) == 3 + assert len(block_ids_by_group[invalid_group_idx]) > invalid_block_idx + + invalid_block_id = block_ids_by_group[invalid_group_idx][invalid_block_idx] + model_runner_output = create_model_runner_output( + [request], + invalid_block_ids={invalid_block_id}, + use_eos=False, + ) + + scheduler.update_from_output( + scheduler_output, + model_runner_output, + ) + + invalid_block_start = invalid_block_idx * 32 + expected_recompute_from = ( + invalid_block_start // scheduler_block_size * scheduler_block_size + ) + + assert invalid_block_start == 96 + assert expected_recompute_from == 64 + assert request.num_computed_tokens == expected_recompute_from + assert request.status == RequestStatus.RUNNING + assert request in scheduler.running + assert request.request_id in scheduler.requests diff --git a/tests/v1/kv_connector/unit/utils.py b/tests/v1/kv_connector/unit/utils.py index b1fb43353374..a2908b7ab2e4 100644 --- a/tests/v1/kv_connector/unit/utils.py +++ b/tests/v1/kv_connector/unit/utils.py @@ -153,6 +153,7 @@ def create_scheduler( vllm_config: VllmConfig, num_blocks: int = 10000, kv_cache_config: KVCacheConfig | None = None, + hash_block_size: int | None = None, ) -> Scheduler | AsyncScheduler: """Initialize Scheduler For Testing.""" block_size = vllm_config.cache_config.block_size @@ -183,6 +184,7 @@ def create_scheduler( log_stats=True, structured_output_manager=StructuredOutputManager(vllm_config), block_size=block_size, + hash_block_size=hash_block_size, ) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ae08df9c35b2..dcb2cd1a5fe1 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -2914,50 +2914,73 @@ def _update_requests_with_invalid_blocks( is_affected = False marked_invalid_block = False req_id = request.request_id - # TODO (davidb): add support for hybrid memory allocator - (req_block_ids,) = self.kv_cache_manager.get_block_ids(req_id) - # We iterate only over blocks that may contain externally computed - # tokens + req_block_ids_by_group = self.kv_cache_manager.get_block_ids(req_id) + # We iterate only over blocks that may contain externally + # computed tokens. req_num_computed_tokens = ( request.num_computed_tokens - num_scheduled_tokens.get(req_id, 0) ) + computed_block_ids_by_group = ( + self.kv_cache_manager.get_block_ids_for_computed_tokens( + req_id, + req_num_computed_tokens, + ) + ) - req_num_computed_blocks = ( - req_num_computed_tokens + self.block_size - 1 - ) // self.block_size - for idx, block_id in zip(range(req_num_computed_blocks), req_block_ids): - if block_id not in invalid_block_ids: - continue - - is_affected = True - - if block_id in marked_invalid_block_ids: - # This invalid block is shared with a previous request - # and was already marked for recomputation. - # This means this request can still consider this block - # as computed when rescheduled. - # Currently this only applies to sync loading; Async - # loading does not yet support block sharing - continue - - marked_invalid_block_ids.add(block_id) + # Map each invalid block to the earliest scheduler-aligned token + # boundary from which this request must be recomputed. + invalid_block_boundaries: dict[int, int] = {} + for group, group_block_ids in zip( + self.kv_cache_config.kv_cache_groups, + computed_block_ids_by_group, + strict=True, + ): + group_block_size = group.kv_cache_spec.block_size + for block_idx, block_id in enumerate(group_block_ids): + if block_id not in invalid_block_ids: + continue - if marked_invalid_block: - # This request has already marked an invalid block for - # recomputation and updated its num_computed_tokens. - continue + block_start = block_idx * group_block_size + recompute_from = block_start // self.block_size * self.block_size + previous_boundary = invalid_block_boundaries.get(block_id) + invalid_block_boundaries[block_id] = ( + recompute_from + if previous_boundary is None + else min(previous_boundary, recompute_from) + ) - marked_invalid_block = True - # Truncate the computed tokens at the first failed block - request.num_computed_tokens = idx * self.block_size - num_affected_tokens = ( - req_num_computed_tokens - request.num_computed_tokens + if invalid_block_boundaries: + is_affected = True + new_invalid_block_ids = ( + invalid_block_boundaries.keys() - marked_invalid_block_ids ) - total_affected_tokens += num_affected_tokens + marked_invalid_block_ids.update(new_invalid_block_ids) - # collect invalid block and all downstream dependent blocks - if evict_blocks: - blocks_to_evict.update(req_block_ids[idx:]) + if new_invalid_block_ids: + marked_invalid_block = True + request.num_computed_tokens = min( + invalid_block_boundaries[block_id] + for block_id in new_invalid_block_ids + ) + num_affected_tokens = ( + req_num_computed_tokens - request.num_computed_tokens + ) + total_affected_tokens += num_affected_tokens + + # Every KV group after the common recomputation boundary + # depends on the failed prefix, so collect downstream + # blocks from all groups. + if evict_blocks: + for group, group_block_ids in zip( + self.kv_cache_config.kv_cache_groups, + req_block_ids_by_group, + strict=True, + ): + group_block_size = group.kv_cache_spec.block_size + first_block_idx = ( + request.num_computed_tokens // group_block_size + ) + blocks_to_evict.update(group_block_ids[first_block_idx:]) if is_affected: if not marked_invalid_block: From 6bafa633153a44ba1aa54eb7c5bafac248fe68e6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 18 Aug 2026 08:45:05 +0000 Subject: [PATCH 13/52] test(scheduler): cover hybrid KV fail policy across 17 groups Model a 17-group hybrid KV layout and report a load failure from the final group. The test requires failure_policy=fail to finish only the affected request, emit an error result, and schedule a subsequent healthy request.\n\nValidation: 20 KV load-failure tests and 7 hybrid/Mamba scheduler tests pass in the CUDA 13.3 PyTorch 2.13 runtime. --- .../unit/test_invalid_blocks_correctness.py | 105 ++++++++++++++++++ 1 file changed, 105 insertions(+) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index d2a14b363984..9c41bd407828 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -17,6 +17,9 @@ import pytest import torch +from vllm.distributed.kv_transfer.kv_connector.v1.nixl.connector import ( + NixlBaseConnector, +) from vllm.v1.core.sched.scheduler import Scheduler from vllm.v1.kv_cache_interface import ( FullAttentionSpec, @@ -565,6 +568,108 @@ def test_sync_recompute_handles_invalid_block_in_second_kv_cache_group(): assert request.request_id in scheduler.requests +def test_sync_fail_handles_invalid_block_in_seventeenth_kv_cache_group(): + """A hybrid-cache load failure must fail one request, not the scheduler.""" + block_size = 16 + num_blocks = 512 + num_prompt_blocks = 8 + num_external_computed_blocks = 7 + invalid_block_idx = 3 + num_kv_cache_groups = 17 + + vllm_config = create_vllm_config( + block_size=block_size, + kv_load_failure_policy="fail", + ) + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + *[ + KVCacheGroupSpec( + [f"full_attention_layer_{group_idx}"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + for group_idx in range(num_kv_cache_groups - 1) + ], + KVCacheGroupSpec( + ["sliding_window_layer"], + SlidingWindowSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + sliding_window=num_prompt_blocks * block_size, + ), + ), + ], + ) + scheduler = create_scheduler( + vllm_config, + num_blocks=num_blocks, + kv_cache_config=kv_cache_config, + ) + + failed_request = create_request( + num_tokens=num_prompt_blocks * block_size, + block_size=block_size, + ) + scheduler.add_request(failed_request) + + num_external_computed_tokens = num_external_computed_blocks * block_size + scheduler.connector = Mock(spec=NixlBaseConnector) + scheduler.connector.get_num_new_matched_tokens.side_effect = ( + _make_get_num_new_matched_tokens( + { + failed_request.request_id: num_external_computed_tokens, + }, + False, + ) + ) + scheduler.connector.request_finished.return_value = (False, None) + scheduler.connector.request_finished_all_groups.return_value = ( + False, + None, + ) + scheduler.connector.take_events.return_value = () + + scheduler_output = scheduler.schedule() + block_ids_by_group = scheduler_output.scheduled_new_reqs[0].block_ids + + assert len(block_ids_by_group) == num_kv_cache_groups + invalid_block_id = block_ids_by_group[-1][invalid_block_idx] + model_runner_output = create_model_runner_output( + [failed_request], + invalid_block_ids={invalid_block_id}, + use_eos=False, + ) + + outputs = scheduler.update_from_output(scheduler_output, model_runner_output) + + assert failed_request.status == RequestStatus.FINISHED_ERROR + assert failed_request.request_id not in scheduler.requests + assert failed_request not in scheduler.running + assert len(outputs) == 1 + engine_output = next(iter(outputs.values())).outputs[0] + assert engine_output.request_id == failed_request.request_id + assert engine_output.finish_reason == FinishReason.ERROR + + healthy_request = create_request( + num_tokens=2 * block_size, + block_size=block_size, + ) + scheduler.add_request(healthy_request) + next_scheduler_output = scheduler.schedule() + + assert healthy_request.request_id in next_scheduler_output.num_scheduled_tokens + assert healthy_request.status == RequestStatus.RUNNING + + def test_sync_recompute_handles_mixed_kv_group_block_sizes(): """Recompute from a common boundary for mixed KV block sizes.""" hash_block_size = 16 From 3c296be28b31dba5c4a0ced714d2e21dabe246af Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 18 Aug 2026 10:26:56 +0000 Subject: [PATCH 14/52] test(scheduler): preserve hybrid KV recovery edge cases --- .../unit/test_invalid_blocks_correctness.py | 120 +++++++++++++++++- 1 file changed, 116 insertions(+), 4 deletions(-) diff --git a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py index 9c41bd407828..14cfec017e67 100644 --- a/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py +++ b/tests/v1/kv_connector/unit/test_invalid_blocks_correctness.py @@ -67,6 +67,121 @@ def recompute_scheduler(): return create_scheduler(vllm_config) +def _create_invalid_block_test_scheduler( + scheduler_block_size: int, + group_block_sizes: tuple[int, ...], +) -> Scheduler: + """Create the scheduler state required by invalid-block mapping tests.""" + scheduler = Scheduler.__new__(Scheduler) + scheduler.block_size = scheduler_block_size + scheduler.kv_cache_manager = Mock() + scheduler.kv_cache_config = KVCacheConfig( + num_blocks=128, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + [f"full_attention_{group_idx}"], + FullAttentionSpec( + block_size=block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ) + for group_idx, block_size in enumerate(group_block_sizes) + ], + ) + return scheduler + + +def test_hybrid_cache_invalid_block_truncates_and_evicts_all_groups(): + """A failed hybrid-cache group invalidates the shared logical suffix.""" + scheduler = _create_invalid_block_test_scheduler(16, (16, 16)) + scheduler.kv_cache_manager.get_block_ids.return_value = ( + [1, 2, 3, 4], + [11, 12, 13, 14], + ) + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.return_value = ( + [1, 2, 3, 4], + [11, 12, 13, 14], + ) + + request = Mock(spec=Request) + request.request_id = "hybrid-request" + request.num_computed_tokens = 64 + + affected, affected_tokens, evicted = scheduler._update_requests_with_invalid_blocks( + [request], + invalid_block_ids={12}, + num_scheduled_tokens={}, + ) + + assert affected == {request.request_id} + assert request.num_computed_tokens == 16 + assert affected_tokens == 48 + assert evicted == {2, 3, 4, 12, 13, 14} + + +def test_hybrid_cache_shared_invalid_block_is_recomputed_once(): + """A shared failed block contributes one recomputation interval.""" + scheduler = _create_invalid_block_test_scheduler(16, (16, 16)) + scheduler.kv_cache_manager.get_block_ids.side_effect = ( + ([1, 2, 3], [11, 12, 13]), + ([21, 22, 23], [31, 12, 33]), + ) + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.side_effect = ( + ([1, 2, 3], [11, 12, 13]), + ([21, 22, 23], [31, 12, 33]), + ) + + first_request = Mock(spec=Request) + first_request.request_id = "first-request" + first_request.num_computed_tokens = 48 + second_request = Mock(spec=Request) + second_request.request_id = "second-request" + second_request.num_computed_tokens = 48 + + affected, affected_tokens, evicted = scheduler._update_requests_with_invalid_blocks( + [first_request, second_request], + invalid_block_ids={12}, + num_scheduled_tokens={}, + ) + + assert affected == {first_request.request_id, second_request.request_id} + assert first_request.num_computed_tokens == 16 + assert second_request.num_computed_tokens == 48 + assert affected_tokens == 32 + assert evicted == {2, 3, 12, 13} + + +def test_hybrid_cache_ignores_invalid_blocks_after_external_prefix(): + """Blocks used only by scheduled local tokens are not load failures.""" + scheduler = _create_invalid_block_test_scheduler(32, (16, 32)) + scheduler.kv_cache_manager.get_block_ids.return_value = ( + [1, 2, 3, 4], + [11, 12], + ) + scheduler.kv_cache_manager.get_block_ids_for_computed_tokens.return_value = ( + [1, 2], + [11], + ) + + request = Mock(spec=Request) + request.request_id = "local-suffix" + request.num_computed_tokens = 64 + + affected, affected_tokens, evicted = scheduler._update_requests_with_invalid_blocks( + [request], + invalid_block_ids={3, 12}, + num_scheduled_tokens={request.request_id: 32}, + ) + + assert affected == set() + assert request.num_computed_tokens == 64 + assert affected_tokens == 0 + assert evicted == set() + + def test_sync_recompute_blocks_not_freed_for_running_requests( recompute_scheduler: Scheduler, ): @@ -774,13 +889,10 @@ def test_sync_recompute_handles_mixed_kv_group_block_sizes(): model_runner_output, ) - invalid_block_start = invalid_block_idx * 32 expected_recompute_from = ( - invalid_block_start // scheduler_block_size * scheduler_block_size + invalid_block_idx * 32 // scheduler_block_size * scheduler_block_size ) - assert invalid_block_start == 96 - assert expected_recompute_from == 64 assert request.num_computed_tokens == expected_recompute_from assert request.status == RequestStatus.RUNNING assert request in scheduler.running From a18a86e1faf246e4d022071d72645e57c90b71cf Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Tue, 18 Aug 2026 22:01:04 +0000 Subject: [PATCH 15/52] fix(structured-output): map compacted DSpark grammar rows --- .../v1/worker/test_gpu_structured_outputs.py | 51 +++++++++++ vllm/v1/worker/gpu/model_runner.py | 2 + vllm/v1/worker/gpu/structured_outputs.py | 84 ++++++++++++++++--- 3 files changed, 124 insertions(+), 13 deletions(-) create mode 100644 tests/v1/worker/test_gpu_structured_outputs.py diff --git a/tests/v1/worker/test_gpu_structured_outputs.py b/tests/v1/worker/test_gpu_structured_outputs.py new file mode 100644 index 000000000000..08867468ed0e --- /dev/null +++ b/tests/v1/worker/test_gpu_structured_outputs.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import numpy as np + +from vllm.v1.worker.gpu.structured_outputs import _build_grammar_row_mapping + + +def test_grammar_mapping_preserves_bonus_rows_after_zero_draft_budget(): + """Zero draft capacity retains each request's scheduled bonus mask.""" + source_indices, logits_indices = _build_grammar_row_mapping( + req_ids=["low", "high", "prefill"], + grammar_req_ids=["low", "high", "prefill"], + grammar_num_spec_tokens=[2, 2, 0], + cu_num_logits_np=np.array([0, 1, 2, 3], dtype=np.int32), + num_draft_tokens_per_req=np.array([0, 0, 0], dtype=np.int32), + num_bonus_tokens=1, + ) + + assert source_indices == [2, 5, 6] + assert logits_indices == [0, 1, 2] + + +def test_grammar_mapping_selects_active_drafts_from_each_source_group(): + """Compaction preserves per-request draft rows and the final bonus row.""" + source_indices, logits_indices = _build_grammar_row_mapping( + req_ids=["plain", "trimmed", "full"], + grammar_req_ids=["trimmed", "full"], + grammar_num_spec_tokens=[3, 3], + cu_num_logits_np=np.array([0, 1, 3, 7], dtype=np.int32), + num_draft_tokens_per_req=np.array([0, 1, 3], dtype=np.int32), + num_bonus_tokens=1, + ) + + assert source_indices == [0, 3, 4, 5, 6, 7] + assert logits_indices == [1, 2, 3, 4, 5, 6] + + +def test_grammar_mapping_supports_non_speculative_batches(): + """A batch without draft tokens maps one bonus row per grammar request.""" + source_indices, logits_indices = _build_grammar_row_mapping( + req_ids=["plain", "grammar"], + grammar_req_ids=["grammar"], + grammar_num_spec_tokens=[0], + cu_num_logits_np=np.array([0, 1, 2], dtype=np.int32), + num_draft_tokens_per_req=None, + num_bonus_tokens=1, + ) + + assert source_indices == [0] + assert logits_indices == [1] diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 4ad1e067588d..32fd8b1827aa 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -500,6 +500,7 @@ def load_model(self, load_dummy_weights: bool = False, *args, **kwargs) -> None: max_num_logits=self.max_num_reqs * self.decode_query_len, vocab_size=self.vocab_size, device=self.device, + num_bonus_tokens=self.model_state.num_new_sampled_tokens_per_step, ) if self.is_pooling_model and self.is_last_pp_rank: @@ -1606,6 +1607,7 @@ def sample( input_batch, grammar_output.structured_output_request_ids, grammar_output.grammar_bitmask, + grammar_output.num_spec_tokens, ) if input_batch.num_draft_tokens == 0 or self.rejection_sampler is None: diff --git a/vllm/v1/worker/gpu/structured_outputs.py b/vllm/v1/worker/gpu/structured_outputs.py index 34f00086d2a2..3be14f97f31b 100644 --- a/vllm/v1/worker/gpu/structured_outputs.py +++ b/vllm/v1/worker/gpu/structured_outputs.py @@ -9,8 +9,61 @@ from vllm.v1.worker.gpu.input_batch import InputBatch +def _build_grammar_row_mapping( + req_ids: list[str], + grammar_req_ids: list[str], + grammar_num_spec_tokens: list[int], + cu_num_logits_np: np.ndarray, + num_draft_tokens_per_req: np.ndarray | None, + num_bonus_tokens: int, +) -> tuple[list[int], list[int]]: + """Map serialized grammar rows to the active compact logits layout.""" + assert len(grammar_req_ids) == len(grammar_num_spec_tokens) + assert num_bonus_tokens in (0, 1) + + req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} + source_indices: list[int] = [] + logits_indices: list[int] = [] + source_offset = 0 + + for grammar_req_id, num_source_drafts in zip( + grammar_req_ids, + grammar_num_spec_tokens, + strict=True, + ): + req_idx = req_id_to_idx[grammar_req_id] + num_active_drafts = ( + 0 + if num_draft_tokens_per_req is None + else int(num_draft_tokens_per_req[req_idx]) + ) + assert 0 <= num_active_drafts <= num_source_drafts + + logits_start = int(cu_num_logits_np[req_idx]) + num_active_logits = int( + cu_num_logits_np[req_idx + 1] - cu_num_logits_np[req_idx] + ) + assert num_active_logits == num_active_drafts + num_bonus_tokens + + source_indices.extend(range(source_offset, source_offset + num_active_drafts)) + logits_indices.extend(range(logits_start, logits_start + num_active_drafts)) + if num_bonus_tokens: + source_indices.append(source_offset + num_source_drafts) + logits_indices.append(logits_start + num_active_drafts) + + source_offset += num_source_drafts + num_bonus_tokens + + return source_indices, logits_indices + + class StructuredOutputsWorker: - def __init__(self, max_num_logits: int, vocab_size: int, device: torch.device): + def __init__( + self, + max_num_logits: int, + vocab_size: int, + device: torch.device, + num_bonus_tokens: int, + ): self.logits_indices = torch.zeros( max_num_logits, dtype=torch.int32, device=device ) @@ -19,6 +72,7 @@ def __init__(self, max_num_logits: int, vocab_size: int, device: torch.device): ) self.device = device self.copy_stream = torch.cuda.Stream() + self.num_bonus_tokens = num_bonus_tokens def apply_grammar_bitmask( self, @@ -26,27 +80,31 @@ def apply_grammar_bitmask( input_batch: InputBatch, grammar_req_ids: list[str], grammar_bitmask: np.ndarray, + grammar_num_spec_tokens: list[int], ) -> None: if not grammar_req_ids: return - # Asynchronously copy the bitmask to GPU. + source_indices, mapping = _build_grammar_row_mapping( + input_batch.req_ids, + grammar_req_ids, + grammar_num_spec_tokens, + input_batch.cu_num_logits_np, + input_batch.num_draft_tokens_per_req, + self.num_bonus_tokens, + ) + expected_source_rows = sum( + num_drafts + self.num_bonus_tokens for num_drafts in grammar_num_spec_tokens + ) + assert grammar_bitmask.shape[0] == expected_source_rows + grammar_bitmask = grammar_bitmask[source_indices] + + # Asynchronously copy the active bitmask rows to GPU. with torch.cuda.stream(self.copy_stream): bitmask = async_copy_to_gpu( grammar_bitmask, out=self.grammar_bitmask[: grammar_bitmask.shape[0]] ) - # Construct bitmask -> logits mapping - mapping: list[int] = [] - req_ids = input_batch.req_ids - cu_num_logits = input_batch.cu_num_logits_np.tolist() - req_id_to_idx = {req_id: i for i, req_id in enumerate(req_ids)} - for grammar_req_id in grammar_req_ids: - req_idx = req_id_to_idx[grammar_req_id] - logits_start_idx = cu_num_logits[req_idx] - logits_end_idx = cu_num_logits[req_idx + 1] - mapping.extend(range(logits_start_idx, logits_end_idx)) - # Asynchronously copy the mapping to GPU. with torch.cuda.stream(self.copy_stream): logits_indices = torch.tensor( From e7ffc38b8ffcc89529bcdfbee6127c9a9c1f12e7 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Wed, 19 Aug 2026 09:48:42 +0000 Subject: [PATCH 16/52] test(dspark): define graph-tail mode in capture fixture --- tests/v1/spec_decode/test_dspark_cudagraph_contract.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/v1/spec_decode/test_dspark_cudagraph_contract.py b/tests/v1/spec_decode/test_dspark_cudagraph_contract.py index c1c1c444e52f..770407440bb9 100644 --- a/tests/v1/spec_decode/test_dspark_cudagraph_contract.py +++ b/tests/v1/spec_decode/test_dspark_cudagraph_contract.py @@ -15,6 +15,7 @@ def test_dspark_generate_draft_accepts_dflash_capture_contract(): speculator = SimpleNamespace( num_query_per_req=5, capacity_activation_batch_size=1, + _markov_outside_cudagraph=False, _speculative_steps_for_query_len=Mock(return_value=5), _run_model=Mock(return_value=head_hidden), _sample_sequential=Mock(), From b115455fd0f3d8984d2149974a9007cda51593b6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 20 Aug 2026 09:41:22 +0000 Subject: [PATCH 17/52] fix(structured-output): stop XGrammar batches at termination Stop accepting speculative token batches when the grammar matcher reaches its terminal state. Preserve terminal-state tracking across validation and acceptance calls so tokens after a complete structured value cannot be committed. This is the Infernal Invocation backport of vllm-project/vllm#52805 commits d8cde608cf1f3de406c75f081a76a0e6eb55a9cb, 1cf6f25351357354cf8c520c0b2976b029429668, and 1856abd22452c3da67364986ece7245fce52c950. Signed-off-by: Martin Vit --- .../spec_decode/test_mtp_structured_output.py | 48 +++++++++++++++++++ vllm/v1/structured_output/backend_xgrammar.py | 19 ++++++-- 2 files changed, 62 insertions(+), 5 deletions(-) diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 619f3ad6fded..8bd599733a24 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -262,6 +262,54 @@ def test_validate_tokens_then_bitmask_round_trip(backend): assert not grammar.is_terminated() +def test_xgrammar_accept_tokens_stops_at_termination(capfd): + """Tokens after a terminating EOS do not reach the matcher.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + processed_before = grammar.num_processed_tokens + + assert grammar.accept_tokens(request.request_id, [eos, trailing]) + assert grammar.is_terminated() + assert grammar.num_processed_tokens == processed_before + 1 + assert "trying to accept new token" not in capfd.readouterr().err + + processed_after_eos = grammar.num_processed_tokens + assert grammar.accept_tokens(request.request_id, [trailing]) + assert grammar.num_processed_tokens == processed_after_eos + assert "trying to accept new token" not in capfd.readouterr().err + + grammar.reset() + assert not grammar.is_terminated() + assert grammar.num_processed_tokens == 0 + + +def test_xgrammar_validate_tokens_stops_at_termination(capfd): + """Validation rolls back after reaching a terminating EOS.""" + tokenizer, _, request, prompt = _make_manager_and_request("xgrammar") + grammar = request.structured_output_request.grammar + + assert grammar.accept_tokens(request.request_id, prompt) + + eos = tokenizer.eos_token_id + trailing = tokenizer.encode("\n")[0] + + assert grammar.validate_tokens([eos, trailing]) == [eos] + assert "trying to accept new token" not in capfd.readouterr().err + # Check matcher state directly to verify validation rolled it back. + assert not grammar.matcher.is_terminated() + + assert grammar.accept_tokens(request.request_id, [eos]) + assert grammar.is_terminated() + + assert grammar.validate_tokens([trailing]) == [] + assert "trying to accept new token" not in capfd.readouterr().err + + class _MarkerReasoner: """Stub reasoner whose reasoning-end marker is a single fixed token.""" diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 58b726bf72a4..29659c4e3ba3 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -156,11 +156,12 @@ class XgrammarGrammar(StructuredOutputGrammar): def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: """Accepts a list of tokens and advances the FSM. - Returns True if the FSM was advanced successfully. - Returns False if the FSM failed to advance. + Returns True if all grammar-constrained tokens were accepted. + Tokens after termination are ignored. Returns False if the FSM + failed to advance. """ if self._is_terminated: - return False + return True for token in tokens: if not self.matcher.accept_token(token): logger.error( @@ -171,7 +172,9 @@ def accept_tokens(self, request_id: str, tokens: list[int]) -> bool: ) return False self.num_processed_tokens += 1 - self._is_terminated = self.matcher.is_terminated() + self._is_terminated = self.matcher.is_terminated() + if self._is_terminated: + break return True def validate_tokens(self, tokens: list[int]) -> list[int]: @@ -180,10 +183,15 @@ def validate_tokens(self, tokens: list[int]) -> list[int]: Returns the prefix list of tokens that are accepted by the FSM. """ + if self._is_terminated: + return [] + accepted_tokens = [] for token in tokens: if self.matcher.accept_token(token): accepted_tokens.append(token) + if self.matcher.is_terminated(): + break else: break if len(accepted_tokens) > 0: @@ -203,8 +211,9 @@ def is_terminated(self) -> bool: return self._is_terminated def reset(self): - self.num_processed_tokens = 0 self.matcher.reset() + self.num_processed_tokens = 0 + self._is_terminated = False # cf https://github.com/mlc-ai/xgrammar/blob/a32ac892676d2eedc0327416105b9b06edfb94b2/cpp/json_schema_converter.cc From 5fe7989f2342b0be8f271d5e730aa402e493679f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 15 Aug 2026 13:51:02 +0000 Subject: [PATCH 18/52] fix(structured-output): validate speculative blocks before commit Structured-output masks are prepared before speculative verification. An accepted block can cross reasoning activation or grammar termination, so its suffix may have been sampled under a grammar state that no longer applies at commit time. Validate the accepted block without advancing the matcher, commit only its valid prefix, and roll scheduler accounting back for resampling. Preserve the unstructured and single-token fast paths, and report only committed draft tokens in speculative metrics. Co-authored-by: Adam Moisa Assisted-by: OpenAI Codex Signed-off-by: Martin Vit (cherry picked from commit fa0777ff775a870a5d6a6ff187f769e50b38c4ad) Signed-off-by: Martin Vit --- tests/v1/core/test_scheduler.py | 166 ++++++++++++++++++ .../spec_decode/test_mtp_structured_output.py | 156 ++++++++++++++++ vllm/v1/core/sched/scheduler.py | 45 ++++- vllm/v1/structured_output/__init__.py | 131 +++++++++++++- 4 files changed, 490 insertions(+), 8 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index e4188e9167d5..d04a1febbb3c 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -839,6 +839,172 @@ def test_stop_via_update_from_output(): assert list(requests[0].output_token_ids) == [EOS_TOKEN_ID, 10, 11] +@pytest.mark.parametrize("async_scheduling", [False, True]) +@pytest.mark.parametrize( + "filtered_tokens,num_grammar_rejected,expected_accepted", + [ + ([10, 11, 12], 1, 3), + ([10, 11], 2, 2), + ([], 4, 0), + ], +) +def test_speculative_grammar_filter_rolls_back_scheduler_state_and_stats( + async_scheduling: bool, + filtered_tokens: list[int], + num_grammar_rejected: int, + expected_accepted: int, +): + """Rejected suffix tokens remain schedulable and are not accepted drafts.""" + scheduler = create_scheduler( + num_speculative_tokens=3, + speculative_method="ngram_gpu", + async_scheduling=async_scheduling, + ) + request = create_requests(num_requests=1)[0] + request.structured_output_request = Mock() + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + 4 + request.num_output_placeholders = 4 if async_scheduling else 0 + scheduler.requests[request.request_id] = request + scheduler.running.append(request) + + manager = Mock() + manager.filter_speculative_grammar_tokens.return_value = ( + filtered_tokens, + num_grammar_rejected, + ) + manager.should_advance.return_value = False + scheduler.structured_output_manager = manager + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 4}, + total_num_scheduled_tokens=4, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={request.request_id: [10, 11, 12]}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[10, 11, 12, 13]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + outputs = scheduler.update_from_output(scheduler_output, model_output) + + assert request.num_computed_tokens == request.num_tokens + assert request.num_output_placeholders == 0 + assert list(request.output_token_ids) == filtered_tokens + if filtered_tokens: + assert outputs[0].outputs[0].new_token_ids == filtered_tokens + manager.filter_speculative_grammar_tokens.assert_called_once_with( + request, [10, 11, 12, 13] + ) + stats = outputs[0].scheduler_stats.spec_decoding_stats + assert stats is not None + assert stats.num_drafts == 1 + assert stats.num_draft_tokens == 3 + assert stats.num_accepted_tokens == expected_accepted + assert stats.num_accepted_tokens_per_pos == [ + int(position < expected_accepted) for position in range(3) + ] + + +def test_speculative_grammar_filter_is_not_called_for_unstructured_requests(): + """Unstructured speculative decoding does not enter grammar validation.""" + scheduler = create_scheduler(num_speculative_tokens=2) + request = create_requests(num_requests=1)[0] + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + 3 + scheduler.requests[request.request_id] = request + scheduler.running.append(request) + + manager = Mock() + manager.should_advance.return_value = False + scheduler.structured_output_manager = manager + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 3}, + total_num_scheduled_tokens=3, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={request.request_id: [10, 11]}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[10, 11, 12]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + scheduler.update_from_output(scheduler_output, model_output) + + assert list(request.output_token_ids) == [10, 11, 12] + manager.filter_speculative_grammar_tokens.assert_not_called() + + +def test_speculative_grammar_filter_commits_and_advances_only_valid_prefix(): + """Scheduler output and grammar state share the filtered token prefix.""" + scheduler = create_scheduler(num_speculative_tokens=2) + request = create_requests(num_requests=1)[0] + grammar = Mock(spec=StructuredOutputGrammar) + grammar.accept_tokens.side_effect = lambda request_id, tokens: tokens == [10, 11] + request.structured_output_request = Mock(grammar=grammar) + request.status = RequestStatus.RUNNING + request.num_computed_tokens = request.num_tokens + 3 + scheduler.requests[request.request_id] = request + scheduler.running.append(request) + + manager = Mock() + manager.filter_speculative_grammar_tokens.return_value = ([10, 11], 1) + manager.should_advance.return_value = True + manager.trim_reasoning_for_advance.return_value = [10, 11] + scheduler.structured_output_manager = manager + + scheduler_output = SchedulerOutput( + scheduled_new_reqs=[], + scheduled_cached_reqs=CachedRequestData.make_empty(), + num_scheduled_tokens={request.request_id: 3}, + total_num_scheduled_tokens=3, + scheduled_encoder_inputs={}, + scheduled_spec_decode_tokens={request.request_id: [10, 11]}, + num_common_prefix_blocks=[], + finished_req_ids=set(), + free_encoder_mm_hashes=[], + ) + model_output = ModelRunnerOutput( + req_ids=[request.request_id], + req_id_to_index={request.request_id: 0}, + sampled_token_ids=[[10, 11, 12]], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=[], + ) + + outputs = scheduler.update_from_output(scheduler_output, model_output) + + assert request.status == RequestStatus.RUNNING + assert request.num_computed_tokens == request.num_tokens + assert list(request.output_token_ids) == [10, 11] + assert outputs[0].outputs[0].new_token_ids == [10, 11] + grammar.accept_tokens.assert_called_once_with(request.request_id, [10, 11]) + stats = outputs[0].scheduler_stats.spec_decoding_stats + assert stats is not None + assert stats.num_draft_tokens == 2 + assert stats.num_accepted_tokens == 2 + + def test_check_stop_min_tokens(): """Test that requests don't stop when min_tokens requirement isn't met.""" from vllm.v1.core.sched.utils import check_stop diff --git a/tests/v1/spec_decode/test_mtp_structured_output.py b/tests/v1/spec_decode/test_mtp_structured_output.py index 619f3ad6fded..0d4ac0909c3f 100644 --- a/tests/v1/spec_decode/test_mtp_structured_output.py +++ b/tests/v1/spec_decode/test_mtp_structured_output.py @@ -2,12 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """grammar_bitmask under spec-decode draft padding (#44006).""" +from collections.abc import Iterable, Sequence +from typing import overload +from unittest.mock import Mock + import pytest from transformers import AutoTokenizer from vllm.config import StructuredOutputsConfig, VllmConfig from vllm.config.model import ModelConfig from vllm.config.speculative import SpeculativeConfig +from vllm.reasoning.step3p5_reasoning_parser import Step3p5ReasoningParser from vllm.sampling_params import SamplingParams, StructuredOutputsParams from vllm.v1.request import Request from vllm.v1.structured_output import StructuredOutputManager @@ -341,3 +346,154 @@ def test_trim_reasoning_for_advance(): next_step = [post, post] request.append_output_token_ids(next_step) assert manager.trim_reasoning_for_advance(request, next_step) == next_step + + +@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) +def test_speculative_grammar_filter_rejects_invalid_boundary_suffix(backend): + """Only the grammar-valid answer prefix may cross the commit boundary.""" + tokenizer, manager, request, _, marker = _setup_boundary_request(backend) + reasoning_token = tokenizer.encode(" ")[0] + valid_answer_token = tokenizer.encode("{")[0] + invalid_answer_token = tokenizer.encode("z")[0] + sampled_tokens = [ + reasoning_token, + marker, + valid_answer_token, + invalid_answer_token, + ] + + filtered, rejected = manager.filter_speculative_grammar_tokens( + request, sampled_tokens + ) + + assert filtered == [reasoning_token, marker, valid_answer_token] + assert rejected == 1 + grammar = request.structured_output_request.grammar + assert grammar.validate_tokens([valid_answer_token]) == [valid_answer_token] + + +@pytest.mark.parametrize("backend", ["xgrammar", "guidance"]) +def test_speculative_grammar_filter_rejects_tokens_after_completion(backend): + """A sampled block cannot commit tokens past a completed grammar value.""" + tokenizer, manager, request, _, _ = _setup_boundary_request(backend) + request.structured_output_request.reasoning_ended = True + complete_object = tokenizer.encode("{}") + invalid_suffix = tokenizer.encode("z")[0] + + filtered, rejected = manager.filter_speculative_grammar_tokens( + request, [*complete_object, invalid_suffix] + ) + + assert filtered == complete_object + assert rejected == 1 + + +def test_reasoning_boundary_scan_does_not_copy_committed_history(): + """Boundary detection can expose long history without materializing it.""" + parser_calls = 0 + + class TokenHistory(Sequence[int]): + def __len__(self) -> int: + return 300_000 + + @overload + def __getitem__(self, index: int) -> int: ... + + @overload + def __getitem__(self, index: slice) -> list[int]: ... + + def __getitem__(self, index: int | slice) -> int | list[int]: + raise AssertionError("committed token history was materialized") + + class EndTokenReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + nonlocal parser_calls + parser_calls += 1 + assert len(input_ids) >= 300_000 + return 99 in delta_ids + + reasoner = EndTokenReasoner() + boundary = StructuredOutputManager._find_reasoning_end_offset( + reasoner, TokenHistory(), [10, 99, 20] + ) + + assert boundary == 1 + assert parser_calls == 3 + + +def test_reasoning_boundary_scan_checks_nontransition_block_once(): + """A sampled block without a transition requires one parser call.""" + parser_calls = 0 + + class NoBoundaryReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + nonlocal parser_calls + parser_calls += 1 + return False + + reasoner = NoBoundaryReasoner() + boundary = StructuredOutputManager._find_reasoning_end_offset( + reasoner, [1, 2, 3], [10, 20, 30] + ) + + assert boundary is None + assert parser_calls == 1 + + +def test_reasoning_boundary_scan_preserves_stateful_parser(): + """Pre-commit probing must not consume Step3.5's pending transition.""" + tokenizer = Mock() + tokenizer.get_vocab.return_value = {"": 1, "": 2} + reasoner = Step3p5ReasoningParser(tokenizer) + reasoner._end_token_pending = True + + boundary = StructuredOutputManager._find_reasoning_end_offset( + reasoner, [1, 2, 3], [10, 20] + ) + + assert boundary == 0 + assert reasoner._end_token_pending + + +def test_reasoning_boundary_scan_locates_multi_token_marker(): + """A parser that examines cumulative deltas locates a multi-token marker.""" + + class MultiTokenReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + delta = list(delta_ids) + return any( + delta[index : index + 2] == [20, 30] for index in range(len(delta) - 1) + ) + + boundary = StructuredOutputManager._find_reasoning_end_offset( + MultiTokenReasoner(), [1, 2, 3], [20, 30, 40] + ) + + assert boundary == 1 + + +def test_reasoning_boundary_scan_handles_marker_across_blocks(): + """A multi-token marker may start in history and finish in the new block.""" + + class CrossBlockReasoner: + def is_reasoning_end_streaming( + self, input_ids: Sequence[int], delta_ids: Iterable[int] + ) -> bool: + tokens = list(input_ids) + delta_len = len(list(delta_ids)) + return any( + tokens[index : index + 3] == [7, 8, 9] + for index in range(max(0, len(tokens) - delta_len - 2), len(tokens)) + ) + + boundary = StructuredOutputManager._find_reasoning_end_offset( + CrossBlockReasoner(), [1, 7, 8], [9, 10] + ) + + assert boundary == 0 diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index ae08df9c35b2..e4aed9aee4fa 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1803,9 +1803,13 @@ def update_from_output( scheduled_spec_token_ids = ( scheduler_output.scheduled_spec_decode_tokens.get(req_id) ) + observed_spec_decode = False + num_draft_tokens = 0 + num_accepted = 0 if scheduled_spec_token_ids and ( generated_token_ids or self.num_sampled_tokens_per_step == 0 ): + observed_spec_decode = True num_draft_tokens = len(scheduled_spec_token_ids) num_sampled = self.num_sampled_tokens_per_step num_accepted = max(len(generated_token_ids) - num_sampled, 0) @@ -1823,13 +1827,6 @@ def update_from_output( request.num_computed_tokens -= num_rejected if request.num_output_placeholders > 0: request.num_output_placeholders -= num_rejected - spec_decoding_stats = self.make_spec_decoding_stats( - spec_decoding_stats, - num_draft_tokens=num_draft_tokens, - num_accepted_tokens=num_accepted, - num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, - request_id=req_id, - ) # Free encoder inputs only after the step has actually executed. if request.has_encoder_inputs: @@ -1845,6 +1842,40 @@ def update_from_output( status_before_stop = request.status num_output_tokens_before = len(request._output_token_ids) + if ( + len(new_token_ids) > 1 + and scheduled_spec_token_ids + and request.use_structured_output + and not output_is_stale + ): + new_token_ids, num_grammar_rejected = ( + self.structured_output_manager.filter_speculative_grammar_tokens( + request, new_token_ids + ) + ) + if num_grammar_rejected > 0: + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_grammar_rejected + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_grammar_rejected + # Target-sampled tokens occupy the end of the output block. + # Removing that tail changes the accepted-draft count only + # when the rejected suffix extends into draft positions. + num_accepted -= max( + num_grammar_rejected - self.num_sampled_tokens_per_step, + 0, + ) + assert num_accepted >= 0 + + if observed_spec_decode: + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, + ) + # Check for stop and update request status. if new_token_ids: new_token_ids, stopped = self._update_request_with_output( diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 5ba4ad5b77e7..d56bc7fb3348 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -4,7 +4,8 @@ import multiprocessing from collections.abc import Iterable, Sequence from concurrent.futures import Future, ThreadPoolExecutor -from typing import TYPE_CHECKING +from copy import copy +from typing import TYPE_CHECKING, overload from vllm.config import VllmConfig from vllm.logger import init_logger @@ -32,6 +33,34 @@ logger = init_logger(__name__) +class _TokenSequenceView(Sequence[int]): + """Read-only concatenation that does not copy committed token history.""" + + def __init__(self, prefix: Sequence[int], suffix: Sequence[int]) -> None: + self.prefix = prefix + self.suffix = suffix + + def __len__(self) -> int: + return len(self.prefix) + len(self.suffix) + + @overload + def __getitem__(self, index: int) -> int: ... + + @overload + def __getitem__(self, index: slice) -> list[int]: ... + + def __getitem__(self, index: int | slice) -> int | list[int]: + if isinstance(index, slice): + return [self[i] for i in range(*index.indices(len(self)))] + if index < 0: + index += len(self) + if index < 0 or index >= len(self): + raise IndexError(index) + if index < len(self.prefix): + return self.prefix[index] + return self.suffix[index - len(self.prefix)] + + class StructuredOutputManager: """Engine-level manager for structured output requests.""" @@ -494,6 +523,106 @@ def trim_reasoning_for_advance( return new_token_ids return new_token_ids[num_reasoning:] + @staticmethod + def _find_reasoning_end_offset( + reasoner: "ReasoningParser", + prior_token_ids: Sequence[int], + new_token_ids: list[int], + ) -> int | None: + """Return where a reasoning-end marker completes in a sampled block. + + The first parser call keeps reasoning-only blocks at one check. When + the block contains a transition, cumulative delta prefixes locate + single-token and multi-token markers, including markers that start in + ``prior_token_ids`` and finish in ``new_token_ids``. + + Args: + reasoner: Request-scoped parser that identifies the transition. + prior_token_ids: Tokens committed before the sampled block. + new_token_ids: Accepted tokens awaiting commit. + + Returns: + The zero-based offset where the marker completes, or ``None`` if + the sampled block remains inside reasoning. + """ + complete_tokens = _TokenSequenceView(prior_token_ids, new_token_ids) + if not copy(reasoner).is_reasoning_end_streaming( + complete_tokens, new_token_ids + ): + return None + + for offset in range(len(new_token_ids)): + delta_ids = new_token_ids[: offset + 1] + token_ids = _TokenSequenceView(prior_token_ids, delta_ids) + # Streaming parsers may retain request-local transition state. + # A probe must not consume that state before the normal commit path. + if copy(reasoner).is_reasoning_end_streaming(token_ids, delta_ids): + return offset + + # Some parsers report only that the complete block crossed a boundary. + # Treating the full block as reasoning avoids exposing an unknown suffix + # to the answer grammar. + return len(new_token_ids) - 1 + + def filter_speculative_grammar_tokens( + self, + request: "Request", + new_token_ids: list[int], + ) -> tuple[list[int], int]: + """Validate an accepted speculative block before it is committed. + + Grammar masks cover scheduled draft positions, but an accepted block + can contain an unconstrained token immediately after a reasoning + transition or after the grammar completes. The filter retains the + grammar-valid prefix and reports the trailing tokens that must be + resampled. ``validate_tokens`` restores the grammar state before + returning, so the scheduler's normal commit path remains the only + operation that advances the matcher. + + Args: + request: Request that owns the grammar and reasoning state. + new_token_ids: Accepted tokens awaiting scheduler commit. + + Returns: + The tokens that are safe to commit and the rejected suffix length. + """ + if self.vllm_config.speculative_config is None: + return new_token_ids, 0 + if len(new_token_ids) < 2 or not request.use_structured_output: + return new_token_ids, 0 + + structured_request = request.structured_output_request + if structured_request is None: + return new_token_ids, 0 + grammar = structured_request.grammar + if not isinstance(grammar, StructuredOutputGrammar): + return new_token_ids, 0 + + reasoner = self._get_reasoner(request) + grammar_start = 0 + if ( + reasoner is not None + and not self.enable_in_reasoning + and not structured_request.reasoning_ended + ): + boundary = self._find_reasoning_end_offset( + reasoner, + request.all_token_ids, + new_token_ids, + ) + if boundary is None: + return new_token_ids, 0 + grammar_start = boundary + 1 + + grammar_tokens = new_token_ids[grammar_start:] + if not grammar_tokens: + return new_token_ids, 0 + valid_grammar_tokens = grammar.validate_tokens(grammar_tokens) + rejected = len(grammar_tokens) - len(valid_grammar_tokens) + if rejected == 0: + return new_token_ids, 0 + return new_token_ids[:grammar_start] + valid_grammar_tokens, rejected + def clear_backend(self) -> None: if self.backend is not None: self.backend.destroy() From fd0237e15f714d4b6aa413f915a45473ec2dec5a Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 20 Aug 2026 09:41:57 +0000 Subject: [PATCH 19/52] test(structured-output): use prompt-aware reasoner contract Infernal Invocation exposes prompt inspection through is_reasoning_end_for_prompt. Make the upstream structured-output regression fixture implement the branch contract so it exercises the production method instead of a stale mock interface. Signed-off-by: Martin Vit --- .../v1/structured_output/test_reasoning_structured_output.py | 4 ++-- vllm/v1/structured_output/__init__.py | 4 +--- 2 files changed, 3 insertions(+), 5 deletions(-) diff --git a/tests/v1/structured_output/test_reasoning_structured_output.py b/tests/v1/structured_output/test_reasoning_structured_output.py index ad5f1d5d7951..f5deffe3e248 100644 --- a/tests/v1/structured_output/test_reasoning_structured_output.py +++ b/tests/v1/structured_output/test_reasoning_structured_output.py @@ -15,7 +15,7 @@ class MockReasoner: def __init__(self, tokenizer): - self.is_reasoning_end = Mock(return_value=False) + self.is_reasoning_end_for_prompt = Mock(return_value=False) self.is_reasoning_end_streaming = Mock(return_value=False) @@ -137,7 +137,7 @@ class KwargReasoner: def __init__(self, tokenizer, chat_template_kwargs=None): self.chat_template_kwargs = chat_template_kwargs or {} - def is_reasoning_end(self, input_ids): + def is_reasoning_end_for_prompt(self, input_ids): return not self.chat_template_kwargs.get("enable_thinking", False) manager = StructuredOutputManager(mock_vllm_config) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index d56bc7fb3348..79867280935b 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -409,9 +409,7 @@ def should_fill_bitmask(self, request: "Request") -> bool: # After unifying the `openai_gptoss` and non-`openai_gptoss` styles, # it can be removed. request.structured_output_request.reasoning_ended = ( - reasoner.is_reasoning_end_for_prompt( - request.prompt_token_ids or [] - ) + reasoner.is_reasoning_end_for_prompt(request.prompt_token_ids or []) ) return request.structured_output_request.reasoning_ended return True From 2e8535c70af889a07c42370c28387172c678212f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 20 Aug 2026 09:57:44 +0000 Subject: [PATCH 20/52] fix(dspark): declare compact-RoPE ownership context Type the conditional Kimi compact-RoPE protection scope through the shared context-manager interface. Both the Kimi protection context and the no-op context retain their existing runtime behavior. Signed-off-by: Martin Vit --- vllm/v1/worker/gpu/spec_decode/dspark/utils.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py index cd23ae1bb559..2bfbe39ecfaf 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/utils.py @@ -1,7 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from contextlib import nullcontext +from contextlib import AbstractContextManager, nullcontext import torch.nn as nn @@ -48,6 +48,7 @@ def load_dspark_model(target_model: nn.Module, vllm_config: VllmConfig) -> nn.Mo if hasattr(target_model, "get_language_model") else target_model ) + rope_ownership: AbstractContextManager[None] if getattr(draft_model_config.hf_config, "model_type", None) == "k3_dspark": from vllm.models.kimi_k3.nvidia.dspark_mla import ( protect_k3_compact_rope_sources, From d931e0d45dc9606d5d23554e676afa5aed075e8f Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Thu, 20 Aug 2026 09:59:55 +0000 Subject: [PATCH 21/52] fix(warmup): narrow enabled DSpark debug events The debug branch initializes the event list before every sweep point. Assert that invariant after detaching the list from the model runner so static analysis can verify indexed event access. Profiling and warmup behavior are unchanged. Signed-off-by: Martin Vit --- vllm/v1/worker/gpu/warmup.py | 1 + 1 file changed, 1 insertion(+) diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 6dba53f2f324..3381655e3d3e 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -583,6 +583,7 @@ def _profile_sps_curve( if sps_debug: events = model_runner._sps_debug_events model_runner._sps_debug_events = None + assert events is not None # Skip the warmup iters; report mean verify/draft GPU ms. timed = events[warmup_iters:] if timed: From c27d1179ffaf14c796df0e1996d5a5b071902a8b Mon Sep 17 00:00:00 2001 From: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Date: Fri, 14 Aug 2026 12:03:42 -0400 Subject: [PATCH 22/52] [Model][Spec Decode] Tap the pre-norm AttnRes mixture as the Kimi K3 DFlash aux state (#50487) Signed-off-by: Rahul Chalamala <22563365+rchalamala@users.noreply.github.com> Co-authored-by: Janelle Cai (cherry picked from commit 03a8d0b1ede68efc20b8bbfc14ee4681a62f8d75) --- .../kimi_k3/test_aux_attn_res_stream.py | 197 ++++++++++++++++++ tests/models/kimi_k3/test_eagle3.py | 58 ++++++ vllm/envs.py | 8 + vllm/models/kimi_k3/nvidia/model.py | 87 +++++++- 4 files changed, 349 insertions(+), 1 deletion(-) create mode 100644 tests/models/kimi_k3/test_aux_attn_res_stream.py diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py new file mode 100644 index 000000000000..7227ed6fd9ea --- /dev/null +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -0,0 +1,197 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Which value the DFlash drafter is fed under AttnRes. + +`_capture_aux_hidden_stream` picks the weights it mixes against from one of +three places depending on where the tapped layer sits, and returns the plain +running prefix when the feature is off. The mixture itself is the kernel's +job and is covered by ``test_attn_res.py``; what is asserted here is the +selection, which is the part that can silently feed the drafter the wrong +tensor. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.models.kimi_k3.nvidia import model as k3_model + +END_LAYER = 4 + + +def _weights(tag: float) -> SimpleNamespace: + """A norm/projection pair that is identifiable by value.""" + return SimpleNamespace( + weight=torch.full((2,), tag), + variance_epsilon=tag, + ) + + +def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: + """A stand-in carrying only what the tap reads. + + Constructing the real model needs a distributed init and weights, and none + of it participates in the selection under test. + """ + consumers = [] + for i in range(END_LAYER): + consumers.append( + SimpleNamespace( + self_attention_res_norm=_weights(float(i)), + self_attention_res_proj=SimpleNamespace( + weight=torch.full((1, 2), float(i)) + ), + prev_valid_blocks=i, + ) + ) + return SimpleNamespace( + _aux_attn_res_stream=enabled, + use_attn_res=use_attn_res, + end_layer=END_LAYER, + layers=consumers, + output_attn_res_norm=_weights(99.0), + output_attn_res_proj=SimpleNamespace(weight=torch.full((1, 2), 99.0)), + num_attn_res_blocks=99, + ) + + +@pytest.fixture +def recorder(monkeypatch): + """Replace the kernel so the call it would have made is inspectable.""" + calls = [] + + def _fake_attn_res( + prefix, + delta, + block_residual, + norm_weight, + proj_weight, + output_norm_weight, + **kwargs, + ): + calls.append( + SimpleNamespace( + prefix=prefix, + delta=delta, + block_residual=block_residual, + norm_weight=norm_weight, + proj_weight=proj_weight, + kwargs=kwargs, + ) + ) + return torch.full_like(prefix, -1.0) + + monkeypatch.setattr(k3_model, "attn_res", _fake_attn_res) + return calls + + +def _set_last_rank(monkeypatch, is_last: bool): + monkeypatch.setattr( + k3_model, + "get_pp_group", + lambda: SimpleNamespace(is_last_rank=is_last), + ) + + +def _call(stub, layer_idx, prefix_sum, pending_mlp_out, block_residual): + return k3_model.KimiLinearModel._capture_aux_hidden_stream( + stub, layer_idx, prefix_sum, pending_mlp_out, block_residual + ) + + +@pytest.mark.parametrize( + "enabled,use_attn_res", [(False, True), (True, False), (False, False)] +) +def test_disabled_reproduces_the_plain_residual_sum( + recorder, monkeypatch, enabled, use_attn_res +): + """Off, the tap must be exactly the sum it replaced. + + Both conditions matter. `use_attn_res` is what constructs the norm and + projection weights, so without it the lookups below would raise rather + than fall back. + """ + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + got = _call( + _stub_model(enabled=enabled, use_attn_res=use_attn_res), + 0, + prefix_sum, + pending, + torch.zeros(2), + ) + + torch.testing.assert_close(got, prefix_sum + pending) + assert not recorder, "the kernel must not run when the tap is off" + + +def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): + """The value the next layer reads is the mixture against *its* weights, + so the tap has to reach forward rather than use the current layer's.""" + _set_last_rank(monkeypatch, True) + + _call(_stub_model(enabled=True), 1, torch.zeros(2), None, torch.zeros(2)) + + assert len(recorder) == 1 + call = recorder[0] + # Layer 2's weights, not layer 1's. + torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0)) + assert call.kwargs["num_blocks"] == 2 + + +def test_last_layer_on_the_final_rank_uses_the_output_aggregation( + recorder, monkeypatch +): + """Nothing downstream but the model's own output-side mixture.""" + _set_last_rank(monkeypatch, True) + + _call( + _stub_model(enabled=True), END_LAYER - 1, torch.zeros(2), None, torch.zeros(2) + ) + + assert len(recorder) == 1 + torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0)) + assert recorder[0].kwargs["num_blocks"] == 99 + + +def test_last_layer_of_a_non_final_stage_falls_back(recorder, monkeypatch): + """The consumer lives on the next rank and the output aggregation only + exists on the last one, so there is nothing here to mix against. + + This is the case that would otherwise reach for weights this rank never + constructs. The forward guard is `layer_idx + 1 < end_layer`, where + `end_layer` is the rank's own exclusive bound from `get_pp_indices`, so a + `PPMissingLayer` is unreachable by construction -- the fallback below is + what makes that true rather than merely likely. + """ + _set_last_rank(monkeypatch, False) + prefix_sum = torch.tensor([3.0, 4.0]) + + got = _call( + _stub_model(enabled=True), END_LAYER - 1, prefix_sum, None, torch.zeros(2) + ) + + torch.testing.assert_close(got, prefix_sum) + assert not recorder, "no weights exist on this rank to mix against" + + +def test_pending_mlp_output_is_folded_in_rather_than_passed_as_delta( + recorder, monkeypatch +): + """The kernel writes an applied delta back into the prefix in place, which + would double-add it into the live residual stream, so the pending output + has to arrive already summed into the prefix with `delta` left None.""" + _set_last_rank(monkeypatch, True) + prefix_sum = torch.tensor([1.0, 2.0]) + pending = torch.tensor([0.5, 0.25]) + + _call(_stub_model(enabled=True), 0, prefix_sum, pending, torch.zeros(2)) + + assert len(recorder) == 1 + assert recorder[0].delta is None + torch.testing.assert_close(recorder[0].prefix, prefix_sum + pending) + # And the caller's tensor is not mutated on the way. + torch.testing.assert_close(prefix_sum, torch.tensor([1.0, 2.0])) diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 7af025c91ae4..42c25c028f57 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -22,6 +22,7 @@ def _make_kimi_linear_model() -> KimiLinearModel: object.__setattr__(model, "aux_hidden_state_layers", (2,)) object.__setattr__(model, "use_sequence_parallel", False) object.__setattr__(model, "reuse_attn_res_output", True) + object.__setattr__(model, "use_attn_res", False) return model @@ -190,6 +191,63 @@ def finish_auxiliary_stream(self): ] assert len(aux_hidden_states) == 1 assert aux_hidden_states[0] is projected + + +def test_attn_res_stream_capture_receives_layer_outputs_in_order(monkeypatch): + """Verify the positional contract between ``forward`` and the capture tap.""" + model = _make_kimi_linear_model() + initial_hidden_states = torch.tensor([[1.0, 2.0]]) + layer_hidden_states = torch.tensor([[3.0, 4.0]]) + prefix_sum = torch.tensor([[5.0, 6.0]]) + block_residual = torch.tensor([[[7.0, 8.0]]]) + captured = torch.tensor([[11.0, 12.0]]) + + object.__setattr__(model, "start_layer", 0) + object.__setattr__(model, "end_layer", 1) + object.__setattr__( + model, + "layers", + [Mock(return_value=(layer_hidden_states, prefix_sum, block_residual))], + ) + object.__setattr__(model, "aux_hidden_state_layers", (1,)) + object.__setattr__(model, "use_attn_res", True) + object.__setattr__(model, "num_attn_res_blocks", 1) + object.__setattr__( + model, + "output_attn_res_norm", + SimpleNamespace(weight=torch.ones(2), variance_epsilon=1e-5), + ) + object.__setattr__( + model, + "output_attn_res_proj", + SimpleNamespace(weight=torch.ones(1, 2)), + ) + monkeypatch.setattr( + kimi_model, + "get_pp_group", + lambda: SimpleNamespace(is_first_rank=True, is_last_rank=True), + ) + monkeypatch.setattr(kimi_model, "attn_res", Mock(return_value=torch.zeros(1, 2))) + monkeypatch.setenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "1") + + capture = Mock(return_value=captured) + monkeypatch.setattr(KimiLinearModel, "_capture_aux_hidden_stream", capture) + + _, aux_hidden_states = model.forward( + input_ids=None, + positions=torch.tensor([0]), + intermediate_tensors=None, + inputs_embeds=initial_hidden_states, + ) + + layer_idx, got_prefix, got_pending, got_residual = capture.call_args.args + assert layer_idx == 0 + assert got_prefix is prefix_sum + assert got_pending is layer_hidden_states + assert got_residual is block_residual + torch.testing.assert_close(aux_hidden_states[0], captured) + + def test_kimi_attn_res_workspace_is_reused_and_sliced(): model = _make_kimi_linear_model() object.__setattr__(model, "num_attn_res_blocks", 3) diff --git a/vllm/envs.py b/vllm/envs.py index 8bcd9d0bc7f2..92d8c3e0098a 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -247,6 +247,7 @@ VLLM_USE_FUSED_MOE_GROUPED_TOPK: bool = True VLLM_MOE_SKIP_PADDING: bool = True VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT: bool = False + VLLM_KIMI_K3_AUX_ATTN_RES_STREAM: bool = False VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER: bool = True VLLM_USE_FLASHINFER_MOE_INT4: bool = False VLLM_FLASHINFER_AUTOTUNE_CACHE_DIR: str | None = None @@ -1844,6 +1845,13 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT": lambda: bool( int(os.getenv("VLLM_KIMI_K3_SHARD_SP_SHARED_EXPERT", "0")) ), + # Kimi K3 only, and unrelated to the MoE flags above. Tap the pre-norm + # AttnRes mixture, rather than the post-mixture sum, as the auxiliary + # hidden state handed to a DFlash drafter. This changes the numerics the + # speculator sees, so it is off by default while the effect is measured. + "VLLM_KIMI_K3_AUX_ATTN_RES_STREAM": lambda: bool( + int(os.getenv("VLLM_KIMI_K3_AUX_ATTN_RES_STREAM", "0")) + ), # Allow use of FlashInfer FP8 block-scale GEMM for linear layers. # This uses TensorRT-LLM kernels and requires SM90+ (Hopper). "VLLM_BLOCKSCALE_FP8_GEMM_FLASHINFER": lambda: bool( diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index bbceb40cec19..a853f3b3a7d2 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1891,6 +1891,85 @@ def make_empty_intermediate_tensors( } ) + def _set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None: + super()._set_aux_hidden_state_layers(layers) + if self.use_attn_res: + # Emitted once, at configuration time. Which layers are tapped and + # which convention is in force are the two things you need to + # confirm from a running process, and neither is recoverable from + # the served output. + logger.info_once( + "Kimi-K3 aux hidden capture: layers=%s mode=%s " + "(VLLM_KIMI_K3_AUX_ATTN_RES_STREAM=%d)", + layers, + "attn_res_stream" if self._aux_attn_res_stream else "prefix_only", + int(self._aux_attn_res_stream), + ) + + @property + def _aux_attn_res_stream(self) -> bool: + return envs.VLLM_KIMI_K3_AUX_ATTN_RES_STREAM + + def _capture_aux_hidden_stream( + self, + layer_idx: int, + prefix_sum: torch.Tensor, + pending_mlp_out: torch.Tensor | None, + block_residual: torch.Tensor, + ) -> torch.Tensor: + """Auxiliary feature tapped after ``layer_idx`` under AttnRes. + + The wire between layers only carries the current block's running prefix; + the committed blocks live in the bank. The value the next consumer + actually reads is the pre-norm AttnRes mixture over + ``bank[:num_blocks] + prefix``, which is what the DFlash drafters were + trained against. ``attn_res`` with no delta, no block write and no + output norm computes exactly that and leaves both the prefix and the + bank untouched. + + Folding the pending MLP output into the prefix rather than passing it as + ``delta`` is deliberate: the kernel writes an applied delta back into + the prefix in place, which would double-add it into the live residual + stream. + """ + prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out + # `use_attn_res` is what constructs the norm and projection weights this + # reads; without it there is no mixture to compute and the attribute + # lookups below would raise. + if not (self._aux_attn_res_stream and self.use_attn_res): + return prefix + + if layer_idx + 1 < self.end_layer: + consumer = self.layers[layer_idx + 1] + score_norm = consumer.self_attention_res_norm + score_proj = consumer.self_attention_res_proj + num_blocks = consumer.prev_valid_blocks + elif get_pp_group().is_last_rank: + # Nothing downstream but the model's own output-side aggregation. + score_norm = self.output_attn_res_norm + score_proj = self.output_attn_res_proj + num_blocks = self.num_attn_res_blocks + else: + # Last layer of a non-final pipeline stage: the consumer lives on + # the next rank and the output-side aggregation only exists on the + # last one, so there is nothing here to mix against. Falling back + # to the running prefix keeps the tap defined rather than reaching + # for weights this rank does not construct. + return prefix + + return attn_res( + prefix, + None, + block_residual, + score_norm.weight, + score_proj.weight.squeeze(0), + None, + num_blocks=num_blocks, + block_write_idx=-1, + eps=score_norm.variance_epsilon, + output_norm_eps=0.0, + ) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) @@ -1993,6 +2072,9 @@ def forward( pp_group = get_pp_group() stream_aux_hidden_states = bool( projector is not None + # The streaming projector consumes plain residual sums. DFlash + # AttnRes capture requires the pre-norm mixture computed below. + and not self._aux_attn_res_stream and not self.use_sequence_parallel and pp_group.is_first_rank and pp_group.is_last_rank @@ -2054,7 +2136,10 @@ def forward( projector.accumulate_auxiliary_state(hidden_states, residual) elif self.use_attn_res: assert prefix_sum is not None - aux_hidden_state = prefix_sum + hidden_states + assert residual is not None + aux_hidden_state = self._capture_aux_hidden_stream( + layer_idx, prefix_sum, hidden_states, residual + ) aux_hidden_states.append(aux_hidden_state) else: assert residual is not None From 0a00dacd4f1e6c724c9de1b4ba65b26a2ad2f37c Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 01:31:38 +0000 Subject: [PATCH 23/52] test(kimi-k3): strengthen DFlash capture contracts Verify that disabled AttnRes capture returns before reading unavailable weights and that enabled capture selects both normalization and projection weights from the correct consumer. Document the capture interface parameters and return value. --- .../kimi_k3/test_aux_attn_res_stream.py | 24 ++++++++++++------- vllm/models/kimi_k3/nvidia/model.py | 9 +++++++ 2 files changed, 24 insertions(+), 9 deletions(-) diff --git a/tests/models/kimi_k3/test_aux_attn_res_stream.py b/tests/models/kimi_k3/test_aux_attn_res_stream.py index 7227ed6fd9ea..09b203fc4050 100644 --- a/tests/models/kimi_k3/test_aux_attn_res_stream.py +++ b/tests/models/kimi_k3/test_aux_attn_res_stream.py @@ -34,6 +34,14 @@ def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: Constructing the real model needs a distributed init and weights, and none of it participates in the selection under test. """ + model = SimpleNamespace( + _aux_attn_res_stream=enabled, + use_attn_res=use_attn_res, + end_layer=END_LAYER, + ) + if not use_attn_res: + return model + consumers = [] for i in range(END_LAYER): consumers.append( @@ -45,15 +53,11 @@ def _stub_model(*, enabled: bool, use_attn_res: bool = True) -> SimpleNamespace: prev_valid_blocks=i, ) ) - return SimpleNamespace( - _aux_attn_res_stream=enabled, - use_attn_res=use_attn_res, - end_layer=END_LAYER, - layers=consumers, - output_attn_res_norm=_weights(99.0), - output_attn_res_proj=SimpleNamespace(weight=torch.full((1, 2), 99.0)), - num_attn_res_blocks=99, - ) + model.layers = consumers + model.output_attn_res_norm = _weights(99.0) + model.output_attn_res_proj = SimpleNamespace(weight=torch.full((1, 2), 99.0)) + model.num_attn_res_blocks = 99 + return model @pytest.fixture @@ -139,6 +143,7 @@ def test_taps_the_consumer_layer_when_one_follows(recorder, monkeypatch): call = recorder[0] # Layer 2's weights, not layer 1's. torch.testing.assert_close(call.norm_weight, torch.full((2,), 2.0)) + torch.testing.assert_close(call.proj_weight, torch.full((2,), 2.0)) assert call.kwargs["num_blocks"] == 2 @@ -154,6 +159,7 @@ def test_last_layer_on_the_final_rank_uses_the_output_aggregation( assert len(recorder) == 1 torch.testing.assert_close(recorder[0].norm_weight, torch.full((2,), 99.0)) + torch.testing.assert_close(recorder[0].proj_weight, torch.full((2,), 99.0)) assert recorder[0].kwargs["num_blocks"] == 99 diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index a853f3b3a7d2..c12eeb712004 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1931,6 +1931,15 @@ def _capture_aux_hidden_stream( ``delta`` is deliberate: the kernel writes an applied delta back into the prefix in place, which would double-add it into the live residual stream. + + Args: + layer_idx: Index of the layer that produced the pending MLP output. + prefix_sum: Running prefix for the active AttnRes block. + pending_mlp_out: MLP output to fold into the running prefix, if any. + block_residual: Committed AttnRes block bank for the active rows. + + Returns: + Auxiliary hidden states for the configured DFlash capture mode. """ prefix = prefix_sum if pending_mlp_out is None else prefix_sum + pending_mlp_out # `use_attn_res` is what constructs the norm and projection weights this From 71964b3cc3e7c450971fa1c5142958bd3247dde0 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 00:09:55 +0000 Subject: [PATCH 24/52] Bound Kimi vision RoPE allocation to input grids Compute MoonViT rotary frequencies only for the image grid sizes present in each request instead of materializing the configured 512x512 ceiling. This reduces the measured first-image CUDA allocation peak from 340,018,176 bytes to 1,990,656 bytes for a 36x36 grid while preserving bit-identical CPU and CUDA output. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/kimi_k3/test_vision_warmup.py | 37 +++++++++++++ vllm/model_executor/models/kimi_k25_vit.py | 63 +++++++++++----------- 2 files changed, 67 insertions(+), 33 deletions(-) diff --git a/tests/models/kimi_k3/test_vision_warmup.py b/tests/models/kimi_k3/test_vision_warmup.py index f7e4af39839e..bdd5b88a8f91 100644 --- a/tests/models/kimi_k3/test_vision_warmup.py +++ b/tests/models/kimi_k3/test_vision_warmup.py @@ -13,6 +13,43 @@ ) +def _full_grid_rope_reference( + rope: kimi_k25_vit.Rope2DPosEmbRepeated, + shapes: list[list[int]], +) -> torch.Tensor: + flat_pos = torch.arange(rope.max_height * rope.max_width).float() + x_pos = flat_pos % rope.max_width + y_pos = flat_pos // rope.max_width + dim_range = torch.arange(0, rope.dim, 4).float() + freqs = 1.0 / (rope.theta_base ** (dim_range / rope.dim)) + x_freqs = torch.outer(x_pos, freqs).float() + y_freqs = torch.outer(y_pos, freqs).float() + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) + table = torch.cat( + [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 + ).reshape(rope.max_height, rope.max_width, -1) + return torch.cat( + [table[:h, :w].reshape(-1, rope.dim // 2).repeat(t, 1) for t, h, w in shapes] + ) + + +def test_vision_rope_materializes_only_requested_grids() -> None: + rope = kimi_k25_vit.Rope2DPosEmbRepeated( + dim=32, + max_height=11, + max_width=13, + ) + shapes = [[1, 3, 5], [2, 4, 2], [1, 3, 5]] + + actual = rope.get_freqs_cis(shapes, device=torch.device("cpu")) + expected = _full_grid_rope_reference(rope, shapes) + + assert torch.equal(actual, expected) + assert actual.shape == (46, 16) + assert not hasattr(rope, "freqs_cis") + + def test_warm_vision_position_interpolation(monkeypatch) -> None: model = torch.nn.Sequential( kimi_k25_vit.Learnable2DInterpPosEmbDivided_fixed( diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index a3da0fe82a79..f63df30d8772 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -282,27 +282,28 @@ def extra_repr(self): f"max_width={self.max_width}, theta_base={self.theta_base}" ) - def _precompute_freqs_cis(self, device: torch.device) -> torch.Tensor: - """Calculate the cis(freqs) for each position in the 2D grid.""" - N = self.max_height * self.max_width - flat_pos = torch.arange(0, N).float().to(device) - x_pos = flat_pos % self.max_width - y_pos = flat_pos // self.max_width - dim_range = ( - torch.arange(0, self.dim, 4)[: (self.dim // 4)].float().to(device) - ) # C/4 + def _compute_grid_freqs_cis( + self, height: int, width: int, device: torch.device + ) -> torch.Tensor: + """Calculate rotary frequencies for one requested image grid.""" + x_pos = torch.arange(width, dtype=torch.float32, device=device) + y_pos = torch.arange(height, dtype=torch.float32, device=device) + dim_range = torch.arange(0, self.dim, 4, dtype=torch.float32, device=device)[ + : (self.dim // 4) + ] freqs = 1.0 / (self.theta_base ** (dim_range / self.dim)) - x_freqs = torch.outer(x_pos, freqs).float() # N, C/4 - y_freqs = torch.outer(y_pos, freqs).float() # N, C/4 - x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) # N, C/4 - y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) # N, C/4 - # N, C/4, 2 - freqs_cis = torch.cat( - [x_cis.unsqueeze(dim=-1), y_cis.unsqueeze(dim=-1)], dim=-1 - ) - # max_height, max_width, C/2 - freqs_cis = freqs_cis.reshape(self.max_height, self.max_width, -1) - return freqs_cis + x_freqs = torch.outer(x_pos, freqs).float() + y_freqs = torch.outer(y_pos, freqs).float() + x_cis = torch.polar(torch.ones_like(x_freqs), x_freqs) + y_cis = torch.polar(torch.ones_like(y_freqs), y_freqs) + freqs_cis = torch.stack( + ( + x_cis.unsqueeze(0).expand(height, -1, -1), + y_cis.unsqueeze(1).expand(-1, width, -1), + ), + dim=-1, + ) + return freqs_cis.reshape(height * width, self.dim // 2) def get_freqs_cis( self, grid_thws: torch.Tensor | list[list[int]], device: torch.device @@ -314,11 +315,6 @@ def get_freqs_cis( Returns: freqs_cis: tensor of shape (sum(t * height * width), dim//2) """ - if not hasattr(self, "freqs_cis"): - self.register_buffer( - "freqs_cis", self._precompute_freqs_cis(device), persistent=False - ) - shapes = grid_thws if isinstance(grid_thws, list) else grid_thws.tolist() assert all( 1 <= h <= self.max_height and 1 <= w <= self.max_width for t, h, w in shapes @@ -327,14 +323,15 @@ def get_freqs_cis( self.max_height, self.max_width, ) - freqs_cis = torch.cat( - [ - self.freqs_cis[:h, :w].reshape(-1, self.dim // 2).repeat(t, 1) - for t, h, w in shapes - ], - dim=0, - ) - return freqs_cis + grids: dict[tuple[int, int], torch.Tensor] = {} + result = [] + for t, h, w in shapes: + grid = grids.get((h, w)) + if grid is None: + grid = self._compute_grid_freqs_cis(h, w, device) + grids[(h, w)] = grid + result.append(grid.repeat(t, 1)) + return torch.cat(result, dim=0) class MLP2(nn.Module): From 18d9e2789b1939fd4158b2e35ce4340437995a08 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 02:05:13 +0000 Subject: [PATCH 25/52] Bound Kimi projector transients per image Project independent Kimi vision features separately so MXFP8/Marlin workspace scales with the largest image instead of the sum of all scheduled images. Preserve output order, shape, activation dtype, and numerical results while reducing the measured TP16 three-image transient peak by 32.52 MiB. Co-authored-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/models/kimi_k3/test_vision_projector.py | 42 +++++++++++++++++-- vllm/model_executor/models/kimi_k25_vit.py | 23 +++++----- 2 files changed, 51 insertions(+), 14 deletions(-) diff --git a/tests/models/kimi_k3/test_vision_projector.py b/tests/models/kimi_k3/test_vision_projector.py index 9f2750809a40..c89c79a3b4ce 100644 --- a/tests/models/kimi_k3/test_vision_projector.py +++ b/tests/models/kimi_k3/test_vision_projector.py @@ -80,11 +80,11 @@ def __init__(self): requires_grad=False, ) self.pre_norm = nn.LayerNorm(4, dtype=torch.bfloat16) - self.input_dtype: torch.dtype | None = None + self.inputs: list[torch.Tensor] = [] def forward(self, inputs: torch.Tensor) -> torch.Tensor: - self.input_dtype = inputs.dtype - return inputs + self.inputs.append(inputs) + return inputs + len(self.inputs) def test_kimi_projector_uses_norm_activation_dtype_for_fp8_weights(): @@ -95,9 +95,43 @@ def test_kimi_projector_uses_norm_activation_dtype_for_fp8_weights(): [torch.randn(2, 4), torch.randn(1, 4)], ) - assert projector.input_dtype == torch.bfloat16 + assert len(projector.inputs) == 2 + assert [inputs.shape for inputs in projector.inputs] == [(2, 4), (1, 4)] + assert all(inputs.dtype == torch.bfloat16 for inputs in projector.inputs) assert [output.shape for output in outputs] == [(2, 4), (1, 4)] assert all(output.dtype == torch.bfloat16 for output in outputs) + assert torch.equal(outputs[0], projector.inputs[0] + 1) + assert torch.equal(outputs[1], projector.inputs[1] + 2) + + +def test_kimi_projector_rejects_empty_vision_output(): + with pytest.raises( + ValueError, match="Kimi vision projection requires at least one image feature" + ): + mm_projector_forward(_SerializedFp8Projector(), []) + + +class _DeterministicProjector(nn.Module): + def __init__(self): + super().__init__() + self.pre_norm = nn.LayerNorm(4) + self.linear = nn.Linear(4, 3, bias=False) + + def forward(self, inputs: torch.Tensor) -> torch.Tensor: + return self.linear(self.pre_norm(inputs)) + + +def test_kimi_projector_preserves_batched_projection_results(): + torch.manual_seed(1) + projector = _DeterministicProjector() + inputs = [torch.randn(2, 4), torch.randn(3, 4)] + expected = torch.split(projector(torch.cat(inputs)), [2, 3]) + + outputs = mm_projector_forward(projector, inputs) + + assert len(outputs) == len(expected) + for output, reference in zip(outputs, expected): + torch.testing.assert_close(output, reference) def test_kimi_vision_rope_reuses_packed_qk_buffers(): diff --git a/vllm/model_executor/models/kimi_k25_vit.py b/vllm/model_executor/models/kimi_k25_vit.py index f63df30d8772..6c0bbd1deb33 100644 --- a/vllm/model_executor/models/kimi_k25_vit.py +++ b/vllm/model_executor/models/kimi_k25_vit.py @@ -840,21 +840,24 @@ def prepare_encoder_cudagraph_metadata( @torch.inference_mode() def mm_projector_forward(mm_projector: torch.nn.Module, vt_output: list[torch.Tensor]): - """Apply MM projector to vision tower outputs.""" - num_embedding_list = [x.shape[0] for x in vt_output] - batched = torch.cat(vt_output, dim=0) + """Apply the projector without concatenating independent image features.""" + if not vt_output: + raise ValueError("Kimi vision projection requires at least one image feature") + projector_norm = getattr(mm_projector, "pre_norm", None) if projector_norm is None: projector_norm = getattr(mm_projector, "post_norm", None) projector_dtype = ( - projector_norm.weight.dtype if projector_norm is not None else batched.dtype + projector_norm.weight.dtype if projector_norm is not None else None ) - if batched.dtype != projector_dtype: - batched = batched.to(projector_dtype) - proj_out = mm_projector(batched) - proj_out = proj_out.reshape(-1, proj_out.shape[-1]) - proj_out = torch.split(proj_out, num_embedding_list) - return proj_out + + projected = [] + for image_features in vt_output: + if projector_dtype is not None and image_features.dtype != projector_dtype: + image_features = image_features.to(projector_dtype) + output = mm_projector(image_features) + projected.append(output.reshape(-1, output.shape[-1])) + return tuple(projected) @torch.inference_mode() From 4f34748d5862cf1ba97a588c99bd63956ce4f884 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 08:43:58 +0000 Subject: [PATCH 26/52] [II] Let cache specs own DCP block-table geometry Define token-position DCP shard count on each cache specification and use max_num_blocks_per_req as the worker block-table width contract. Attention caches retain full, partial, or replicated DCP layouts; recurrent caches report one token-position shard and preserve their mode-specific table width. This removes the model runner's cache-type special case while retaining the 1,310-column Mamba align table required by a 1,000,000-token model length with 768-token blocks and seven speculative blocks. Assisted-by: OpenAI Codex Signed-off-by: Martin Vit --- tests/v1/core/test_kv_cache_utils.py | 18 +++++- vllm/v1/kv_cache_interface.py | 94 ++++++++++++++++++++-------- vllm/v1/worker/gpu/model_runner.py | 15 ++--- 3 files changed, 89 insertions(+), 38 deletions(-) diff --git a/tests/v1/core/test_kv_cache_utils.py b/tests/v1/core/test_kv_cache_utils.py index f080d6aa03e3..517bab5fd65d 100644 --- a/tests/v1/core/test_kv_cache_utils.py +++ b/tests/v1/core/test_kv_cache_utils.py @@ -58,6 +58,7 @@ SlidingWindowMLASpec, SlidingWindowSpec, UniformTypeKVCacheSpecs, + get_kv_cache_dcp_shard_count, get_kv_cache_spec_kind, get_kv_cache_spec_sliding_window, ) @@ -197,6 +198,16 @@ def new_mamba_spec( ) +def test_mamba_cache_has_one_dcp_token_position_shard(): + spec = new_mamba_spec(block_size=768, num_speculative_blocks=7) + vllm_config = SimpleNamespace( + cache_config=SimpleNamespace(mamba_cache_mode="align") + ) + + assert get_kv_cache_dcp_shard_count(spec, dcp_world_size=16) == 1 + assert spec.max_num_blocks_per_req(vllm_config, max_len=1_000_000) == 1310 + + def test_unify_kv_cache_spec_page_size_uses_lcm_for_non_divisible_pages(): mimo_spec = FullAttentionSpec( block_size=64, @@ -1309,8 +1320,11 @@ def test_uniform_type_spec_block_table_width_matches_layer_spec( # The runner sizes the block table from the group spec while the metadata # builders are constructed from the per-layer spec, so the aggregate must # report the same width as the layers it wraps. - vllm_config = VllmConfig(model_config=ModelConfig(max_model_len=1024)) - vllm_config.parallel_config.decode_context_parallel_size = dcp_size + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size), + cache_config=SimpleNamespace(mamba_cache_mode="none"), + model_config=SimpleNamespace(max_model_len=1024), + ) if layer_type == "mla": layer_spec = new_mla_spec() elif layer_type == "replicated": diff --git a/vllm/v1/kv_cache_interface.py b/vllm/v1/kv_cache_interface.py index 14f238b4c735..d3ebdb6fe3ec 100644 --- a/vllm/v1/kv_cache_interface.py +++ b/vllm/v1/kv_cache_interface.py @@ -148,6 +148,21 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: """ return cdiv(max_len, self.block_size) + def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int: + """Return the number of unique token-position shards under DCP. + + Cache types that store recurrent or otherwise rank-local state do not + shard that state by token position. Attention cache specifications + override this method because their default layout is DCP-sharded. + """ + configured_dcp = int(dcp_world_size) + if configured_dcp < 1: + raise ValueError( + "Configured decode-context-parallel size must be positive: " + f"{configured_dcp}" + ) + return 1 + def copy_with_new_block_size(self, block_size: int) -> Self: """ Create a new KVCacheSpec from self but replacing the block size. @@ -240,11 +255,38 @@ def real_page_size_bytes(self) -> int: def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: parallel_config = vllm_config.parallel_config - kv_shard_count = get_kv_cache_dcp_shard_count( - self, parallel_config.decode_context_parallel_size + kv_shard_count = self.get_num_dcp_kv_shards( + parallel_config.decode_context_parallel_size ) return cdiv(max_len, self.block_size * kv_shard_count) + def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int: + """Return the configured or explicitly overridden attention shard count.""" + configured_dcp = int(dcp_world_size) + if configured_dcp < 1: + raise ValueError( + "Configured decode-context-parallel size must be positive: " + f"{configured_dcp}" + ) + replicated = bool(getattr(self, "dcp_replicated", False)) + override = getattr(self, "dcp_kv_shard_count", None) + if replicated: + if override not in (None, 1): + raise ValueError( + "dcp_replicated cannot be combined with " + f"dcp_kv_shard_count={override}" + ) + return 1 + if override is None: + return configured_dcp + override = int(override) + if override < 1 or override > configured_dcp or configured_dcp % override != 0: + raise ValueError( + "dcp_kv_shard_count must be a positive divisor of the configured " + f"DCP size, got shards={override}, DCP={configured_dcp}" + ) + return override + @dataclass(frozen=True, kw_only=True) class FullAttentionSpec(AttentionSpec): @@ -398,36 +440,24 @@ def get_kv_cache_dcp_shard_count( dcp_world_size: int, ) -> int: """Return the number of unique DCP token-position shards for a cache group.""" - configured_dcp = int(dcp_world_size) - if configured_dcp < 1: - raise ValueError( - f"Configured decode-context-parallel size must be positive: " - f"{configured_dcp}" - ) - replicated = bool(getattr(spec, "dcp_replicated", False)) - override = getattr(spec, "dcp_kv_shard_count", None) - if replicated: - if override not in (None, 1): - raise ValueError( - f"dcp_replicated cannot be combined with dcp_kv_shard_count={override}" - ) - return 1 - if override is None: - return configured_dcp - override = int(override) - if override < 1 or override > configured_dcp or configured_dcp % override != 0: - raise ValueError( - "dcp_kv_shard_count must be a positive divisor of the configured " - f"DCP size, got shards={override}, DCP={configured_dcp}" - ) - return override + return spec.get_num_dcp_kv_shards(dcp_world_size) def has_nondefault_kv_dcp_layout( spec: KVCacheSpec, dcp_world_size: int, ) -> bool: - return get_kv_cache_dcp_shard_count(spec, dcp_world_size) != int(dcp_world_size) + layer_specs = ( + spec.kv_cache_specs.values() + if isinstance(spec, UniformTypeKVCacheSpecs) + else (spec,) + ) + is_attention_group = all( + isinstance(layer_spec, AttentionSpec) for layer_spec in layer_specs + ) + return is_attention_group and ( + get_kv_cache_dcp_shard_count(spec, dcp_world_size) != int(dcp_world_size) + ) def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec): @@ -1046,6 +1076,18 @@ def max_num_blocks_per_req(self, vllm_config: VllmConfig, max_len: int) -> int: ) return next(iter(widths)) + def get_num_dcp_kv_shards(self, dcp_world_size: int) -> int: + shard_counts = { + spec.get_num_dcp_kv_shards(dcp_world_size) + for spec in self.kv_cache_specs.values() + } + if len(shard_counts) != 1: + raise ValueError( + "All layers in a uniform KV cache group must use the same " + f"number of DCP KV shards, got {sorted(shard_counts)}." + ) + return next(iter(shard_counts)) + @classmethod def is_uniform_type(cls, kv_cache_specs: dict[str, KVCacheSpec]) -> bool: """ diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 16685bd7a2a4..53bfa20a77bd 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -65,7 +65,6 @@ from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask -from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -619,19 +618,15 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) - # One local block covers `block_size * dcp_shard_count` tokens in - # the global sequence. Replicated groups keep the full cache on - # every rank instead. group_cp_size = get_kv_cache_dcp_shard_count(spec, self.dcp_size) group_cp_sizes.append(group_cp_size) - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * group_cp_size + # Cache specifications own their block-table geometry. Attention + # caches account for their token-position DCP shards, while + # recurrent state and replicated attention caches remain unscaled. + max_num_blocks = spec.max_num_blocks_per_req( + self.vllm_config, block_table_max_model_len ) - # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens if isinstance(spec, MambaSpec): - max_num_blocks = ( - max_num_blocks if self.cache_config.enable_prefix_caching else 1 - ) + spec.num_speculative_blocks max_num_blocks = get_block_table_width( max_num_blocks, spec.block_size, token_alignment=None ) From edb1042d476275d1f1230d5db3dddd19cc003cb6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 12:30:13 +0000 Subject: [PATCH 27/52] fix(mamba): resume prefix hits on checkpoint grid Signed-off-by: Martin Vit --- .../worker/test_mamba_hybrid_model_state.py | 35 +++++++++++++++++++ .../worker/gpu/model_states/mamba_hybrid.py | 12 ++++++- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/tests/v1/worker/test_mamba_hybrid_model_state.py b/tests/v1/worker/test_mamba_hybrid_model_state.py index 545d23f90912..ad6b2e712c04 100644 --- a/tests/v1/worker/test_mamba_hybrid_model_state.py +++ b/tests/v1/worker/test_mamba_hybrid_model_state.py @@ -1,10 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch from vllm.platforms import current_platform +from vllm.v1.core.sched.output import NewRequestData from vllm.v1.worker.gpu.model_states.mamba_hybrid import MambaHybridModelState @@ -27,3 +30,35 @@ def test_postprocess_state_scalar_with_int32_mapping( [expected_value, 9, expected_value, 9], dtype=torch.int32, device="cuda" ) torch.testing.assert_close(state.num_accepted_tokens_gpu, expected) + + +@pytest.mark.parametrize( + ("num_computed_tokens", "expected_state_index"), + [(0, -1), (110_592, 8), (110_593, 9)], +) +def test_prefix_hit_uses_mamba_checkpoint_cadence( + num_computed_tokens: int, expected_state_index: int +) -> None: + """A resumed request indexes recurrent checkpoints, not attention pages.""" + state = object.__new__(MambaHybridModelState) + state.rope_state = None + state._align_mode = True + state.cache_config = SimpleNamespace(block_size=768, mamba_block_size=12_288) + state._mamba_block_size = 12_288 + state._mamba_state_idx_gpu = torch.zeros(1, dtype=torch.int32) + state.num_accepted_tokens_gpu = torch.full((1,), 9, dtype=torch.int32) + request = NewRequestData( + req_id="prefix-hit", + prompt_token_ids=[], + mm_features=[], + sampling_params=None, + pooling_params=None, + block_ids=(), + num_computed_tokens=num_computed_tokens, + lora_request=None, + ) + + state.add_request(0, request) + + assert state._mamba_state_idx_gpu.item() == expected_state_index + assert state.num_accepted_tokens_gpu.item() == 1 diff --git a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py index e72e3afbcb11..6ceb8371b0a3 100644 --- a/vllm/v1/worker/gpu/model_states/mamba_hybrid.py +++ b/vllm/v1/worker/gpu/model_states/mamba_hybrid.py @@ -73,6 +73,12 @@ def __init__( # columns and the running state_idx are kept GPU-resident. self._align_mode = self.cache_config.mamba_cache_mode == "align" if self._align_mode: + # The physical attention page and the logical recurrent-state + # checkpoint can have different token widths. Prefix-hit requests + # must resume from the recurrent-state grid used by MambaSpec. + self._mamba_block_size = ( + self.cache_config.mamba_block_size or self.cache_config.block_size + ) self._mamba_state_idx_gpu = torch.zeros( self.max_num_reqs, dtype=torch.int32, device=self.device ) @@ -93,7 +99,7 @@ def add_request(self, req_index: int, new_req_data: NewRequestData) -> None: if self._align_mode: # Seed the running state block from the resumed/prefilled position. self._mamba_state_idx_gpu[req_index].fill_( - (new_req_data.num_computed_tokens - 1) // self.cache_config.block_size + (new_req_data.num_computed_tokens - 1) // self._mamba_block_size ) def _get_mamba_group_info( @@ -109,6 +115,10 @@ def _get_mamba_group_info( specs.append(spec) assert specs, "no mamba layers in the model" assert all(specs[0] == s for s in specs) + assert specs[0].block_size == self._mamba_block_size, ( + "Mamba state migration and cache allocation must use the same " + "checkpoint cadence" + ) self._mamba_group_ids = group_ids self._mamba_spec = specs[0] return self._mamba_group_ids, self._mamba_spec From 832b1c929e2c8cedca261817396533eb6d51b266 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 15:41:11 +0000 Subject: [PATCH 28/52] fix(multimodal): gather uneven vision shards without TP padding Gather each tensor-parallel vision shard at its produced row count instead of padding every rank to the largest shard. This preserves embedding order and the uniform-size fast path while preventing the transient allocation from scaling with TP size when a request contains fewer images than ranks. Validate zero-length PyNccl inputs, single-image output parity, empty inputs, uneven four-GPU assignments, and multi-image assignments. A TP16 Kimi-K3-shaped harness reduces the collective output from 224 MiB to 14 MiB per GPU with bit-exact gathered content. Signed-off-by: Martin Vit --- tests/distributed/test_pynccl.py | 4 +- vllm/distributed/communication_op.py | 10 ++++ vllm/model_executor/models/vision.py | 79 ++++++++++++---------------- 3 files changed, 46 insertions(+), 47 deletions(-) diff --git a/tests/distributed/test_pynccl.py b/tests/distributed/test_pynccl.py index 62eba4843a0b..dd2be5cfd6a3 100644 --- a/tests/distributed/test_pynccl.py +++ b/tests/distributed/test_pynccl.py @@ -257,7 +257,9 @@ def all_gatherv_worker_fn(): device = f"cuda:{pynccl_comm.rank}" assert world_size <= 8 - sizes = [81, 20, 57, 52, 81, 5, 49, 49][:world_size] + # A zero-length rank is required when fewer multimodal inputs than TP + # ranks are distributed across the model-parallel group. + sizes = [81, 0, 57, 52, 81, 5, 49, 49][:world_size] num_elems = sizes[rank] tensor = torch.arange(num_elems, dtype=torch.float32, device=device) + rank * 100 result = torch.zeros(sum(sizes), dtype=torch.float32, device=device) diff --git a/vllm/distributed/communication_op.py b/vllm/distributed/communication_op.py index 1bc270f3084f..33abaa45ed5d 100644 --- a/vllm/distributed/communication_op.py +++ b/vllm/distributed/communication_op.py @@ -26,6 +26,16 @@ def tensor_model_parallel_all_gather( return get_tp_group().all_gather(input_, dim) +def tensor_model_parallel_all_gatherv( + input_: torch.Tensor, sizes: list[int], dim: int = 0 +) -> torch.Tensor: + """All-gather variable-length tensor slices across the model-parallel group.""" + tp_group = get_tp_group() + if tp_group.world_size == 1: + return input_ + return tp_group.all_gatherv(input_, dim=dim, sizes=sizes) + + def tensor_model_parallel_reduce_scatter( input_: torch.Tensor, dim: int = -1 ) -> torch.Tensor: diff --git a/vllm/model_executor/models/vision.py b/vllm/model_executor/models/vision.py index ff62b34ec787..8598ac4c6386 100644 --- a/vllm/model_executor/models/vision.py +++ b/vllm/model_executor/models/vision.py @@ -17,6 +17,7 @@ get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, tensor_model_parallel_all_gather, + tensor_model_parallel_all_gatherv, ) from vllm.logger import init_logger from vllm.platforms import current_platform @@ -442,9 +443,8 @@ def run_dp_sharded_mrope_vision_model( # Get load balancing assignment with all metadata # image_to_tp_rank = [0, 2, 1, 3] # gpu_sample_counts = [1, 3] - # grouped_pixel_values_len = [1000, 350] - (image_to_tp_rank, gpu_sample_counts, grouped_pixel_values_len) = ( - get_load_balance_assignment(patches_per_image, tp_size) + (image_to_tp_rank, gpu_sample_counts, _) = get_load_balance_assignment( + patches_per_image, tp_size ) # cu_gpu_sample_counts = [0, 1, 4] @@ -481,11 +481,23 @@ def run_dp_sharded_mrope_vision_model( vision_model.spatial_merge_size * vision_model.spatial_merge_size ) - # Find the max length across all ranks - # The output embedding of every DP rank has to be - # padded to this length for tensor_model_parallel_all_gather - # to work - max_len_per_rank = max(grouped_pixel_values_len) // embed_dim_reduction_factor + patches_per_output_image = [ + patch_size // embed_dim_reduction_factor for patch_size in patches_per_image + ] + output_lengths_per_rank = [] + rank_image_offset = 0 + for count in gpu_sample_counts: + rank_image_indices = image_to_tp_rank[ + rank_image_offset : rank_image_offset + count + ] + output_lengths_per_rank.append( + sum(patches_per_output_image[i] for i in rank_image_indices) + ) + rank_image_offset += count + + if not grid_thw_list: + return () + local_grid_thw_list = [grid_thw_list[i] for i in image_idxs_local] # Run the vision model on the local pixel_values_local @@ -514,46 +526,21 @@ def run_dp_sharded_mrope_vision_model( dtype=pixel_values.dtype, ) - # Pad the output based on max_len_per_rank - # for tensor_model_parallel_all_gather to work - current_len = image_embeds_local.shape[0] - if current_len < max_len_per_rank: - padding_size = max_len_per_rank - current_len - if rope_type == "rope_2d": - padding = torch.empty( - ( - padding_size, - image_embeds_local.shape[1], - image_embeds_local.shape[2], - ), - dtype=image_embeds_local.dtype, - device=image_embeds_local.device, - ) - else: - padding = torch.empty( - (padding_size, image_embeds_local.shape[1]), - dtype=image_embeds_local.dtype, - device=image_embeds_local.device, - ) - image_embeds_local_padded = torch.cat([image_embeds_local, padding], dim=0) - else: - image_embeds_local_padded = image_embeds_local - - # Do all_gather to collect embeddings from all ranks - gathered_embeds = tensor_model_parallel_all_gather(image_embeds_local_padded, dim=0) - - # Remove padding and reconstruct per-rank embeddings - rank_embeddings = list[torch.Tensor]() - for rank in range(tp_size): - start_idx = rank * max_len_per_rank - end_idx = start_idx + ( - grouped_pixel_values_len[rank] // embed_dim_reduction_factor + expected_local_len = output_lengths_per_rank[tp_rank_local] + if image_embeds_local.shape[0] != expected_local_len: + raise ValueError( + "Vision encoder output length does not match the image-grid metadata: " + f"rank {tp_rank_local} produced {image_embeds_local.shape[0]} rows, " + f"expected {expected_local_len}" ) - rank_embeddings.append(gathered_embeds[start_idx:end_idx]) - patches_per_output_image = [ - (patch_size // embed_dim_reduction_factor) for patch_size in patches_per_image - ] + # Gather only the rows produced by each rank. Padding every rank to the + # largest shard can multiply the transient allocation by the TP size when + # a request contains fewer images than tensor-parallel ranks. + gathered_embeds = tensor_model_parallel_all_gatherv( + image_embeds_local, sizes=output_lengths_per_rank, dim=0 + ) + rank_embeddings = list(gathered_embeds.split(output_lengths_per_rank, dim=0)) # Reconstruct embeddings in the original order original_order_embeddings = [None] * len(grid_thw_list) From d19ef454b57b0800a4eb34785cb0964bd6c6da61 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 22 Aug 2026 15:33:12 +0000 Subject: [PATCH 29/52] fix(kimi-k3): preserve final AttnRes block Signed-off-by: Martin Vit --- tests/models/kimi_k3/test_eagle3.py | 26 ++++++++++++++++++++++++++ vllm/models/kimi_k3/nvidia/model.py | 14 +++++++++++++- 2 files changed, 39 insertions(+), 1 deletion(-) diff --git a/tests/models/kimi_k3/test_eagle3.py b/tests/models/kimi_k3/test_eagle3.py index 7af025c91ae4..ebb392d55cea 100644 --- a/tests/models/kimi_k3/test_eagle3.py +++ b/tests/models/kimi_k3/test_eagle3.py @@ -273,6 +273,7 @@ def _make_attn_res_decoder_layer(*, block_write: bool): object.__setattr__(layer, "use_attn_res", True) object.__setattr__(layer, "reuse_attn_res_output", True) object.__setattr__(layer, "is_block_write_layer", block_write) + object.__setattr__(layer, "is_final_block_write_layer", False) object.__setattr__(layer, "block_write_idx", 0) object.__setattr__(layer, "prev_valid_blocks", 0) object.__setattr__( @@ -361,3 +362,28 @@ def record_attn_res(*args, **kwargs): assert regular_hidden is attention_output assert regular_prefix is prefix assert outputs == [prefix, attention_output] + + +def test_kimi_post_attn_norm_preserves_final_committed_block(monkeypatch): + prefix = torch.randn(2, 4) + attention_output = torch.randn(2, 4) + blocks = torch.randn(2, 3, 4) + allocated_output = torch.randn(2, 4) + outputs = [] + + def record_attn_res(*args, **kwargs): + outputs.append(kwargs["output"]) + return allocated_output + + monkeypatch.setattr(kimi_model, "attn_res", record_attn_res) + layer = _make_attn_res_decoder_layer(block_write=True) + object.__setattr__(layer, "is_final_block_write_layer", True) + + hidden, next_prefix, next_blocks = layer._post_attn_norm( + attention_output, blocks, prefix + ) + + assert hidden is allocated_output + assert next_prefix is attention_output + assert next_blocks is blocks + assert outputs == [None] diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index bbceb40cec19..0f5c91108858 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -1579,6 +1579,11 @@ def __init__( self.attn_res_block_size = attn_res_block_size self.is_block_write_layer = layer_idx % self.attn_res_block_size == 0 self.block_write_idx = layer_idx // self.attn_res_block_size + self.is_final_block_write_layer = ( + self.is_block_write_layer + and self.block_write_idx + == cdiv(config.num_hidden_layers, self.attn_res_block_size) - 1 + ) self.prev_valid_blocks = cdiv(layer_idx, self.attn_res_block_size) self.self_attention_res_norm = RMSNorm( config.hidden_size, eps=config.rms_norm_eps @@ -1694,7 +1699,14 @@ def _post_attn_norm( assert prefix_sum is not None if self.is_block_write_layer: - output = prefix_sum if self.reuse_attn_res_output else None + # The old prefix becomes the last committed residual block at the + # final block boundary. It must remain immutable for every later + # AttnRes mixture and therefore cannot also hold the new delta. + output = ( + prefix_sum + if self.reuse_attn_res_output and not self.is_final_block_write_layer + else None + ) prefix_sum = hidden_states prefix_delta = None else: From 7aa4d61593c4f26f44b1ae1b00c040f13408afa1 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 22 Aug 2026 10:28:21 +0000 Subject: [PATCH 30/52] fix(attention): preserve aliased LSE during state merge Cache each head's prefix and suffix log-sum-exp values before any output write when the thread group fits inside a CUDA block. This preserves chunked-attention accumulators that pass the running LSE tensor as both prefix input and output destination, while retaining the direct-load path for head groups that cross block boundaries. Index all cached values through the declared tensor strides.\n\nAdd exact in-place versus disjoint-output coverage for the six-head, 128-element MLA geometry at 256 and 4096 tokens.\n\nThe shared-memory loading structure adapts vLLM PR #45778 (commit c71576fd587a93a7568ff40c004a685231415d37) to the strided-LSE kernel contract.\n\nCo-authored-by: nicole-lihui Signed-off-by: Martin Vit --- .../attention/merge_attn_states.cu | 46 ++++++++++++---- .../attention/test_merge_attn_states.py | 54 +++++++++++++++++++ 2 files changed, 91 insertions(+), 9 deletions(-) diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index cc89397f68a1..b15f147ba88d 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -39,15 +39,36 @@ __global__ void merge_attn_states_kernel( const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x; const uint token_head_threads = num_tokens * num_heads * threads_per_head; - if (global_idx >= token_head_threads) return; - - // global_idx -> token_idx + head_idx + pack_idx + // Derive indices before the block barrier so every thread reaches it. const uint token_head_idx = global_idx / threads_per_head; const uint pack_idx = global_idx % threads_per_head; - const uint token_idx = token_head_idx / num_heads; const uint head_idx = token_head_idx % num_heads; + // A running chunked-attention LSE may be both prefix_lse and output_lse. + // Load it once per head before any thread can overwrite the destination. + // Groups that cross block boundaries retain independent global loads. + __shared__ float shared_prefix_lse[NUM_THREADS]; + __shared__ float shared_suffix_lse[NUM_THREADS]; + const bool group_fits_in_block = NUM_THREADS % threads_per_head == 0; + const bool is_valid = global_idx < token_head_threads; + const uint group_idx = threadIdx.x / threads_per_head; + + if (group_fits_in_block) { + if (is_valid && pack_idx == 0 && token_idx < prefix_num_tokens) { + shared_prefix_lse[group_idx] = + prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + shared_suffix_lse[group_idx] = + suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; + } + __syncthreads(); + if (!is_valid) return; + } else if (!is_valid) { + return; + } + const uint pack_offset = pack_idx * pack_size; // (0~15)*8, etc. const uint src_head_offset = token_idx * num_heads * prefix_head_stride + head_idx * prefix_head_stride; @@ -95,11 +116,18 @@ __global__ void merge_attn_states_kernel( return; } - // For tokens within prefix range, merge prefix and suffix - float p_lse = prefix_lse[head_idx * prefix_lse_head_stride + - token_idx * prefix_lse_token_stride]; - float s_lse = suffix_lse[head_idx * suffix_lse_head_stride + - token_idx * suffix_lse_token_stride]; + // For tokens within prefix range, merge prefix and suffix. + float p_lse; + float s_lse; + if (group_fits_in_block) { + p_lse = shared_prefix_lse[group_idx]; + s_lse = shared_suffix_lse[group_idx]; + } else { + p_lse = prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + s_lse = suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; + } p_lse = std::isinf(p_lse) ? -std::numeric_limits::infinity() : p_lse; s_lse = std::isinf(s_lse) ? -std::numeric_limits::infinity() : s_lse; diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index 1394ea9df405..fbf109c85dc6 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -99,6 +99,60 @@ def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: assert not output.isnan().any() +@pytest.mark.parametrize("num_tokens", [256, 4096]) +@torch.inference_mode() +def test_merge_attn_states_cuda_inplace_accumulator(num_tokens: int) -> None: + """The CUDA kernel supports a running partial as input and destination. + + Chunked attention folds each suffix partial into one prefix allocation. + Both the attention output and its log-sum-exp tensor therefore alias their + corresponding destinations. The result must match a merge into disjoint + output allocations exactly, including at Kimi-K3's MLA head geometry. + """ + if not current_platform.is_cuda(): + pytest.skip("The custom merge-attention kernel requires CUDA") + + torch.manual_seed(0) + num_heads = 6 + head_size = 128 + shape = (num_tokens, num_heads, head_size) + + prefix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + suffix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda") + prefix_lse = torch.randn( + (num_heads, num_tokens), dtype=torch.float32, device="cuda" + ) + suffix_lse = torch.randn( + (num_heads, num_tokens), dtype=torch.float32, device="cuda" + ) + + reference_output = torch.empty_like(prefix_output) + reference_lse = torch.empty_like(prefix_lse) + merge_attn_states_cuda( + reference_output, + prefix_output, + prefix_lse, + suffix_output, + suffix_lse, + reference_lse, + ) + + inplace_output = prefix_output.clone() + inplace_lse = prefix_lse.clone() + merge_attn_states_cuda( + inplace_output, + inplace_output, + inplace_lse, + suffix_output, + suffix_lse, + inplace_lse, + ) + torch.accelerator.synchronize() + + torch.testing.assert_close(inplace_output, reference_output, rtol=0, atol=0) + torch.testing.assert_close(inplace_lse, reference_lse, rtol=0, atol=0) + + def generate_markdown_table(): global all_case_info table_header = ( From 7671c6660f7ceb368acaad6b7c2d9d5932403ac6 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Sat, 22 Aug 2026 11:30:12 +0000 Subject: [PATCH 31/52] fix(attention): align merge blocks to head groups Choose a block size that is an integer multiple of the input packs required for one attention head. This keeps every token-head group within one CUDA block, so all threads read an aliased prefix LSE before the group writes its merged destination. Reject head sizes that require more than the kernel's 128-thread limit and cover the six-head, 192-element Kimi-K3 chunked-context geometry in the in-place accumulator test. Signed-off-by: Martin Vit --- .../attention/merge_attn_states.cu | 55 +++++++++---------- .../attention/test_merge_attn_states.py | 10 +++- 2 files changed, 32 insertions(+), 33 deletions(-) diff --git a/csrc/libtorch_stable/attention/merge_attn_states.cu b/csrc/libtorch_stable/attention/merge_attn_states.cu index b15f147ba88d..8af7f6846b41 100644 --- a/csrc/libtorch_stable/attention/merge_attn_states.cu +++ b/csrc/libtorch_stable/attention/merge_attn_states.cu @@ -36,7 +36,7 @@ __global__ void merge_attn_states_kernel( const uint pack_size = 16 / sizeof(scalar_t); const uint threads_per_head = head_size / pack_size; - const uint global_idx = blockIdx.x * NUM_THREADS + threadIdx.x; + const uint global_idx = blockIdx.x * blockDim.x + threadIdx.x; const uint token_head_threads = num_tokens * num_heads * threads_per_head; // Derive indices before the block barrier so every thread reaches it. @@ -46,28 +46,23 @@ __global__ void merge_attn_states_kernel( const uint head_idx = token_head_idx % num_heads; // A running chunked-attention LSE may be both prefix_lse and output_lse. - // Load it once per head before any thread can overwrite the destination. - // Groups that cross block boundaries retain independent global loads. + // The launcher aligns block boundaries to complete head groups, allowing + // every group to load its LSE values before any thread overwrites them. __shared__ float shared_prefix_lse[NUM_THREADS]; __shared__ float shared_suffix_lse[NUM_THREADS]; - const bool group_fits_in_block = NUM_THREADS % threads_per_head == 0; const bool is_valid = global_idx < token_head_threads; const uint group_idx = threadIdx.x / threads_per_head; - if (group_fits_in_block) { - if (is_valid && pack_idx == 0 && token_idx < prefix_num_tokens) { - shared_prefix_lse[group_idx] = - prefix_lse[head_idx * prefix_lse_head_stride + - token_idx * prefix_lse_token_stride]; - shared_suffix_lse[group_idx] = - suffix_lse[head_idx * suffix_lse_head_stride + - token_idx * suffix_lse_token_stride]; - } - __syncthreads(); - if (!is_valid) return; - } else if (!is_valid) { - return; + if (is_valid && pack_idx == 0 && token_idx < prefix_num_tokens) { + shared_prefix_lse[group_idx] = + prefix_lse[head_idx * prefix_lse_head_stride + + token_idx * prefix_lse_token_stride]; + shared_suffix_lse[group_idx] = + suffix_lse[head_idx * suffix_lse_head_stride + + token_idx * suffix_lse_token_stride]; } + __syncthreads(); + if (!is_valid) return; const uint pack_offset = pack_idx * pack_size; // (0~15)*8, etc. const uint src_head_offset = token_idx * num_heads * prefix_head_stride + @@ -117,17 +112,8 @@ __global__ void merge_attn_states_kernel( } // For tokens within prefix range, merge prefix and suffix. - float p_lse; - float s_lse; - if (group_fits_in_block) { - p_lse = shared_prefix_lse[group_idx]; - s_lse = shared_suffix_lse[group_idx]; - } else { - p_lse = prefix_lse[head_idx * prefix_lse_head_stride + - token_idx * prefix_lse_token_stride]; - s_lse = suffix_lse[head_idx * suffix_lse_head_stride + - token_idx * suffix_lse_token_stride]; - } + float p_lse = shared_prefix_lse[group_idx]; + float s_lse = shared_suffix_lse[group_idx]; p_lse = std::isinf(p_lse) ? -std::numeric_limits::infinity() : p_lse; s_lse = std::isinf(s_lse) ? -std::numeric_limits::infinity() : s_lse; @@ -335,10 +321,19 @@ void merge_attn_states_launcher( // Process one pack elements per thread. for float, the // pack_size is 4 for half/bf16, the pack_size is 8. const uint threads_per_head = head_size / pack_size; + STD_TORCH_CHECK( + threads_per_head <= NUM_THREADS, + "headsize requires more threads than the merge kernel block supports: ", + head_size); const uint total_threads = num_tokens * num_heads * threads_per_head; + // Keep each token-head group inside one block. This is required when + // output_lse aliases prefix_lse because the whole group must read the input + // LSE before its first thread writes the merged value. + const uint block_threads = + (NUM_THREADS / threads_per_head) * threads_per_head; - dim3 block(NUM_THREADS); - dim3 grid((total_threads + NUM_THREADS - 1) / NUM_THREADS); + dim3 block(block_threads); + dim3 grid((total_threads + block_threads - 1) / block_threads); const torch::stable::accelerator::DeviceGuard device_guard( prefix_output.get_device_index()); diff --git a/tests/kernels/attention/test_merge_attn_states.py b/tests/kernels/attention/test_merge_attn_states.py index fbf109c85dc6..4f8eba22d231 100644 --- a/tests/kernels/attention/test_merge_attn_states.py +++ b/tests/kernels/attention/test_merge_attn_states.py @@ -100,21 +100,25 @@ def test_merge_attn_states_both_empty(merge_fn, output_dtype) -> None: @pytest.mark.parametrize("num_tokens", [256, 4096]) +@pytest.mark.parametrize("head_size", [128, 192, 512]) @torch.inference_mode() -def test_merge_attn_states_cuda_inplace_accumulator(num_tokens: int) -> None: +def test_merge_attn_states_cuda_inplace_accumulator( + num_tokens: int, head_size: int +) -> None: """The CUDA kernel supports a running partial as input and destination. Chunked attention folds each suffix partial into one prefix allocation. Both the attention output and its log-sum-exp tensor therefore alias their corresponding destinations. The result must match a merge into disjoint - output allocations exactly, including at Kimi-K3's MLA head geometry. + output allocations exactly. The 192-element case is Kimi-K3's chunked + context-merge geometry and does not divide the kernel's 128-thread limit; + the 512-element case spans multiple warps per head. """ if not current_platform.is_cuda(): pytest.skip("The custom merge-attention kernel requires CUDA") torch.manual_seed(0) num_heads = 6 - head_size = 128 shape = (num_tokens, num_heads, head_size) prefix_output = torch.randn(shape, dtype=torch.bfloat16, device="cuda") From 62bead59e89eb6b8cc4f117d00832d4d701ea670 Mon Sep 17 00:00:00 2001 From: Martin Vit Date: Fri, 21 Aug 2026 20:36:54 +0000 Subject: [PATCH 32/52] fix(kimi-k3): bound chunked MLA output lifetime Copy the semantic suffix output into caller-owned storage before allocating the padded context result, then reuse consumed query bytes for the compact context result. This removes one concurrent 192 MiB FlashAttention output at the 4,096-token Kimi-K3 prefill shape without changing attention arithmetic, scheduler capacity, or KV allocation.\n\nValidated with BF16 and FP8 storage tests, exact-shape CUDA alias and graph replay harnesses, and a 237k-token full-model TP16/DCP16 prefill that previously failed in FlashAttention allocation. Signed-off-by: Martin Vit --- tests/models/kimi_k3/test_mla_padding.py | 28 ++++++++++++++++++ vllm/models/kimi_k3/nvidia/mla.py | 37 +++++++++++++++++++++--- 2 files changed, 61 insertions(+), 4 deletions(-) diff --git a/tests/models/kimi_k3/test_mla_padding.py b/tests/models/kimi_k3/test_mla_padding.py index 8a2d859b931b..007556801152 100644 --- a/tests/models/kimi_k3/test_mla_padding.py +++ b/tests/models/kimi_k3/test_mla_padding.py @@ -3,6 +3,7 @@ from types import SimpleNamespace +import pytest import torch @@ -122,6 +123,33 @@ def write_active_prefill(*args): torch.testing.assert_close(output[2:], torch.zeros_like(output[2:])) +@pytest.mark.parametrize("query_dtype", [torch.bfloat16, torch.float8_e4m3fn]) +def test_kimi_mla_context_output_reuses_consumed_query_bytes(query_dtype): + from vllm.models.kimi_k3.nvidia import mla + + query = torch.empty((4, 2, 256), dtype=query_dtype) + output = torch.randn((4, 2, 128), dtype=torch.bfloat16) + + compact = mla._reuse_consumed_query_for_context_output(query, output) + compact.copy_(output) + + assert compact.data_ptr() == query.data_ptr() + assert compact.shape == output.shape + assert compact.dtype == output.dtype + assert compact.is_contiguous() + torch.testing.assert_close(compact, output) + + +def test_kimi_mla_context_output_rejects_insufficient_query_storage(): + from vllm.models.kimi_k3.nvidia import mla + + query = torch.empty((4, 2, 64), dtype=torch.bfloat16) + output = torch.empty((4, 2, 128), dtype=torch.bfloat16) + + with pytest.raises(ValueError, match="too small"): + mla._reuse_consumed_query_for_context_output(query, output) + + def test_kimi_mla_caller_output_selection_preserves_decode_and_sp_paths(): from vllm.models.kimi_k3.nvidia import mla diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 8e107646ad66..766bd1570bfc 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -151,6 +151,24 @@ def _restore_merged_output_order( ) +def _reuse_consumed_query_for_context_output( + query: torch.Tensor, + output: torch.Tensor, +) -> torch.Tensor: + """Return contiguous semantic-output storage backed by a consumed query.""" + if not query.is_contiguous(): + raise ValueError("Kimi-K3 MLA prefill query storage must be contiguous") + required_bytes = output.numel() * output.element_size() + query_bytes = query.view(torch.uint8).flatten() + if query_bytes.numel() < required_bytes: + raise ValueError( + "Kimi-K3 MLA prefill query storage is too small for compact context " + f"output: available={query_bytes.numel()} bytes, " + f"required={required_bytes} bytes" + ) + return query_bytes[:required_bytes].view(output.dtype).view_as(output) + + class KimiShardedMergedColumnParallelLinear(MergedColumnParallelLinear): """Merged column projection with one gather and logical-shard reorder.""" @@ -1066,6 +1084,16 @@ def _forward_prefill_fused( ) if has_context: + suffix_output, suffix_lse = output_prefill + out = out.view(-1, self.num_local_heads, self.v_head_dim) + # FlashAttention 2 pads Kimi-K3's 128-wide V to the 256-wide + # query/key head dimension. Preserve only the semantic V slice in + # caller-owned output storage before context attention allocates + # its equally large padded result. The merge kernel supports + # output aliasing its suffix input, so both padded results never + # need to be live at the same time. + out.copy_(suffix_output[..., : self.v_head_dim]) + del output_prefill, suffix_output if self.dcp_world_size > 1: context_output, context_lse = ( self.impl._context_parallel_compute_prefill_context( # type: ignore[attr-defined] @@ -1080,13 +1108,14 @@ def _forward_prefill_fused( context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] q, self._attn_read_kv_cache(), attn_metadata, self._k_scale ) - suffix_output, suffix_lse = output_prefill - out = out.view(-1, self.num_local_heads, self.v_head_dim) + compact_context_output = _reuse_consumed_query_for_context_output(q, out) + compact_context_output.copy_(context_output[..., : self.v_head_dim]) + del context_output merge_attn_states( output=out, - prefix_output=context_output[..., : self.v_head_dim], + prefix_output=compact_context_output, prefix_lse=context_lse, - suffix_output=suffix_output[..., : self.v_head_dim], + suffix_output=out, suffix_lse=suffix_lse, ) elif not writes_out: From f7180bc7d20415d1d5f40ed69b0feddd4ce7d5dc Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 17 Aug 2026 01:29:41 -0500 Subject: [PATCH 33/52] [Bugfix][Mamba] Fix overlapping state copy race (#50729) Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex Signed-off-by: Alessandra005 --- tests/v1/worker/test_mamba_utils.py | 189 ++++++++++++++++++++++------ vllm/v1/worker/mamba_utils.py | 102 +++++++++++---- 2 files changed, 228 insertions(+), 63 deletions(-) diff --git a/tests/v1/worker/test_mamba_utils.py b/tests/v1/worker/test_mamba_utils.py index a338b934b54f..2ba1ce1c5937 100644 --- a/tests/v1/worker/test_mamba_utils.py +++ b/tests/v1/worker/test_mamba_utils.py @@ -18,6 +18,7 @@ from vllm.v1.worker.mamba_utils import ( MambaCopyBuffers, MambaSpecDecodeGPUContext, + batch_memcpy, collect_mamba_copy_meta, do_mamba_copy_block, preprocess_mamba, @@ -536,6 +537,34 @@ def device(self): def test_config(self): return _TestConfig() + def test_batch_memcpy_left_overlap_has_memmove_semantics(self, device): + batch = 128 + row_bytes = 32 * 1024 + shift = 16 + copy_size = row_bytes - shift + + pattern = (torch.arange(row_bytes, dtype=torch.int32, device=device) % 251).to( + torch.uint8 + ) + state = pattern.expand(batch, -1).clone() + snapshot = state.clone() + + row_stride_bytes = state.stride(0) * state.element_size() + row_offsets = ( + torch.arange(batch, dtype=torch.int64, device=device) * row_stride_bytes + ) + dst_ptrs = (row_offsets + state.data_ptr()).to(torch.uint64) + src_ptrs = (row_offsets + state.data_ptr() + shift).to(torch.uint64) + sizes = torch.full((batch,), copy_size, dtype=torch.int32, device=device) + + expected = snapshot.clone() + expected[:, :copy_size].copy_(snapshot[:, shift:]) + for _ in range(10): + state.copy_(snapshot) + batch_memcpy(src_ptrs, dst_ptrs, sizes) + torch.accelerator.synchronize() + torch.testing.assert_close(state, expected, rtol=0, atol=0) + def test_matches_python_postprocess_mamba(self, device, test_config): """ Golden test: GPU kernel produces identical results to Python impl. @@ -1190,12 +1219,27 @@ def test_same_block_idx_with_offset_copies_then_sets_accepted_to_1( # --- Verify Python behavior (ground truth) --- dest_block_id = block_ids_per_req[0][1] # dest_block_idx = 1 - # Conv state should be modified (shifted copy within block) - conv_changed = not torch.allclose( - conv_state_py[dest_block_id], conv_state_orig[dest_block_id] + # This is an overlapping in-place left shift, so comparing only the + # Python and fused paths can hide the same memcpy race in both. Build + # the memmove result from the untouched snapshot and check each path + # independently. + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-1].copy_( + conv_state_orig[dest_block_id, 1:] ) - assert conv_changed, ( - "Python: Conv state should be modified when accept_token_bias > 0" + torch.testing.assert_close( + conv_state_py, + expected_conv_state, + rtol=0, + atol=0, + msg="Python: overlapping conv copy should have memmove semantics", + ) + torch.testing.assert_close( + conv_state_gpu, + expected_conv_state, + rtol=0, + atol=0, + msg="GPU: overlapping conv copy should have memmove semantics", ) # Temporal state should be modified (copy from different block) @@ -2122,13 +2166,10 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): actual_src_block_idx = src_block_idx + accept_token_bias actual_src_block_id = block_table[req, actual_src_block_idx] - All prior regression tests exercise only ``bias == 1``, i.e. they - only ever read one slot ahead of ``src_block_idx`` in the block - table. An off-by-one (or missing scale) in the address computation - on line 143 of ``mamba_utils.py`` would be invisible to every - existing test but would silently read the wrong physical block on - any speculative-decode cycle that accepts multiple tokens across a - block boundary, feeding a stale hidden state forward one step. + A ``bias == 1`` case only reads one slot ahead of ``src_block_idx`` + in the block table. This test isolates the larger-stride case, where + an off-by-one would read the wrong physical block after multiple + tokens are accepted across a block boundary. Setup (block_size=16): - running = 28 + 2 - 0 = 30 @@ -2141,8 +2182,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): With identity block_ids = [0,1,2,3,...], an off-by-one that used bias=1 would copy from block_ids[2]=2 instead of block_ids[3]=3, - producing a clear state-value mismatch against the Python - reference. + producing a clear mismatch against the untouched snapshot. """ cfg = test_config torch.manual_seed(7002) @@ -2166,6 +2206,7 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): fwd_py, fwd_gpu, ) = _make_dual_layer_state(cfg, device) + conv_state_orig = conv_state_py.clone() temporal_state_orig = temporal_state_py.clone() # --- Python reference --- @@ -2212,12 +2253,22 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): device=device, ) - # --- Ground truth: Python must have sourced temporal from block 3 --- + # --- Ground truth from untouched snapshots --- actual_src_block_id = block_ids_per_req[0][3] # == 3 dest_block_id = block_ids_per_req[0][1] # == 1 + expected_conv_state = conv_state_orig.clone() + expected_conv_state[dest_block_id, :-2].copy_( + conv_state_orig[dest_block_id, 2:] + ) + torch.testing.assert_close(conv_state_py, expected_conv_state, rtol=0, atol=0) + torch.testing.assert_close(conv_state_gpu, expected_conv_state, rtol=0, atol=0) + + # Python must have sourced temporal from block 3. torch.testing.assert_close( temporal_state_py[dest_block_id], temporal_state_orig[actual_src_block_id], + rtol=0, + atol=0, msg=( "Python reference did not copy from block_ids[src+bias]=3; " "test preconditions are wrong" @@ -2251,21 +2302,44 @@ def test_temporal_copy_with_bias_ge_2(self, device, test_config): msg="num_accepted_tokens mismatch at accept_token_bias=2", ) - def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( - self, device, test_config, monkeypatch + @pytest.mark.parametrize( + "same_physical_block", [True, False], ids=["same", "distinct"] + ) + @pytest.mark.parametrize("accept_token_bias", [1, 2, 3]) + @pytest.mark.parametrize( + "dtype", + [torch.float16, torch.float32, torch.float64], + ids=["fp16", "fp32", "fp64"], + ) + def test_sd_and_ds_conv_layouts_match_snapshot( + self, + device, + test_config, + monkeypatch, + accept_token_bias, + same_physical_block, + dtype, ): - """DS conv postprocess should match SD when accept_token_bias > 0.""" + """SD and DS copies should independently match memmove semantics.""" from vllm.model_executor.layers.mamba import mamba_utils as model_mamba_utils cfg = test_config + cfg.dtype = dtype torch.manual_seed(38898) req_ids = ["req_0"] - num_computed_tokens = [30] - num_scheduled_tokens = {"req_0": 1} + # Keep new_num_computed on an aligned boundary while varying how far + # below it the running state starts. This makes the copy bias exactly + # ``accept_token_bias`` for each case. The 32 boundary keeps source and + # destination in logical block 1; the 64 boundary copies block 2 -> 3. + aligned_boundary = 32 if same_physical_block else 64 + num_computed_tokens = [aligned_boundary - 2 * accept_token_bias] + num_scheduled_tokens = {"req_0": accept_token_bias} num_draft_tokens: dict[str, int] = {} - num_accepted_tokens = [2] # Results in accept_token_bias = 1 - mamba_state_idx = [1] # src_block_idx = 1 = dest_block_idx + num_accepted_tokens = [accept_token_bias + 1] + dest_block_idx = aligned_boundary // cfg.block_size - 1 + src_block_idx = dest_block_idx if same_physical_block else dest_block_idx - 1 + mamba_state_idx = [src_block_idx] block_ids_per_req = [list(range(8))] layer_names = ["layer_0"] @@ -2288,7 +2362,8 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( cfg.num_blocks, cfg.temporal_state_dim, dtype=cfg.dtype, device=device ) - # SD GPU path. Default layout is SD. + # SD GPU path. + monkeypatch.delenv("VLLM_SSM_CONV_STATE_LAYOUT", raising=False) model_mamba_utils.get_conv_state_layout.cache_clear() sd_conv = sd_source_conv.clone() sd_temporal = sd_source_temporal.clone() @@ -2312,9 +2387,32 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( ) torch.accelerator.synchronize() - # Sanity: SD path actually modified the state (copy was performed). - assert not torch.equal(sd_conv, sd_source_conv), ( - "SD baseline did not modify conv state; test setup is wrong" + src_block_id = block_ids_per_req[0][src_block_idx] + dest_block_id = block_ids_per_req[0][dest_block_idx] + expected_conv = sd_source_conv.clone() + expected_conv[dest_block_id, :-accept_token_bias].copy_( + sd_source_conv[src_block_id, accept_token_bias:] + ) + torch.testing.assert_close( + sd_conv, + expected_conv, + rtol=0, + atol=0, + msg="SD conv copy did not match the untouched source snapshot", + ) + + actual_temporal_src_idx = src_block_idx + accept_token_bias + actual_temporal_src_id = block_ids_per_req[0][actual_temporal_src_idx] + expected_temporal = sd_source_temporal.clone() + expected_temporal[dest_block_id].copy_( + sd_source_temporal[actual_temporal_src_id] + ) + torch.testing.assert_close( + sd_temporal, + expected_temporal, + rtol=0, + atol=0, + msg="SD temporal copy did not match the untouched source snapshot", ) # DS GPU path on the DS twin. @@ -2346,22 +2444,39 @@ def test_ds_conv_layout_bias_gt_0_byte_equal_to_sd( # Reset the lru cache so other tests see the default layout again. model_mamba_utils.get_conv_state_layout.cache_clear() - # DS bytes, un-permuted, should match the SD result. + # Validate DS independently against the snapshot; otherwise a shared + # SD/DS bug would remain invisible. + ds_conv_sd_layout = ds_conv.permute(0, 2, 1).contiguous() torch.testing.assert_close( - ds_conv.permute(0, 2, 1).contiguous(), - sd_conv, - msg=( - "DS conv post-kernel does not match SD baseline; the DS " - "row-loop in postprocess_mamba_fused_kernel is wrong." - ), + ds_conv_sd_layout, + expected_conv, + rtol=0, + atol=0, + msg="DS conv copy did not match the untouched source snapshot", ) torch.testing.assert_close( ds_temporal, - sd_temporal, - msg="DS temporal state diverged from SD", + expected_temporal, + rtol=0, + atol=0, + msg="DS temporal copy did not match the untouched source snapshot", + ) + + expected_accepted = 1 if same_physical_block else accept_token_bias + 1 + expected_accepted_tensor = torch.tensor( + [expected_accepted], dtype=torch.int32, device=device ) torch.testing.assert_close( - gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], gpu_ctx_sd.num_accepted_tokens_out[:num_reqs], - msg="DS num_accepted_tokens diverged from SD", + expected_accepted_tensor, + rtol=0, + atol=0, + msg="SD num_accepted_tokens result is wrong", + ) + torch.testing.assert_close( + gpu_ctx_ds.num_accepted_tokens_out[:num_reqs], + expected_accepted_tensor, + rtol=0, + atol=0, + msg="DS num_accepted_tokens result is wrong", ) diff --git a/vllm/v1/worker/mamba_utils.py b/vllm/v1/worker/mamba_utils.py index f87a34a364f9..4add17c7e851 100644 --- a/vllm/v1/worker/mamba_utils.py +++ b/vllm/v1/worker/mamba_utils.py @@ -188,19 +188,46 @@ def _copy_mamba_state_block( src_block_id = tl.load(block_table_base + src_col).to(tl.int64) dim_rows = tl.load(state_dim_row_count_ptr + state_idx) row_stride = tl.load(state_dim_row_stride_ptr + state_idx) - per_row_bytes = (conv_width - token_bias).to(tl.int64) * state_elem_size - bias_bytes = token_bias.to(tl.int64) * state_elem_size src_block_addr = state_base_addr + src_block_id * state_block_stride offsets = tl.arange(0, COPY_BLOCK_SIZE) - for d in range(0, dim_rows): - row_src = src_block_addr + d * row_stride + bias_bytes - row_dst = dst_addr + d * row_stride - for i in range(0, per_row_bytes, COPY_BLOCK_SIZE): - mask = (i + offsets) < per_row_bytes - curr_src = (row_src + i + offsets).to(tl.pointer_type(tl.uint8)) - curr_dst = (row_dst + i + offsets).to(tl.pointer_type(tl.uint8)) - data = tl.load(curr_src, mask=mask) - tl.store(curr_dst, data, mask=mask) + + # Stable row-to-lane ownership makes left shifts memmove-safe while + # exposing the dimension rows in parallel. All addresses retain + # state_elem_size alignment: tensor strides and token offsets are + # measured in whole elements before conversion to bytes. + num_dst_tokens = conv_width - token_bias + for token_idx in range(0, num_dst_tokens): + for row_base in range(0, dim_rows, COPY_BLOCK_SIZE): + rows = row_base + offsets + mask = rows < dim_rows + src_byte_addr = ( + src_block_addr + + rows * row_stride + + (token_idx + token_bias) * state_elem_size + ) + dst_byte_addr = ( + dst_addr + rows * row_stride + token_idx * state_elem_size + ) + if state_elem_size == 2: + src_u16 = src_byte_addr.to(tl.pointer_type(tl.uint16)) + dst_u16 = dst_byte_addr.to(tl.pointer_type(tl.uint16)) + data_u16 = tl.load(src_u16, mask=mask) + tl.store(dst_u16, data_u16, mask=mask) + elif state_elem_size == 4: + src_u32 = src_byte_addr.to(tl.pointer_type(tl.uint32)) + dst_u32 = dst_byte_addr.to(tl.pointer_type(tl.uint32)) + data_u32 = tl.load(src_u32, mask=mask) + tl.store(dst_u32, data_u32, mask=mask) + else: + for byte_idx in range(0, state_elem_size): + src_u8 = (src_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + dst_u8 = (dst_byte_addr + byte_idx).to( + tl.pointer_type(tl.uint8) + ) + data_u8 = tl.load(src_u8, mask=mask) + tl.store(dst_u8, data_u8, mask=mask) return if is_conv_state: @@ -209,22 +236,40 @@ def _copy_mamba_state_block( # SD conv: copy # state[bt[src_col], token_bias:] -> # state[bt[dst_col], :conv_width - token_bias] - # Small per-block bytes (~60-80 KiB) make tiling degenerate, so - # conv runs as a single-CTA memcpy (NUM_TILES=1). src_block_id = tl.load(block_table_base + src_col).to(tl.int64) - src_offset = token_bias.to(tl.int64) * state_inner_size * state_elem_size - src_addr = state_base_addr + src_block_id * state_block_stride + src_offset - copy_size = ( - (conv_width - token_bias).to(tl.int64) * state_inner_size * state_elem_size - ) - _memcpy_u64_tiled( - src_addr, - dst_addr, - copy_size, - tile_idx, - COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, - NUM_TILES=1, - ) + src_block_addr = state_base_addr + src_block_id * state_block_stride + token_bytes = state_inner_size * state_elem_size + num_dst_tokens = conv_width - token_bias + + # Distinct blocks and exact self-copies cannot have a destructive + # overlap, so retain the u64-vectorized single-CTA copy. + if src_block_id != dest_block_id or token_bias == 0: + src_addr = src_block_addr + token_bias.to(tl.int64) * token_bytes + copy_size = num_dst_tokens.to(tl.int64) * token_bytes + _memcpy_u64_tiled( + src_addr, + dst_addr, + copy_size, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) + return + + # Copy tokens from low to high. Each token-sized source and destination + # region is disjoint, so same-block left shifts are memmove-safe + # without a barrier. + for token_idx in range(0, num_dst_tokens): + src_token = src_block_addr + (token_idx + token_bias) * token_bytes + dst_token = dst_addr + token_idx * token_bytes + _memcpy_u64_tiled( + src_token, + dst_token, + token_bytes, + tile_idx, + COPY_BLOCK_SIZE=COPY_BLOCK_SIZE, + NUM_TILES=1, + ) return # Temporal state: copy state[bt[src_col + token_bias]] -> state[bt[dst_col]] @@ -521,6 +566,7 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): src_ptr = tl.load(src_ptrs + pid) dst_ptr = tl.load(dst_ptrs + pid) size = tl.load(sizes + pid) + is_left_overlap = dst_ptr < src_ptr and dst_ptr + size > src_ptr offsets = tl.arange(0, BLOCK_SIZE) for i in range(0, size, BLOCK_SIZE): @@ -530,6 +576,10 @@ def batch_memcpy_kernel(src_ptrs, dst_ptrs, sizes, BLOCK_SIZE: tl.constexpr): curr_dst_ptr = (dst_ptr + i + offsets).to(tl.pointer_type(tl.uint8)) data = tl.load(curr_src_ptr, mask=mask) + if is_left_overlap: + # Preserve each lane's source before a lower-address lane stores + # over it. The condition is uniform within the program. + tl.debug_barrier() tl.store(curr_dst_ptr, data, mask=mask) From a653e74b10d53aaf64555ad2ed991b34081d9547 Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Sat, 22 Aug 2026 05:51:10 +0900 Subject: [PATCH 34/52] feat(spec_decode): run the Kimi-K3 DSpark draft on a dedicated remote GPU Adds a verifier-side proxy (RemoteK3DSparkSpeculator) and a standalone draft server (vllm.entrypoints.k3_dspark_standalone + k3_dspark_rpc) so the DSpark draft model executes on its own single GPU while the target runs TP/DCP on separate GPUs. Draft weights, KV, Markov head, and CUDA graphs live entirely on the draft process; the target exchanges context and proposals over a versioned ZMQ/TCP protocol (PROTOCOL_VERSION=2). Behavior and invariants: - VLLM_K3_DRAFT_REMOTE_ADDRESS selects the remote path at speculator construction; unset preserves the existing local DSpark/DFlash path. - propose() matches BaseSpeculator's signature; rank 0 performs RPC and all ranks consume the broadcast result. - Fail closed: any RPC failure fills draft tokens with -1 (no speculation for the step) and disables affected requests until they leave the batch; FREE remains safe for never-created remote state. - Retained-prefix reconnection validates a target prefix-cache hit against retained draft state via a host-visible view of the request token table (InputBatch.all_token_ids_cpu, backed by StagedWriteTensor.cpu). - CUDA-graph capture interface preserved: init_cudagraph_manager and capture(capture_phase=...) conform to BaseSpeculator. Compatibility: no change when the remote address is unset; draft side supports DSpark and DFlash checkpoints on a single GPU including Ampere-class cards. Validation: 19 new CPU unit tests pass (test_k3_dspark_remote_speculator.py, test_k3_dspark_standalone.py); production-qualified serving lukealonso/Kimi-K3-QSRT-K2 TP8/DCP8 with an Inferact BF16 DSpark draft on a dedicated RTX 3090. Limitations: one remote draft process (draft TP1); TCP transport; greedy draft sampling with block rejection sampling on the verifier. AI assistance was used in the preparation of this change; every line was reviewed and the listed tests were run by the submitter. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com> --- .../test_k3_dspark_remote_speculator.py | 126 ++ .../spec_decode/test_k3_dspark_standalone.py | 172 +++ vllm/entrypoints/k3_dspark_rpc.py | 1363 +++++++++++++++++ vllm/entrypoints/k3_dspark_standalone.py | 847 ++++++++++ vllm/v1/worker/gpu/buffer_utils.py | 6 + vllm/v1/worker/gpu/input_batch.py | 5 + vllm/v1/worker/gpu/model_runner.py | 1 + vllm/v1/worker/gpu/spec_decode/__init__.py | 26 + .../spec_decode/dspark/remote_speculator.py | 770 ++++++++++ 9 files changed, 3316 insertions(+) create mode 100644 tests/v1/spec_decode/test_k3_dspark_remote_speculator.py create mode 100644 tests/v1/spec_decode/test_k3_dspark_standalone.py create mode 100644 vllm/entrypoints/k3_dspark_rpc.py create mode 100644 vllm/entrypoints/k3_dspark_standalone.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py diff --git a/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py b/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py new file mode 100644 index 000000000000..dcfe1da8fd1c --- /dev/null +++ b/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import numpy as np +import pytest +import torch + +from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import ( + RemoteK3DSparkSpeculator, + _anchor_positions_from_context, + _build_valid_context_plan, + _contiguous_draft_output, + _RetainedRequestPrefix, +) + + +def test_build_valid_context_plan_drops_rejected_tail_rows(): + batch = SimpleNamespace( + num_reqs=2, + num_scheduled_tokens=np.array([4, 3], dtype=np.int32), + num_computed_tokens_np=np.array([10, 20], dtype=np.int32), + ) + + indices, counts = _build_valid_context_plan(batch, [2, 0]) + + assert indices == [0, 1, 4, 5, 6] + assert counts == [2, 3] + + +def test_anchor_positions_follow_actual_valid_context_rows(): + positions = torch.tensor([24, 25, 26, 80, 81], dtype=torch.int64) + + anchors = _anchor_positions_from_context([3, 2], positions) + + assert anchors == [27, 82] + + +def test_remote_tokens_copy_supports_adaptive_depth(): + proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator) + proxy.device = torch.device("cpu") + proxy.draft_tokens = torch.full((3, 8), -1, dtype=torch.int64) + + proxy._copy_tokens_from_response( + {"tokens": [[11, 12], [21, 22]]}, + active_indices=[0, 2], + num_speculative_tokens=2, + ) + + assert proxy.draft_tokens.tolist() == [ + [11, 12, -1, -1, -1, -1, -1, -1], + [-1, -1, -1, -1, -1, -1, -1, -1], + [21, 22, -1, -1, -1, -1, -1, -1], + ] + + +def test_adaptive_depth_output_is_contiguous_for_tp_broadcast(): + draft_tokens = torch.arange(24, dtype=torch.int64).view(3, 8) + + output = _contiguous_draft_output(draft_tokens, 2, 3) + + assert output.is_contiguous() + assert output.tolist() == [[0, 1, 2], [8, 9, 10]] + + +@pytest.mark.parametrize("rejected", [[5, 0], [-1, 0]]) +def test_build_valid_context_plan_rejects_invalid_counts(rejected): + batch = SimpleNamespace( + num_reqs=2, + num_scheduled_tokens=np.array([4, 3], dtype=np.int32), + num_computed_tokens_np=np.array([0, 0], dtype=np.int32), + ) + + with pytest.raises(ValueError, match="Invalid valid-context length"): + _build_valid_context_plan(batch, rejected) + + +def _make_prefix_matcher() -> RemoteK3DSparkSpeculator: + proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator) + proxy._known_requests = {"old"} + proxy._remote_block_size = 16 + proxy._remote_window_size = 32 + proxy._remote_prefix_cache_tokens = 128 + proxy._retained_prefixes = { + "old": _RetainedRequestPrefix( + token_ids=torch.arange(96, dtype=torch.int32), + committed_end=96, + context_start=0, + serial=1, + ) + } + return proxy + + +def test_remote_prefix_match_requires_exact_token_identity(): + proxy = _make_prefix_matcher() + matching = torch.arange(80, dtype=torch.int32) + + assert proxy._find_reconnect_source(matching, 80, {"new"}) == "old" + + mismatched = matching.clone() + mismatched[40] = -1 + assert proxy._find_reconnect_source(mismatched, 80, {"new"}) is None + + +def test_remote_prefix_match_rejects_range_evicted_from_projected_cache(): + proxy = _make_prefix_matcher() + proxy._remote_prefix_cache_tokens = 48 + matching = torch.arange(40, dtype=torch.int32) + + assert proxy._find_reconnect_source(matching, 40, {"new"}) is None + + +def test_remote_prefix_match_rejects_history_before_cold_bootstrap(): + proxy = _make_prefix_matcher() + proxy._retained_prefixes["old"].context_start = 64 + + assert ( + proxy._find_reconnect_source(torch.arange(80, dtype=torch.int32), 80, {"new"}) + is None + ) + assert ( + proxy._find_reconnect_source(torch.arange(96, dtype=torch.int32), 96, {"new"}) + == "old" + ) diff --git a/tests/v1/spec_decode/test_k3_dspark_standalone.py b/tests/v1/spec_decode/test_k3_dspark_standalone.py new file mode 100644 index 000000000000..39e6f6c7d95c --- /dev/null +++ b/tests/v1/spec_decode/test_k3_dspark_standalone.py @@ -0,0 +1,172 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import json + +import pytest +import torch + +from vllm.entrypoints.k3_dspark_rpc import ( + DraftKVSlotAllocator, + ProjectedContextCache, +) +from vllm.entrypoints.k3_dspark_standalone import ( + EMBED_TENSOR, + LM_HEAD_TENSOR, + resolve_shared_weight_files, +) + + +def _write_index(root, weight_map): + (root / "model.safetensors.index.json").write_text( + json.dumps({"weight_map": weight_map}) + ) + + +def test_resolve_shared_weight_files_requires_both_target_tensors(tmp_path): + shard = tmp_path / "shared.safetensors" + shard.touch() + _write_index(tmp_path, {EMBED_TENSOR: shard.name}) + + with pytest.raises(KeyError, match=LM_HEAD_TENSOR): + resolve_shared_weight_files(tmp_path) + + +def test_resolve_shared_weight_files_resolves_checkpoint_shards(tmp_path): + shard = tmp_path / "shared.safetensors" + shard.touch() + _write_index( + tmp_path, + { + EMBED_TENSOR: shard.name, + LM_HEAD_TENSOR: shard.name, + }, + ) + + resolved = resolve_shared_weight_files(tmp_path) + + assert resolved == { + EMBED_TENSOR: shard.resolve(), + LM_HEAD_TENSOR: shard.resolve(), + } + + +def test_resolve_shared_weight_files_rejects_path_escape(tmp_path): + outside = tmp_path.parent / "outside.safetensors" + outside.touch() + _write_index( + tmp_path, + { + EMBED_TENSOR: f"../{outside.name}", + LM_HEAD_TENSOR: f"../{outside.name}", + }, + ) + + with pytest.raises(ValueError, match="escapes"): + resolve_shared_weight_files(tmp_path) + + +def test_draft_kv_slot_allocator_keeps_rolling_blocks_unique(): + allocator = DraftKVSlotAllocator( + num_cache_blocks=11, + block_size=4, + window_size=16, + max_requests=2, + ) + first, created = allocator.get_or_allocate("first") + second, _ = allocator.get_or_allocate("second") + + assert created + assert allocator.physical_block_range(first) == slice(1, 6) + assert allocator.physical_block_range(second) == slice(6, 11) + blocks, local_len = allocator.block_table(first, 21) + assert blocks == [2, 3, 4, 5, 1] + assert len(blocks) == len(set(blocks)) + assert local_len == 17 + + +def test_draft_kv_slot_allocator_reuses_freed_request_slot(): + allocator = DraftKVSlotAllocator( + num_cache_blocks=6, + block_size=4, + window_size=16, + max_requests=1, + ) + original, _ = allocator.get_or_allocate("original") + with pytest.raises(RuntimeError, match="capacity exhausted"): + allocator.get_or_allocate("other") + + assert allocator.free("original") is original + replacement, created = allocator.get_or_allocate("replacement") + assert created + assert replacement.slot == original.slot + + +def test_draft_kv_slot_allocator_rebinds_without_changing_slot(): + allocator = DraftKVSlotAllocator( + num_cache_blocks=6, + block_size=4, + window_size=16, + max_requests=1, + ) + original, _ = allocator.get_or_allocate("original") + + rebound = allocator.rebind("original", "replacement") + + assert rebound is original + assert rebound.request_id == "replacement" + assert allocator.get("original") is None + assert allocator.get("replacement") is rebound + + +def test_projected_context_cache_rewinds_and_overwrites_exact_prefix(): + cache = ProjectedContextCache(hidden_size=2, max_tokens=8, chunk_size=4) + initial = torch.arange(16, dtype=torch.bfloat16).view(8, 2) + cache.append(0, initial) + + replacement = torch.tensor([[100, 101], [102, 103]], dtype=torch.bfloat16) + cache.append(5, replacement) + + assert cache.start_position == 0 + assert cache.end_position == 7 + assert torch.equal(cache.read(0, 5), initial[:5]) + assert torch.equal(cache.read(5, 7), replacement) + + +def test_projected_context_cache_evicts_old_rows_without_regaining_them(): + cache = ProjectedContextCache(hidden_size=1, max_tokens=6, chunk_size=4) + cache.append(0, torch.arange(8, dtype=torch.bfloat16).view(8, 1)) + + assert cache.start_position == 2 + assert cache.has_range(2, 8) + assert not cache.has_range(1, 8) + + cache.append(5, torch.tensor([[50], [60]], dtype=torch.bfloat16)) + assert cache.start_position == 2 + assert cache.end_position == 7 + with pytest.raises(ValueError, match="unavailable"): + cache.read(1, 7) + + +def test_projected_context_cache_tracks_configured_device(): + cache = ProjectedContextCache( + hidden_size=2, + max_tokens=4, + chunk_size=2, + device=torch.device("cpu"), + ) + states = torch.arange(8, dtype=torch.bfloat16).view(4, 2) + + cache.append(0, states) + + assert cache.device == torch.device("cpu") + assert cache.read(0, 4).device == cache.device + assert cache.allocated_bytes == states.numel() * states.element_size() + + +def test_projected_context_cache_rejects_device_mismatch(): + cache = ProjectedContextCache(hidden_size=2, max_tokens=4) + states = torch.empty((1, 2), dtype=torch.bfloat16, device="meta") + + with pytest.raises(ValueError, match="device mismatch"): + cache.append(0, states) diff --git a/vllm/entrypoints/k3_dspark_rpc.py b/vllm/entrypoints/k3_dspark_rpc.py new file mode 100644 index 000000000000..96f049205573 --- /dev/null +++ b/vllm/entrypoints/k3_dspark_rpc.py @@ -0,0 +1,1363 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Host-staged RPC for a dedicated Kimi-K3 draft GPU. + +The verifier and RTX 3090 do not have CUDA peer access on the target host, so +the first transport deliberately uses ZMQ multipart frames backed by host +memory. The protocol is small and versioned so a verifier-side proxy can be +added without coupling the standalone process to the generic EAGLE draft +server protocol. +""" + +from __future__ import annotations + +import json +import os +import threading +import time +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import torch +import zmq + +from vllm.config.vllm import set_current_vllm_config +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonDecodeMetadata, + MLACommonMetadata, +) +from vllm.v1.attention.backend import CommonAttentionMetadata +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) +from vllm.v1.worker.gpu.spec_decode.utils import get_parallel_drafting_token_id + +if TYPE_CHECKING: + from vllm.entrypoints.k3_dspark_standalone import StandaloneRuntime + +logger = init_logger(__name__) + +PROTOCOL_VERSION = 2 + + +class ProjectedContextCache: + """Bounded, chunked cache of projected DSpark context states. + + The standalone draft server keeps this cache on the draft device. Keeping + the projected rows on CPU forced a blocking D2H copy after every proposal, + even though the rows are normally consumed again by the same GPU during a + prefix reconnect. CPU remains the default for lightweight unit tests and + callers which explicitly want host storage. + """ + + def __init__( + self, + *, + hidden_size: int, + max_tokens: int, + chunk_size: int = 256, + initial_position: int = 0, + device: torch.device | str = "cpu", + ) -> None: + if hidden_size <= 0 or max_tokens <= 0 or chunk_size <= 0: + raise ValueError("Projected context cache dimensions must be positive") + if initial_position < 0: + raise ValueError("Projected context cache position cannot be negative") + self.hidden_size = hidden_size + self.max_tokens = max_tokens + self.chunk_size = chunk_size + self.device = torch.device(device) + self.start_position = initial_position + self.end_position = initial_position + self._chunks: dict[int, torch.Tensor] = {} + + def _truncate(self, end_position: int) -> None: + if not self.start_position <= end_position <= self.end_position: + raise ValueError( + "Cannot truncate projected context outside its retained range: " + f"retained=[{self.start_position}, {self.end_position}), " + f"requested_end={end_position}" + ) + first_discarded_chunk = (end_position + self.chunk_size - 1) // self.chunk_size + for chunk_idx in list(self._chunks): + if chunk_idx >= first_discarded_chunk: + del self._chunks[chunk_idx] + self.end_position = end_position + + def append(self, first_position: int, states: torch.Tensor) -> None: + if states.device != self.device: + raise ValueError( + "Projected context cache device mismatch: " + f"cache={self.device}, states={states.device}" + ) + if states.dtype != torch.bfloat16 or states.ndim != 2: + raise ValueError("Projected context states must be a 2D BF16 tensor") + if states.shape[1] != self.hidden_size: + raise ValueError( + f"Projected context width is {states.shape[1]}, expected " + f"{self.hidden_size}" + ) + if first_position < self.start_position or first_position > self.end_position: + raise ValueError( + "Projected context append is not contiguous with retained state: " + f"retained=[{self.start_position}, {self.end_position}), " + f"first={first_position}" + ) + if first_position < self.end_position: + self._truncate(first_position) + + offset = 0 + num_rows = int(states.shape[0]) + while offset < num_rows: + position = first_position + offset + chunk_idx, chunk_offset = divmod(position, self.chunk_size) + count = min(num_rows - offset, self.chunk_size - chunk_offset) + chunk = self._chunks.get(chunk_idx) + if chunk is None: + chunk = torch.empty( + (self.chunk_size, self.hidden_size), + dtype=torch.bfloat16, + device=self.device, + ) + self._chunks[chunk_idx] = chunk + chunk[chunk_offset : chunk_offset + count].copy_( + states[offset : offset + count] + ) + offset += count + + self.end_position = first_position + num_rows + self.start_position = max( + self.start_position, + self.end_position - self.max_tokens, + ) + for chunk_idx in list(self._chunks): + if (chunk_idx + 1) * self.chunk_size <= self.start_position: + del self._chunks[chunk_idx] + + def has_range(self, start_position: int, end_position: int) -> bool: + return ( + self.start_position <= start_position <= end_position <= self.end_position + ) + + def read(self, start_position: int, end_position: int) -> torch.Tensor: + if not self.has_range(start_position, end_position): + raise ValueError( + "Projected context range is unavailable: " + f"retained=[{self.start_position}, {self.end_position}), " + f"requested=[{start_position}, {end_position})" + ) + output = torch.empty( + (end_position - start_position, self.hidden_size), + dtype=torch.bfloat16, + device=self.device, + ) + offset = 0 + while start_position + offset < end_position: + position = start_position + offset + chunk_idx, chunk_offset = divmod(position, self.chunk_size) + count = min( + end_position - position, + self.chunk_size - chunk_offset, + ) + chunk = self._chunks.get(chunk_idx) + if chunk is None: + raise RuntimeError( + f"Projected context chunk {chunk_idx} is unexpectedly missing" + ) + output[offset : offset + count].copy_( + chunk[chunk_offset : chunk_offset + count] + ) + offset += count + return output + + def truncate(self, end_position: int) -> None: + if end_position < self.end_position: + self._truncate(end_position) + + @property + def allocated_bytes(self) -> int: + return sum( + chunk.numel() * chunk.element_size() for chunk in self._chunks.values() + ) + + +@dataclass +class DraftRequestState: + request_id: str + slot: int + committed_end: int = 0 + context_start: int = 0 + context_cache: ProjectedContextCache | None = None + + +class DraftKVSlotAllocator: + """Assign fixed rolling MLA block ranges to a small request batch.""" + + def __init__( + self, + *, + num_cache_blocks: int, + block_size: int, + window_size: int, + max_requests: int, + ) -> None: + if window_size <= 0 or window_size % block_size != 0: + raise ValueError( + "DSpark KV window must be a positive block-size multiple, got " + f"window={window_size}, block_size={block_size}" + ) + self.block_size = block_size + self.window_size = window_size + self.max_requests = max_requests + # The vLLM rolling window may retain window + block_size - 1 tokens + # while it waits for the next whole-block shift. + self.blocks_per_request = window_size // block_size + 1 + required = 1 + max_requests * self.blocks_per_request + if required > num_cache_blocks: + raise ValueError( + "Dedicated draft KV cache is too small for the requested rolling " + f"slots: required_blocks={required}, available={num_cache_blocks}" + ) + self._free_slots = list(range(max_requests)) + self._states: dict[str, DraftRequestState] = {} + + def get_or_allocate(self, request_id: str) -> tuple[DraftRequestState, bool]: + state = self._states.get(request_id) + if state is not None: + return state, False + if not self._free_slots: + raise RuntimeError( + f"DSpark request capacity exhausted (max={self.max_requests})" + ) + slot = self._free_slots.pop(0) + state = DraftRequestState(request_id=request_id, slot=slot) + self._states[request_id] = state + return state, True + + def free(self, request_id: str) -> DraftRequestState | None: + state = self._states.pop(request_id, None) + if state is not None: + self._free_slots.append(state.slot) + self._free_slots.sort() + return state + + def get(self, request_id: str) -> DraftRequestState | None: + return self._states.get(request_id) + + def rebind(self, source_request_id: str, request_id: str) -> DraftRequestState: + state = self._states.get(source_request_id) + if state is None: + raise KeyError(f"Unknown DSpark source request {source_request_id!r}") + if source_request_id == request_id: + return state + if request_id in self._states: + raise ValueError(f"DSpark request {request_id!r} already exists") + del self._states[source_request_id] + state.request_id = request_id + self._states[request_id] = state + return state + + def physical_block(self, state: DraftRequestState, position: int) -> int: + absolute_block = position // self.block_size + return ( + 1 + + state.slot * self.blocks_per_request + + absolute_block % self.blocks_per_request + ) + + def cache_slot(self, state: DraftRequestState, position: int) -> int: + return self.physical_block(state, position) * self.block_size + ( + position % self.block_size + ) + + def block_table( + self, state: DraftRequestState, sequence_end: int + ) -> tuple[list[int], int]: + if sequence_end <= 0: + raise ValueError(f"sequence_end must be positive, got {sequence_end}") + first_block = ( + max(state.context_start, sequence_end - self.window_size) // self.block_size + ) + end_block = (sequence_end + self.block_size - 1) // self.block_size + blocks = [ + 1 + + state.slot * self.blocks_per_request + + absolute_block % self.blocks_per_request + for absolute_block in range(first_block, end_block) + ] + local_sequence_len = sequence_end - first_block * self.block_size + if local_sequence_len > self.window_size + self.block_size - 1: + raise AssertionError("rolling DSpark sequence length exceeded its window") + if len(blocks) > self.blocks_per_request: + raise AssertionError("rolling DSpark block table aliases a live block") + return blocks, local_sequence_len + + def physical_block_range(self, state: DraftRequestState) -> slice: + start = 1 + state.slot * self.blocks_per_request + return slice(start, start + self.blocks_per_request) + + @property + def active_requests(self) -> int: + return len(self._states) + + @property + def request_ids(self) -> list[str]: + return list(self._states) + + +@dataclass +class _DraftCudaGraphState: + """Persistent inputs and output for one standalone draft graph shape.""" + + batch_size: int + num_speculative_tokens: int + query_len: int + input_ids: torch.Tensor + positions: torch.Tensor + slots: torch.Tensor + seq_lens: torch.Tensor + block_table: torch.Tensor + output_tokens: torch.Tensor + input_ids_host: torch.Tensor + positions_host: torch.Tensor + slots_host: torch.Tensor + seq_lens_host: torch.Tensor + block_table_host: torch.Tensor + attn_metadata: dict[str, Any] + slot_mapping: dict[str, torch.Tensor] + graph: torch.cuda.CUDAGraph | None = None + captured_hidden: torch.Tensor | None = None + captured_logits: torch.Tensor | None = None + + def stage( + self, + *, + input_ids: list[int], + positions: list[int], + slots: list[int], + block_rows: list[list[int]], + seq_lens: list[int], + ) -> None: + """Copy one request batch into address-stable graph inputs.""" + expected_tokens = self.batch_size * self.query_len + if not ( + len(input_ids) == len(positions) == len(slots) == expected_tokens + and len(block_rows) == len(seq_lens) == self.batch_size + ): + raise ValueError("Draft CUDA graph input shape mismatch") + + self.input_ids_host.copy_(torch.tensor(input_ids, dtype=torch.int64)) + self.positions_host.copy_(torch.tensor(positions, dtype=torch.int64)) + self.slots_host.copy_(torch.tensor(slots, dtype=torch.int64)) + self.seq_lens_host.copy_(torch.tensor(seq_lens, dtype=torch.int32)) + self.block_table_host.zero_() + for row_idx, row in enumerate(block_rows): + if len(row) > self.block_table_host.shape[1]: + raise ValueError( + "Draft block table exceeds CUDA graph capacity: " + f"row={len(row)}, capacity={self.block_table_host.shape[1]}" + ) + self.block_table_host[row_idx, : len(row)].copy_( + torch.tensor(row, dtype=torch.int32) + ) + + # All copies and the replay are enqueued on the same stream. The + # proposal's final query event synchronizes before these pinned host + # buffers can be reused by the next (serialized) RPC. + self.input_ids.copy_(self.input_ids_host, non_blocking=True) + self.positions.copy_(self.positions_host, non_blocking=True) + self.slots.copy_(self.slots_host, non_blocking=True) + self.seq_lens.copy_(self.seq_lens_host, non_blocking=True) + self.block_table.copy_(self.block_table_host, non_blocking=True) + + +class K3DSparkDraftEngine: + """Minimal greedy K3 draft scheduler backed by the 3090 KV cache.""" + + def __init__( + self, + runtime: StandaloneRuntime, + *, + max_requests: int, + window_size: int, + device: torch.device, + ) -> None: + self.runtime = runtime + self.model = runtime.model + self.method = runtime.method + self.device = device + self.max_model_len = int(runtime.vllm_config.model_config.max_model_len) + first_cache = next(iter(runtime.kv_caches.values())) + self.allocator = DraftKVSlotAllocator( + num_cache_blocks=int(first_cache.shape[0]), + block_size=runtime.kv_cache_block_size, + window_size=window_size, + max_requests=max_requests, + ) + speculative_config = runtime.vllm_config.speculative_config + assert speculative_config is not None + draft_config = speculative_config.draft_model_config.hf_config + self.hidden_size = int(draft_config.hidden_size) + aux_layers = get_eagle3_aux_layers_from_config(speculative_config) + if not aux_layers: + raise ValueError( + f"K3 {self.method} config does not declare target auxiliary layers" + ) + self.num_aux_layers = len(aux_layers) + target_hidden_size = int( + getattr(draft_config, "target_hidden_size", None) + or draft_config.hidden_size + ) + self.raw_context_width = int(target_hidden_size * self.num_aux_layers) + self.mask_token_id = get_parallel_drafting_token_id(draft_config) + self.max_speculative_tokens = int( + runtime.vllm_config.speculative_config.num_speculative_tokens + ) + self.max_context_tokens = int( + runtime.vllm_config.scheduler_config.max_num_batched_tokens + ) + self.prefix_cache_tokens = int( + os.environ.get( + "VLLM_K3_DRAFT_PREFIX_CACHE_TOKENS", + os.environ.get("VLLM_K3_DSPARK_PREFIX_CACHE_TOKENS", "131072"), + ) + ) + if self.prefix_cache_tokens < self.allocator.window_size: + raise ValueError( + "VLLM_K3_DRAFT_PREFIX_CACHE_TOKENS must be at least the " + f"draft KV window ({self.allocator.window_size}), got " + f"{self.prefix_cache_tokens}" + ) + self._positions_staging = torch.empty( + self.max_context_tokens, + dtype=torch.int64, + pin_memory=True, + ) + self._context_staging = torch.empty( + self.max_context_tokens * self.raw_context_width, + dtype=torch.bfloat16, + pin_memory=True, + ) + self._lock = threading.Lock() + self.proposal_count = 0 + self.last_latency_ms = 0.0 + self.last_timing_ms: dict[str, float] = {} + self._timing_totals_ms: dict[str, float] = {} + self.cold_bootstrap_count = 0 + self.reconnect_count = 0 + self.last_reconnect_latency_ms = 0.0 + self.cuda_graph_enabled = False + self.cuda_graph_capture_seconds = 0.0 + self.cuda_graph_memory_gib = 0.0 + self.cuda_graph_replay_count = 0 + self.cuda_graph_eager_fallback_count = 0 + self._cuda_graphs: dict[tuple[int, int], _DraftCudaGraphState] = {} + + def _make_cuda_graph_state( + self, + batch_size: int, + num_speculative_tokens: int, + ) -> _DraftCudaGraphState: + if self.method == "dflash" and self.runtime.attn_metadata_builder is None: + raise RuntimeError("K3 DFlash attention metadata builder is missing") + + query_len = ( + num_speculative_tokens + if self.method == "dspark" + else 1 + num_speculative_tokens + ) + num_tokens = batch_size * query_len + max_blocks = self.allocator.blocks_per_request + input_ids = torch.empty(num_tokens, dtype=torch.int64, device=self.device) + positions = torch.empty(num_tokens, dtype=torch.int64, device=self.device) + slots = torch.empty(num_tokens, dtype=torch.int64, device=self.device) + seq_lens = torch.empty(batch_size, dtype=torch.int32, device=self.device) + block_table = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + output_tokens = torch.empty( + (batch_size, num_speculative_tokens), + dtype=torch.int64, + device=self.device, + ) + + input_ids_host = torch.empty(num_tokens, dtype=torch.int64, pin_memory=True) + positions_host = torch.empty(num_tokens, dtype=torch.int64, pin_memory=True) + slots_host = torch.empty(num_tokens, dtype=torch.int64, pin_memory=True) + seq_lens_host = torch.empty(batch_size, dtype=torch.int32, pin_memory=True) + block_table_host = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, pin_memory=True + ) + + # The graph's launch topology is shape-static, while seq_lens remains + # a live tensor. Triton BF16 attention uses seq_lens to bound the KV + # scan; this conservative upper bound therefore does not force every + # replay to scan the full rolling window. + max_seq_len = self.allocator.window_size + self.allocator.block_size - 1 + query_start_cpu = torch.arange( + 0, + (batch_size + 1) * query_len, + query_len, + dtype=torch.int32, + ) + query_start_gpu = query_start_cpu.to(self.device) + if self.method == "dflash": + common = CommonAttentionMetadata( + query_start_loc=query_start_gpu, + query_start_loc_cpu=query_start_cpu, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=torch.full( + (batch_size,), max_seq_len, dtype=torch.int32 + ), + max_seq_len=max_seq_len, + num_reqs=batch_size, + num_actual_tokens=num_tokens, + max_query_len=query_len, + block_table_tensor=block_table, + slot_mapping=slots, + causal=True, + ) + assert self.runtime.attn_metadata_builder is not None + metadata = self.runtime.attn_metadata_builder.build(0, common) + else: + metadata = MLACommonMetadata( + num_reqs=batch_size, + max_query_len=query_len, + max_seq_len=max_seq_len, + num_actual_tokens=num_tokens, + query_start_loc=query_start_gpu, + slot_mapping=slots, + num_decodes=batch_size, + num_decode_tokens=num_tokens, + num_prefills=0, + causal=False, + head_dim=int(next(iter(self.runtime.kv_caches.values())).shape[-1]), + prefill=None, + decode=MLACommonDecodeMetadata( + block_table=block_table, + seq_lens=seq_lens, + dcp_tot_seq_lens=None, + ), + ) + attn_metadata = {layer_name: metadata for layer_name in self.runtime.kv_caches} + slot_mapping = {layer_name: slots for layer_name in self.runtime.kv_caches} + state = _DraftCudaGraphState( + batch_size=batch_size, + num_speculative_tokens=num_speculative_tokens, + query_len=query_len, + input_ids=input_ids, + positions=positions, + slots=slots, + seq_lens=seq_lens, + block_table=block_table, + output_tokens=output_tokens, + input_ids_host=input_ids_host, + positions_host=positions_host, + slots_host=slots_host, + seq_lens_host=seq_lens_host, + block_table_host=block_table_host, + attn_metadata=attn_metadata, + slot_mapping=slot_mapping, + ) + + # Seed capture with valid, isolated dummy sequences. Runtime replay + # overwrites every graph input before use. + dummy_input_ids: list[int] = [] + dummy_positions: list[int] = [] + dummy_slots: list[int] = [] + dummy_rows: list[list[int]] = [] + for request_idx in range(batch_size): + dummy_input_ids.append(0) + dummy_input_ids.extend([self.mask_token_id] * (query_len - 1)) + dummy_positions.extend(range(query_len)) + dummy_slots.extend( + request_idx * self.allocator.block_size + position + for position in range(query_len) + ) + dummy_rows.append([request_idx]) + state.stage( + input_ids=dummy_input_ids, + positions=dummy_positions, + slots=dummy_slots, + block_rows=dummy_rows, + seq_lens=[query_len] * batch_size, + ) + return state + + def _run_cuda_graph_state( + self, + state: _DraftCudaGraphState, + ) -> None: + num_tokens = state.batch_size * state.query_len + with ( + set_current_vllm_config(self.runtime.draft_vllm_config), + set_forward_context( + state.attn_metadata, + self.runtime.draft_vllm_config, + num_tokens=num_tokens, + skip_compiled=True, + slot_mapping=state.slot_mapping, + ), + ): + hidden = self.model( + input_ids=state.input_ids, + positions=state.positions, + ) + if self.method == "dflash": + sample_hidden = hidden.view(state.batch_size, state.query_len, -1)[:, 1:] + logits = self.model.compute_logits( + sample_hidden.reshape( + state.batch_size * state.num_speculative_tokens, -1 + ) + ) + state.output_tokens.copy_( + logits.argmax(dim=-1).view( + state.batch_size, state.num_speculative_tokens + ) + ) + else: + logits = self.model.compute_draft_logits(hidden).view( + state.batch_size, state.query_len, -1 + ) + previous = state.input_ids.view(state.batch_size, state.query_len)[:, 0] + for step in range(state.query_len): + markov = self.model.markov_bias(self.model.markov_embed(previous)) + previous = (logits[:, step] + markov).argmax(dim=-1) + state.output_tokens[:, step].copy_(previous) + # Retain graph-owned outputs so their backing allocations cannot be + # recycled while captured nodes still reference them. + state.captured_hidden = hidden + state.captured_logits = logits + + @torch.inference_mode() + def capture_cuda_graphs(self, *, warmups: int = 2) -> None: + """Capture all DSpark or DFlash shapes selectable by this server.""" + if warmups < 1: + raise ValueError("Draft CUDA graph capture requires at least one warmup") + if self._cuda_graphs: + return + + started = time.perf_counter() + allocated_before = torch.cuda.memory_allocated(self.device) + capture_stream = torch.cuda.Stream(device=self.device) + capture_stream.wait_stream(torch.cuda.current_stream(self.device)) + # Capture the largest shape first. Triton MLA grows shared workspace + # buffers on demand; capturing a smaller shape first and then resizing + # that workspace for B2/K3 leaves the earlier graph with stale device + # pointers and causes an illegal access on replay. + shapes = [ + (batch_size, depth) + for batch_size in range( + self.runtime.vllm_config.scheduler_config.max_num_seqs, + 0, + -1, + ) + for depth in range(self.max_speculative_tokens, 0, -1) + ] + logger.info( + "Capturing %d standalone K3 %s CUDA graphs on %s: %s", + len(shapes), + self.method, + self.device, + ", ".join(f"B{batch_size}K{depth}" for batch_size, depth in shapes), + ) + with torch.cuda.stream(capture_stream): + for batch_size, depth in shapes: + state = self._make_cuda_graph_state(batch_size, depth) + for _ in range(warmups): + self._run_cuda_graph_state(state) + capture_stream.synchronize() + + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph( + graph, + stream=capture_stream, + ): + self._run_cuda_graph_state(state) + state.graph = graph + self._cuda_graphs[(batch_size, depth)] = state + + torch.cuda.current_stream(self.device).wait_stream(capture_stream) + torch.cuda.synchronize(self.device) + # Dummy capture rows only touch these reserved low blocks. Block zero + # is never assigned; live request allocation clears its own full range. + max_dummy_blocks = int(self.runtime.vllm_config.scheduler_config.max_num_seqs) + for cache in self.runtime.kv_caches.values(): + cache[:max_dummy_blocks].zero_() + torch.cuda.synchronize(self.device) + + self.cuda_graph_enabled = True + self.cuda_graph_capture_seconds = time.perf_counter() - started + self.cuda_graph_memory_gib = max( + 0.0, + (torch.cuda.memory_allocated(self.device) - allocated_before) / 1024**3, + ) + logger.info( + "Standalone K3 %s CUDA graphs ready in %.2fs; allocated_delta=%.3f GiB", + self.method, + self.cuda_graph_capture_seconds, + self.cuda_graph_memory_gib, + ) + + @property + def cuda_graph_shapes(self) -> list[str]: + return [ + f"B{batch_size}K{depth}" for batch_size, depth in sorted(self._cuda_graphs) + ] + + def _clear_state_cache( + self, + state: DraftRequestState, + *, + clear_context: bool = True, + ) -> None: + block_range = self.allocator.physical_block_range(state) + for cache in self.runtime.kv_caches.values(): + cache[block_range].zero_() + state.committed_end = 0 + state.context_start = 0 + if clear_context: + state.context_cache = None + + def reset(self, request_ids: list[str]) -> None: + with self._lock: + for request_id in request_ids: + state, _ = self.allocator.get_or_allocate(request_id) + self._clear_state_cache(state) + + def free(self, request_ids: list[str]) -> None: + with self._lock: + for request_id in request_ids: + self.allocator.free(request_id) + + def clear(self) -> None: + self.free(self.allocator.request_ids) + + @property + def prefix_cache_bytes(self) -> int: + return sum( + state.context_cache.allocated_bytes + for request_id in self.allocator.request_ids + if (state := self.allocator.get(request_id)) is not None + and state.context_cache is not None + ) + + @property + def prefix_cache_host_bytes(self) -> int: + return sum( + state.context_cache.allocated_bytes + for request_id in self.allocator.request_ids + if (state := self.allocator.get(request_id)) is not None + and state.context_cache is not None + and state.context_cache.device.type == "cpu" + ) + + @property + def prefix_cache_device_bytes(self) -> int: + return self.prefix_cache_bytes - self.prefix_cache_host_bytes + + @property + def mean_timing_ms(self) -> dict[str, float]: + if self.proposal_count <= 0: + return {} + return { + key: value / self.proposal_count + for key, value in self._timing_totals_ms.items() + } + + def _record_timing(self, timing_ms: dict[str, float]) -> None: + self.last_timing_ms = timing_ms + for key, value in timing_ms.items(): + self._timing_totals_ms[key] = self._timing_totals_ms.get(key, 0.0) + value + + def _restore_projected_context( + self, + state: DraftRequestState, + prefix_end: int, + ) -> int: + context_cache = state.context_cache + if context_cache is None: + raise ValueError( + f"No projected context is retained for {state.request_id!r}" + ) + restore_start = max(0, prefix_end - self.allocator.window_size) + restore_start = ( + restore_start // self.allocator.block_size * self.allocator.block_size + ) + if not context_cache.has_range(restore_start, prefix_end): + raise ValueError( + f"Projected context for {state.request_id!r} cannot restore " + f"prefix_end={prefix_end}; retained=" + f"[{context_cache.start_position}, {context_cache.end_position})" + ) + + self._clear_state_cache(state, clear_context=False) + for start in range(restore_start, prefix_end, self.max_context_tokens): + end = min(prefix_end, start + self.max_context_tokens) + context_states = context_cache.read(start, end) + positions = torch.arange(start, end, dtype=torch.int64) + context_gpu = context_states.to(self.device, non_blocking=True) + positions_gpu = positions.to(self.device, non_blocking=True) + slots = torch.tensor( + [ + self.allocator.cache_slot(state, position) + for position in range(start, end) + ], + dtype=torch.int64, + device=self.device, + ) + self.model.precompute_and_store_context_kv( + context_gpu, + positions_gpu, + slots, + ) + state.committed_end = prefix_end + state.context_start = restore_start + context_cache.truncate(prefix_end) + return restore_start + + def reconnect( + self, + source_request_id: str, + request_id: str, + prefix_end: int, + ) -> dict[str, Any]: + started = time.perf_counter() + with self._lock: + state = self.allocator.get(source_request_id) + if state is None: + raise KeyError(f"Unknown DSpark source request {source_request_id!r}") + context_cache = state.context_cache + restore_start = max(0, prefix_end - self.allocator.window_size) + restore_start = ( + restore_start // self.allocator.block_size * self.allocator.block_size + ) + if ( + prefix_end <= 0 + or context_cache is None + or not context_cache.has_range(restore_start, prefix_end) + ): + retained = ( + None + if context_cache is None + else [context_cache.start_position, context_cache.end_position] + ) + raise ValueError( + f"Cannot reconnect {source_request_id!r} at {prefix_end}; " + f"retained={retained}" + ) + state = self.allocator.rebind(source_request_id, request_id) + restored_start = self._restore_projected_context(state, prefix_end) + torch.accelerator.synchronize() + self.reconnect_count += 1 + self.last_reconnect_latency_ms = (time.perf_counter() - started) * 1000 + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "request_id": request_id, + "restored_start": restored_start, + "prefix_end": prefix_end, + "latency_ms": self.last_reconnect_latency_ms, + "active_requests": self.allocator.active_requests, + } + + def _decode_host_tensor( + self, + frame: bytes, + *, + dtype: torch.dtype, + shape: tuple[int, ...], + ) -> torch.Tensor: + element_size = torch.tensor([], dtype=dtype).element_size() + expected = element_size + for dim in shape: + expected *= dim + if len(frame) != expected: + raise ValueError( + f"Tensor frame has {len(frame)} bytes, expected {expected} for " + f"shape={shape}, dtype={dtype}" + ) + source = torch.frombuffer(frame, dtype=dtype) + if dtype == torch.int64: + staging = self._positions_staging + elif dtype == torch.bfloat16: + staging = self._context_staging + else: + raise TypeError(f"Unsupported DSpark RPC tensor dtype: {dtype}") + if source.numel() > staging.numel(): + raise ValueError( + f"Tensor frame exceeds pinned staging capacity: " + f"elements={source.numel()}, capacity={staging.numel()}" + ) + output = staging[: source.numel()] + output.copy_(source) + return output.view(shape) + + def _append_context( + self, + states: list[DraftRequestState], + context_counts: list[int], + positions_cpu: torch.Tensor, + context_cpu: torch.Tensor, + *, + projected: bool, + ) -> None: + if positions_cpu.numel() == 0: + return + positions = positions_cpu.to(self.device, non_blocking=True) + context_input = context_cpu.to(self.device, non_blocking=True) + context_states = ( + context_input + if projected + else self.model.combine_hidden_states(context_input) + ) + + slots: list[int] = [] + offset = 0 + for state, count in zip(states, context_counts): + req_positions = positions_cpu[offset : offset + count] + if count: + first = int(req_positions[0]) + if first > state.committed_end: + raise ValueError( + f"Context gap for {state.request_id!r}: " + f"expected <= {state.committed_end}, got {first}" + ) + expected = torch.arange(first, first + count, dtype=torch.int64) + if not torch.equal(req_positions, expected): + raise ValueError( + f"Context positions for {state.request_id!r} are not contiguous" + ) + state.committed_end = int(req_positions[-1]) + 1 + slots.extend( + self.allocator.cache_slot(state, int(position)) + for position in req_positions + ) + offset += count + slot_mapping = torch.tensor(slots, dtype=torch.int64, device=self.device) + self.model.precompute_and_store_context_kv( + context_states, + positions, + slot_mapping, + ) + offset = 0 + for state, count in zip(states, context_counts): + if count: + first_position = int(positions_cpu[offset]) + if state.context_cache is None: + state.context_cache = ProjectedContextCache( + hidden_size=self.hidden_size, + max_tokens=self.prefix_cache_tokens, + initial_position=first_position, + device=self.device, + ) + state.context_cache.append( + first_position, + context_states[offset : offset + count], + ) + offset += count + + def _run_query_block( + self, + states: list[DraftRequestState], + anchor_positions: list[int], + anchor_token_ids: list[int], + num_speculative_tokens: int, + ) -> torch.Tensor: + batch_size = len(states) + sample_from_anchor = self.method == "dspark" + query_len = ( + num_speculative_tokens if sample_from_anchor else 1 + num_speculative_tokens + ) + input_ids: list[int] = [] + positions: list[int] = [] + slots: list[int] = [] + block_rows: list[list[int]] = [] + seq_lens: list[int] = [] + + for state, anchor_position, anchor_token_id in zip( + states, anchor_positions, anchor_token_ids + ): + if anchor_position != state.committed_end: + raise ValueError( + f"Anchor position for {state.request_id!r} must equal the " + f"committed context end ({state.committed_end}), got " + f"{anchor_position}" + ) + sequence_end = anchor_position + query_len + if sequence_end > self.max_model_len: + raise ValueError( + f"Draft query exceeds max_model_len={self.max_model_len}: " + f"end={sequence_end}" + ) + input_ids.append(anchor_token_id) + input_ids.extend([self.mask_token_id] * (query_len - 1)) + req_positions = range(anchor_position, sequence_end) + positions.extend(req_positions) + slots.extend( + self.allocator.cache_slot(state, position) + for position in range(anchor_position, sequence_end) + ) + row, local_seq_len = self.allocator.block_table(state, sequence_end) + block_rows.append(row) + seq_lens.append(local_seq_len) + + if self.cuda_graph_enabled: + graph_state = self._cuda_graphs.get((batch_size, num_speculative_tokens)) + if graph_state is not None: + graph_state.stage( + input_ids=input_ids, + positions=positions, + slots=slots, + block_rows=block_rows, + seq_lens=seq_lens, + ) + assert graph_state.graph is not None + graph_state.graph.replay() + self.cuda_graph_replay_count += 1 + return graph_state.output_tokens + self.cuda_graph_eager_fallback_count += 1 + + max_blocks = max(len(row) for row in block_rows) + block_table = torch.zeros( + (batch_size, max_blocks), dtype=torch.int32, device=self.device + ) + for row_idx, row in enumerate(block_rows): + block_table[row_idx, : len(row)] = torch.tensor( + row, dtype=torch.int32, device=self.device + ) + input_ids_gpu = torch.tensor(input_ids, dtype=torch.int64, device=self.device) + positions_gpu = torch.tensor(positions, dtype=torch.int64, device=self.device) + slots_gpu = torch.tensor(slots, dtype=torch.int64, device=self.device) + seq_lens_gpu = torch.tensor(seq_lens, dtype=torch.int32, device=self.device) + query_start_loc = torch.arange( + 0, + (batch_size + 1) * query_len, + query_len, + dtype=torch.int32, + device=self.device, + ) + if self.method == "dflash": + query_start_loc_cpu = torch.arange( + 0, + (batch_size + 1) * query_len, + query_len, + dtype=torch.int32, + ) + common = CommonAttentionMetadata( + query_start_loc=query_start_loc, + query_start_loc_cpu=query_start_loc_cpu, + seq_lens=seq_lens_gpu, + seq_lens_cpu_upper_bound=torch.tensor(seq_lens, dtype=torch.int32), + max_seq_len=max(seq_lens), + num_reqs=batch_size, + num_actual_tokens=batch_size * query_len, + max_query_len=query_len, + block_table_tensor=block_table, + slot_mapping=slots_gpu, + causal=True, + ) + if self.runtime.attn_metadata_builder is None: + raise RuntimeError("K3 DFlash attention metadata builder is missing") + metadata = self.runtime.attn_metadata_builder.build(0, common) + else: + metadata = MLACommonMetadata( + num_reqs=batch_size, + max_query_len=query_len, + max_seq_len=max(seq_lens), + num_actual_tokens=batch_size * query_len, + query_start_loc=query_start_loc, + slot_mapping=slots_gpu, + num_decodes=batch_size, + num_decode_tokens=batch_size * query_len, + num_prefills=0, + causal=False, + head_dim=int(next(iter(self.runtime.kv_caches.values())).shape[-1]), + prefill=None, + decode=MLACommonDecodeMetadata( + block_table=block_table, + seq_lens=seq_lens_gpu, + dcp_tot_seq_lens=None, + ), + ) + attn_metadata = {layer_name: metadata for layer_name in self.runtime.kv_caches} + slot_mapping = {layer_name: slots_gpu for layer_name in self.runtime.kv_caches} + with ( + set_current_vllm_config(self.runtime.draft_vllm_config), + set_forward_context( + attn_metadata, + self.runtime.draft_vllm_config, + num_tokens=batch_size * query_len, + skip_compiled=True, + slot_mapping=slot_mapping, + ), + ): + hidden = self.model(input_ids=input_ids_gpu, positions=positions_gpu) + + if self.method == "dflash": + sample_hidden = hidden.view(batch_size, query_len, -1)[:, 1:] + return ( + self.model.compute_logits( + sample_hidden.reshape(batch_size * num_speculative_tokens, -1) + ) + .argmax(dim=-1) + .view(batch_size, num_speculative_tokens) + ) + + base_logits = self.model.compute_draft_logits(hidden).view( + batch_size, query_len, -1 + ) + previous = torch.tensor(anchor_token_ids, dtype=torch.int64, device=self.device) + draft_tokens = torch.empty( + (batch_size, query_len), dtype=torch.int64, device=self.device + ) + for step in range(query_len): + markov = self.model.markov_bias(self.model.markov_embed(previous)) + previous = (base_logits[:, step] + markov).argmax(dim=-1) + draft_tokens[:, step].copy_(previous) + return draft_tokens + + @torch.inference_mode() + def propose(self, header: dict[str, Any], frames: list[bytes]) -> dict[str, Any]: + started = time.perf_counter() + requests = header.get("requests") + if not isinstance(requests, list) or not requests: + raise ValueError("PROPOSE requires a non-empty requests list") + if len(requests) > self.allocator.max_requests: + raise ValueError( + f"Batch has {len(requests)} requests, max is " + f"{self.allocator.max_requests}" + ) + num_speculative_tokens = int( + header.get("num_speculative_tokens", self.max_speculative_tokens) + ) + if not 1 <= num_speculative_tokens <= self.max_speculative_tokens: + raise ValueError( + f"num_speculative_tokens must be in [1, {self.max_speculative_tokens}]" + ) + projected = bool(header.get("projected", False)) + context_counts = [int(req.get("context_count", 0)) for req in requests] + if any(count < 0 for count in context_counts): + raise ValueError("context_count cannot be negative") + total_context = sum(context_counts) + expected_frames = 2 if total_context else 0 + if len(frames) != expected_frames: + raise ValueError( + f"PROPOSE expected {expected_frames} tensor frames, got {len(frames)}" + ) + context_width = self.hidden_size if projected else self.raw_context_width + if total_context: + positions_cpu = self._decode_host_tensor( + frames[0], dtype=torch.int64, shape=(total_context,) + ) + context_cpu = self._decode_host_tensor( + frames[1], + dtype=torch.bfloat16, + shape=(total_context, context_width), + ) + else: + positions_cpu = torch.empty(0, dtype=torch.int64) + context_cpu = torch.empty((0, context_width), dtype=torch.bfloat16) + decoded_at = time.perf_counter() + + lock_started = time.perf_counter() + with self._lock: + lock_acquired = time.perf_counter() + gpu_start = torch.cuda.Event(enable_timing=True) + context_end = torch.cuda.Event(enable_timing=True) + query_end = torch.cuda.Event(enable_timing=True) + gpu_start.record() + states: list[DraftRequestState] = [] + for req in requests: + request_id = req.get("request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("Every request requires a non-empty request_id") + state, created = self.allocator.get_or_allocate(request_id) + if bool(req.get("reset", False)) or created: + self._clear_state_cache(state) + reset_position = int(req.get("reset_position", 0)) + if not 0 <= reset_position <= self.max_model_len: + raise ValueError( + "Draft reset_position must be in model bounds, got " + f"{reset_position}" + ) + state.committed_end = reset_position + state.context_start = reset_position + if reset_position: + self.cold_bootstrap_count += 1 + states.append(state) + self._append_context( + states, + context_counts, + positions_cpu, + context_cpu, + projected=projected, + ) + context_end.record() + anchor_positions = [int(req["anchor_position"]) for req in requests] + anchor_token_ids = [int(req["anchor_token_id"]) for req in requests] + draft_tokens = self._run_query_block( + states, + anchor_positions, + anchor_token_ids, + num_speculative_tokens, + ) + query_end.record() + submit_done = time.perf_counter() + query_end.synchronize() + sync_done = time.perf_counter() + tokens = draft_tokens.cpu().tolist() + tokens_copied = time.perf_counter() + self.proposal_count += 1 + self.last_latency_ms = (tokens_copied - started) * 1000 + timing_ms = { + "decode_frames": (decoded_at - started) * 1000, + "lock_wait": (lock_acquired - lock_started) * 1000, + "host_submit": (submit_done - lock_acquired) * 1000, + "gpu_context": gpu_start.elapsed_time(context_end), + "gpu_query": context_end.elapsed_time(query_end), + "gpu_wait": (sync_done - submit_done) * 1000, + "tokens_d2h": (tokens_copied - sync_done) * 1000, + "total": self.last_latency_ms, + } + self._record_timing(timing_ms) + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "tokens": tokens, + "latency_ms": self.last_latency_ms, + "timing_ms": timing_ms, + "active_requests": self.allocator.active_requests, + } + + +class K3DSparkZMQServer: + def __init__( + self, + engine: K3DSparkDraftEngine, + *, + address: str, + stop: threading.Event, + ) -> None: + self.engine = engine + self.address = address + self.stop = stop + self.ready = threading.Event() + self.error: str | None = None + self._thread = threading.Thread( + target=self._run, + name="k3-draft-zmq", + daemon=True, + ) + + def start(self) -> None: + self._thread.start() + if not self.ready.wait(timeout=10): + raise TimeoutError(f"K3 draft proposal socket did not bind: {self.address}") + if self.error is not None: + raise RuntimeError(self.error) + + def join(self, timeout: float = 5.0) -> None: + self._thread.join(timeout=timeout) + + def _handle(self, parts: list[bytes]) -> dict[str, Any]: + if not parts: + raise ValueError("Empty DSpark RPC message") + header = json.loads(parts[0]) + if not isinstance(header, dict): + raise ValueError("DSpark RPC header must be a JSON object") + if int(header.get("protocol", -1)) != PROTOCOL_VERSION: + raise ValueError( + f"Unsupported protocol {header.get('protocol')}; " + f"expected {PROTOCOL_VERSION}" + ) + op = str(header.get("op", "")).upper() + if op == "PING": + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "op": "PONG", + "method": self.engine.method, + "active_requests": self.engine.allocator.active_requests, + "max_requests": self.engine.allocator.max_requests, + "block_size": self.engine.allocator.block_size, + "window_size": self.engine.allocator.window_size, + "prefix_cache_tokens": self.engine.prefix_cache_tokens, + "prefix_cache_device": str(self.engine.device), + "cold_bootstrap_count": self.engine.cold_bootstrap_count, + "cuda_graph_enabled": self.engine.cuda_graph_enabled, + "cuda_graph_shapes": self.engine.cuda_graph_shapes, + } + if op == "CLEAR": + self.engine.clear() + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "active_requests": 0, + } + if op in ("RESET", "FREE"): + request_ids = header.get("request_ids") + if not isinstance(request_ids, list) or not all( + isinstance(req_id, str) and req_id for req_id in request_ids + ): + raise ValueError(f"{op} requires request_ids: list[str]") + if op == "RESET": + self.engine.reset(request_ids) + else: + self.engine.free(request_ids) + return { + "ok": True, + "protocol": PROTOCOL_VERSION, + "active_requests": self.engine.allocator.active_requests, + } + if op == "RECONNECT": + source_request_id = header.get("source_request_id") + request_id = header.get("request_id") + if not isinstance(source_request_id, str) or not source_request_id: + raise ValueError("RECONNECT requires source_request_id") + if not isinstance(request_id, str) or not request_id: + raise ValueError("RECONNECT requires request_id") + return self.engine.reconnect( + source_request_id, + request_id, + int(header.get("prefix_end", 0)), + ) + if op == "PROPOSE": + return self.engine.propose(header, parts[1:]) + raise ValueError(f"Unknown K3 draft RPC operation: {op!r}") + + def _run(self) -> None: + context = zmq.Context() + socket = context.socket(zmq.REP) + socket.setsockopt(zmq.LINGER, 0) + try: + socket.bind(self.address) + logger.info( + "K3 %s proposal RPC listening on %s", + self.engine.method, + self.address, + ) + self.ready.set() + poller = zmq.Poller() + poller.register(socket, zmq.POLLIN) + while not self.stop.is_set(): + if not dict(poller.poll(250)).get(socket): + continue + try: + response = self._handle(socket.recv_multipart()) + except Exception as exc: + logger.exception("K3 DSpark proposal request failed") + response = { + "ok": False, + "protocol": PROTOCOL_VERSION, + "error": f"{type(exc).__name__}: {exc}", + } + socket.send_json(response) + except Exception as exc: + self.error = f"{type(exc).__name__}: {exc}" + logger.exception("K3 DSpark proposal server failed") + self.ready.set() + finally: + socket.close() + context.term() diff --git a/vllm/entrypoints/k3_dspark_standalone.py b/vllm/entrypoints/k3_dspark_standalone.py new file mode 100644 index 000000000000..25505ff3fc59 --- /dev/null +++ b/vllm/entrypoints/k3_dspark_standalone.py @@ -0,0 +1,847 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Load a Kimi-K3 draft model on a dedicated single GPU. + +The process loads either DSpark or DFlash plus the target embedding and LM +head, without loading the Kimi-K3 target transformer. Draft weights, KV cache, +attention workspace, and proposal compute stay on this process's GPU. +""" + +from __future__ import annotations + +import argparse +import json +import signal +import threading +import time +from dataclasses import asdict, dataclass +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from pathlib import Path +from typing import Any + +import torch +import torch.nn as nn +from safetensors import safe_open + +from vllm.config.vllm import set_current_vllm_config +from vllm.engine.arg_utils import EngineArgs +from vllm.forward_context import set_forward_context +from vllm.logger import init_logger +from vllm.model_executor.layers.vocab_parallel_embedding import ( + ParallelLMHead, + VocabParallelEmbedding, +) +from vllm.utils.torch_utils import set_default_torch_dtype +from vllm.v1.worker.workspace import init_workspace_manager + +logger = init_logger(__name__) + +EMBED_TENSOR = "language_model.model.embed_tokens.weight" +LM_HEAD_TENSOR = "language_model.lm_head.weight" +SHARED_TENSORS = (EMBED_TENSOR, LM_HEAD_TENSOR) + + +@dataclass +class RuntimeStatus: + phase: str = "starting" + ready: bool = False + proposal_transport_ready: bool = False + device: str = "" + compute_capability: str = "" + torch_version: str = torch.__version__ + torch_arches: tuple[str, ...] = () + draft_model: str = "" + method: str = "" + target_weights: str = "" + num_speculative_tokens: int = 0 + max_model_len: int = 0 + allocated_gib: float = 0.0 + reserved_gib: float = 0.0 + draft_kv_cache_gib: float = 0.0 + draft_kv_cache_blocks: int = 0 + draft_kv_cache_token_capacity: int = 0 + draft_kv_cache_smoke: bool = False + draft_kv_window: int = 0 + proposal_address: str | None = None + proposal_count: int = 0 + proposal_active_requests: int = 0 + proposal_last_latency_ms: float = 0.0 + smoke_token: int | None = None + load_seconds: float = 0.0 + error: str | None = None + + +class _TargetLanguageModel(nn.Module): + def __init__(self, vocab_size: int, hidden_size: int) -> None: + super().__init__() + self.embed_tokens = VocabParallelEmbedding( + vocab_size, + hidden_size, + params_dtype=torch.bfloat16, + prefix="model.embed_tokens", + disable_tp=True, + ) + + +class StandaloneTargetFacade(nn.Module): + """Only the two frozen target modules that DSpark shares.""" + + def __init__(self, vocab_size: int, hidden_size: int) -> None: + super().__init__() + self.model = _TargetLanguageModel(vocab_size, hidden_size) + self.lm_head = ParallelLMHead( + vocab_size, + hidden_size, + params_dtype=torch.bfloat16, + prefix="lm_head", + disable_tp=True, + ) + + def get_language_model(self) -> StandaloneTargetFacade: + return self + + +@dataclass +class StandaloneRuntime: + vllm_config: Any + draft_vllm_config: Any + target_facade: StandaloneTargetFacade + model: nn.Module + kv_caches: dict[str, torch.Tensor] + kv_cache_block_size: int + method: str + attn_metadata_builder: Any | None = None + + +def resolve_shared_weight_files(target_weights: Path) -> dict[str, Path]: + """Resolve the two shared tensors through a safetensors index.""" + root = target_weights.resolve() + index_path = root / "model.safetensors.index.json" + if not index_path.is_file(): + raise FileNotFoundError(f"Target weight index is missing: {index_path}") + + index = json.loads(index_path.read_text()) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"Invalid safetensors weight_map in {index_path}") + + resolved: dict[str, Path] = {} + for tensor_name in SHARED_TENSORS: + relative = weight_map.get(tensor_name) + if not isinstance(relative, str): + raise KeyError(f"Target tensor is absent from the index: {tensor_name}") + tensor_path = (root / relative).resolve() + if not tensor_path.is_relative_to(root): + raise ValueError( + f"Target tensor path escapes the checkpoint root: {tensor_path}" + ) + if not tensor_path.is_file(): + raise FileNotFoundError(f"Target tensor file is missing: {tensor_path}") + resolved[tensor_name] = tensor_path + return resolved + + +def _load_module_weight( + module: VocabParallelEmbedding, + tensor_name: str, + tensor_path: Path, + device: torch.device, +) -> None: + logger.info("Loading shared target tensor %s from %s", tensor_name, tensor_path) + with safe_open(str(tensor_path), framework="pt", device=str(device)) as handle: + if tensor_name not in set(handle.keys()): + raise KeyError(f"{tensor_name} is absent from {tensor_path}") + loaded_weight = handle.get_tensor(tensor_name) + if loaded_weight.dtype != torch.bfloat16: + raise TypeError( + f"{tensor_name} must be bfloat16, got {loaded_weight.dtype}" + ) + module.weight_loader(module.weight, loaded_weight) + del loaded_weight + torch.accelerator.empty_cache() + + +def load_shared_target_weights( + target: StandaloneTargetFacade, + target_weights: Path, + device: torch.device, +) -> None: + files = resolve_shared_weight_files(target_weights) + _load_module_weight( + target.model.embed_tokens, + EMBED_TENSOR, + files[EMBED_TENSOR], + device, + ) + _load_module_weight( + target.lm_head, + LM_HEAD_TENSOR, + files[LM_HEAD_TENSOR], + device, + ) + + +def _validate_cuda_runtime(device: torch.device) -> tuple[str, tuple[str, ...]]: + if not torch.cuda.is_available(): + raise RuntimeError("CUDA is not available in the DSpark container") + major, minor = torch.cuda.get_device_capability(device) + expected_arch = f"sm_{major}{minor}" + arches = tuple(torch.cuda.get_arch_list()) + if expected_arch not in arches: + raise RuntimeError( + f"PyTorch does not contain {expected_arch}; compiled arches are {arches}" + ) + if major < 8: + raise RuntimeError( + f"Kimi-K3 DSpark BF16 requires compute capability >= 8.0, got " + f"{major}.{minor}" + ) + return f"{major}.{minor}", arches + + +def _init_single_gpu_distributed() -> None: + from vllm.distributed.parallel_state import ( + ensure_model_parallel_initialized, + init_distributed_environment, + model_parallel_is_initialized, + ) + from vllm.utils.network_utils import get_open_port + + if model_parallel_is_initialized(): + return + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", + backend="gloo", + ) + ensure_model_parallel_initialized( + tensor_model_parallel_size=1, + pipeline_model_parallel_size=1, + ) + + +def _build_vllm_config(args: argparse.Namespace): + attention_backend = "TRITON_ATTN" if args.method == "dflash" else "TRITON_MLA" + speculative_config = { + "model": str(args.draft_model), + "method": args.method, + "num_speculative_tokens": args.num_speculative_tokens, + "attention_backend": attention_backend, + "kv_cache_dtype": "bfloat16", + "draft_sample_method": "greedy", + "rejection_sample_method": "block", + "max_model_len": args.max_model_len, + } + engine_args = EngineArgs( + model=str(args.target_config), + tokenizer_mode="skip", + skip_tokenizer_init=True, + dtype="bfloat16", + kv_cache_dtype="bfloat16", + max_model_len=args.max_model_len, + tensor_parallel_size=1, + decode_context_parallel_size=1, + max_num_batched_tokens=args.max_num_batched_tokens, + max_num_seqs=args.max_num_seqs, + block_size=16, + enable_prefix_caching=False, + enforce_eager=True, + compilation_config={"custom_ops": ["none"]}, + kernel_config={ + "ir_op_priority": { + "rms_norm": ["native"], + "fused_add_rms_norm": ["native"], + }, + "linear_backend": "torch", + }, + load_format="safetensors", + use_tqdm_on_load=False, + speculative_config=speculative_config, + ) + return engine_args.create_engine_config(headless=True) + + +def _allocate_draft_kv_cache( + model: nn.Module, + *, + method: str, + block_size: int, + cache_gib: float, + device: torch.device, +) -> tuple[dict[str, torch.Tensor], int]: + """Allocate and bind the draft model's private BF16 KV cache.""" + if cache_gib <= 0: + raise ValueError(f"--draft-kv-cache-gib must be positive, got {cache_gib}") + + if method == "dflash": + attentions = [layer.self_attn.attn for layer in model.model.layers] + else: + attentions = [layer.self_attn for layer in model.model.layers] + if not attentions: + raise RuntimeError("K3 draft model does not expose any attention layers") + + if method == "dflash": + cache_shapes = { + ( + int(attn.impl.num_kv_heads), + int(attn.impl.head_size), + ) + for attn in attentions + } + if len(cache_shapes) != 1: + raise ValueError(f"K3 DFlash layers have mixed KV shapes: {cache_shapes}") + num_kv_heads, head_size = cache_shapes.pop() + bytes_per_block_all_layers = ( + len(attentions) + * block_size + * num_kv_heads + * 2 + * head_size + * torch.tensor([], dtype=torch.bfloat16).element_size() + ) + else: + latent_widths = { + int(attn.kv_lora_rank + attn.qk_rope_head_dim) for attn in attentions + } + if len(latent_widths) != 1: + raise ValueError( + f"K3 DSpark draft layers have mixed MLA widths: {latent_widths}" + ) + latent_width = latent_widths.pop() + bytes_per_block_all_layers = ( + len(attentions) + * block_size + * latent_width + * torch.tensor([], dtype=torch.bfloat16).element_size() + ) + num_blocks = int(cache_gib * 1024**3) // bytes_per_block_all_layers + # Block zero is deliberately never assigned to a live request. + if num_blocks < 2: + minimum_mib = 2 * bytes_per_block_all_layers / 1024**2 + raise ValueError( + "The draft KV cache must fit a null block plus one data block; " + f"minimum={minimum_mib:.2f} MiB" + ) + + caches: dict[str, torch.Tensor] = {} + for attn in attentions: + if method == "dflash": + # TritonAttention exposes logical B,H,N,2D while its preferred + # physical layout is B,N,H,2D on CUDA. + physical = torch.zeros( + (num_blocks, block_size, num_kv_heads, 2 * head_size), + dtype=torch.bfloat16, + device=device, + ) + cache = physical.permute(0, 2, 1, 3) + else: + cache = torch.zeros( + (num_blocks, block_size, latent_width), + dtype=torch.bfloat16, + device=device, + ) + attn.bind_kv_cache(cache) + caches[attn.layer_name] = cache + return caches, num_blocks + + +def _build_dflash_metadata_builder( + model: nn.Module, + draft_vllm_config: Any, + device: torch.device, +) -> Any: + from vllm.v1.worker.utils import AttentionGroup + + attentions = [layer.self_attn.attn for layer in model.model.layers] + layer_names = [attn.layer_name for attn in attentions] + first = attentions[0] + backend = first.get_attn_backend() + if backend.get_name() != "TRITON_ATTN": + raise ValueError( + f"Standalone K3 DFlash requires TRITON_ATTN, got {backend.get_name()}" + ) + kv_spec = first.get_kv_cache_spec(draft_vllm_config) + if kv_spec is None: + raise RuntimeError("K3 DFlash attention did not return a KV cache spec") + group = AttentionGroup(backend, layer_names, kv_spec, 0) + group.create_metadata_builders( + draft_vllm_config, + device, + kernel_block_size=int(draft_vllm_config.cache_config.block_size), + ) + return group.get_metadata_builder() + + +def _load_runtime(args: argparse.Namespace, status: RuntimeStatus) -> StandaloneRuntime: + start = time.perf_counter() + device = torch.device("cuda", args.device) + torch.cuda.set_device(device) + status.device = torch.cuda.get_device_name(device) + status.compute_capability, status.torch_arches = _validate_cuda_runtime(device) + status.phase = "building_config" + + vllm_config = _build_vllm_config(args) + speculative_config = vllm_config.speculative_config + assert speculative_config is not None + draft_config = speculative_config.draft_model_config.hf_config + + status.phase = "loading_shared_target_weights" + with set_current_vllm_config(vllm_config): + _init_single_gpu_distributed() + init_workspace_manager(device, num_lanes=2) + with torch.device(device), set_default_torch_dtype(torch.bfloat16): + target = StandaloneTargetFacade( + vocab_size=draft_config.vocab_size, + hidden_size=draft_config.hidden_size, + ) + load_shared_target_weights(target, args.target_weights, device) + + status.phase = f"loading_{args.method}" + from vllm.v1.worker.gpu.spec_decode.eagle.utils import ( + _create_draft_vllm_config, + ) + + if args.method == "dflash": + from vllm.v1.worker.gpu.spec_decode.dflash.utils import ( + load_dflash_model, + maybe_load_mask_embedding, + ) + + model = load_dflash_model(target, vllm_config) + maybe_load_mask_embedding( + model, + str(args.draft_model), + int(draft_config.dflash_config["mask_token_id"]), + ) + else: + from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model + + model = load_dspark_model(target, vllm_config) + model.eval() + draft_vllm_config = _create_draft_vllm_config(vllm_config) + + status.phase = "allocating_draft_kv_cache" + block_size = int(vllm_config.cache_config.block_size) + kv_caches, num_blocks = _allocate_draft_kv_cache( + model, + method=args.method, + block_size=block_size, + cache_gib=args.draft_kv_cache_gib, + device=device, + ) + status.draft_kv_cache_blocks = num_blocks + status.draft_kv_cache_token_capacity = (num_blocks - 1) * block_size + status.draft_kv_cache_gib = ( + sum(cache.numel() * cache.element_size() for cache in kv_caches.values()) + / 1024**3 + ) + + torch.accelerator.synchronize() + status.load_seconds = time.perf_counter() - start + status.allocated_gib = torch.cuda.memory_allocated(device) / 1024**3 + status.reserved_gib = torch.cuda.memory_reserved(device) / 1024**3 + metadata_builder = ( + _build_dflash_metadata_builder(model, draft_vllm_config, device) + if args.method == "dflash" + else None + ) + return StandaloneRuntime( + vllm_config, + draft_vllm_config, + target, + model, + kv_caches, + block_size, + args.method, + metadata_builder, + ) + + +@torch.inference_mode() +def _run_eager_smoke(runtime: StandaloneRuntime, device: torch.device) -> int: + if runtime.method == "dflash": + return _run_dflash_eager_smoke(runtime, device) + + from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonDecodeMetadata, + MLACommonMetadata, + ) + from vllm.v1.worker.gpu.spec_decode.utils import ( + get_parallel_drafting_token_id, + ) + + model = runtime.model + draft_config = runtime.vllm_config.speculative_config.draft_model_config.hf_config + num_aux_layers = int(draft_config.num_target_layers) + hidden_size = int(draft_config.hidden_size) + + aux = torch.zeros( + (1, num_aux_layers * hidden_size), + dtype=torch.bfloat16, + device=device, + ) + context_len = 1 + positions = torch.zeros(context_len, dtype=torch.int64, device=device) + context = model.combine_hidden_states(aux) + data_block = 1 + context_slots = torch.tensor( + [data_block * runtime.kv_cache_block_size], + dtype=torch.int64, + device=device, + ) + model.precompute_and_store_context_kv(context, positions, context_slots) + + sample_from_anchor = bool(getattr(draft_config, "sample_from_anchor", True)) + query_len = runtime.vllm_config.speculative_config.num_speculative_tokens + if not sample_from_anchor: + query_len += 1 + mask_token_id = get_parallel_drafting_token_id(draft_config) + input_ids = torch.tensor( + [draft_config.bos_token_id] + [mask_token_id] * (query_len - 1), + dtype=torch.int64, + device=device, + ) + query_positions = torch.arange( + context_len, + context_len + query_len, + dtype=torch.int64, + device=device, + ) + query_slots = torch.arange( + data_block * runtime.kv_cache_block_size + context_len, + data_block * runtime.kv_cache_block_size + context_len + query_len, + dtype=torch.int64, + device=device, + ) + query_start_loc = torch.tensor([0, query_len], dtype=torch.int32, device=device) + block_table = torch.tensor([[data_block]], dtype=torch.int32, device=device) + seq_lens = torch.tensor([context_len + query_len], dtype=torch.int32, device=device) + metadata = MLACommonMetadata( + num_reqs=1, + max_query_len=query_len, + max_seq_len=context_len + query_len, + num_actual_tokens=query_len, + query_start_loc=query_start_loc, + slot_mapping=query_slots, + num_decodes=1, + num_decode_tokens=query_len, + num_prefills=0, + causal=False, + # MLA metadata validates the cached latent width, not the full Q/K + # projection width. + head_dim=int(next(iter(runtime.kv_caches.values())).shape[-1]), + prefill=None, + decode=MLACommonDecodeMetadata( + block_table=block_table, + seq_lens=seq_lens, + dcp_tot_seq_lens=None, + ), + ) + attn_metadata = {layer_name: metadata for layer_name in runtime.kv_caches} + slot_mapping = {layer_name: query_slots for layer_name in runtime.kv_caches} + with ( + set_current_vllm_config(runtime.draft_vllm_config), + set_forward_context( + attn_metadata, + runtime.draft_vllm_config, + num_tokens=query_len, + skip_compiled=True, + slot_mapping=slot_mapping, + ), + ): + hidden = model(input_ids=input_ids, positions=query_positions) + base_logits = model.compute_draft_logits(hidden[-1:]) + markov = model.markov_bias(model.markov_embed(input_ids[-1:])) + token = int((base_logits + markov).argmax(dim=-1).item()) + torch.accelerator.synchronize() + for cache in runtime.kv_caches.values(): + cache[data_block].zero_() + return token + + +@torch.inference_mode() +def _run_dflash_eager_smoke( + runtime: StandaloneRuntime, + device: torch.device, +) -> int: + from vllm.v1.attention.backend import CommonAttentionMetadata + from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, + ) + from vllm.v1.worker.gpu.spec_decode.utils import ( + get_parallel_drafting_token_id, + ) + + speculative_config = runtime.vllm_config.speculative_config + assert speculative_config is not None + draft_config = speculative_config.draft_model_config.hf_config + aux_layers = get_eagle3_aux_layers_from_config(speculative_config) + if not aux_layers: + raise ValueError("K3 DFlash config does not declare target auxiliary layers") + + raw_context = torch.zeros( + (1, len(aux_layers) * int(draft_config.hidden_size)), + dtype=torch.bfloat16, + device=device, + ) + context = runtime.model.combine_hidden_states(raw_context) + block_size = runtime.kv_cache_block_size + context_position = torch.zeros(1, dtype=torch.int64, device=device) + context_slot = torch.tensor([block_size], dtype=torch.int64, device=device) + runtime.model.precompute_and_store_context_kv( + context, + context_position, + context_slot, + ) + + query_len = int(speculative_config.num_speculative_tokens) + 1 + sequence_end = 1 + query_len + block_ids = list(range(1, 1 + (sequence_end + block_size - 1) // block_size)) + input_ids = torch.tensor( + [draft_config.bos_token_id] + + [get_parallel_drafting_token_id(draft_config)] * (query_len - 1), + dtype=torch.int64, + device=device, + ) + positions = torch.arange(1, sequence_end, dtype=torch.int64, device=device) + slots = torch.tensor( + [ + block_ids[position // block_size] * block_size + position % block_size + for position in range(1, sequence_end) + ], + dtype=torch.int64, + device=device, + ) + query_start_cpu = torch.tensor([0, query_len], dtype=torch.int32) + query_start_gpu = query_start_cpu.to(device) + seq_lens = torch.tensor([sequence_end], dtype=torch.int32, device=device) + block_table = torch.tensor([block_ids], dtype=torch.int32, device=device) + common = CommonAttentionMetadata( + query_start_loc=query_start_gpu, + query_start_loc_cpu=query_start_cpu, + seq_lens=seq_lens, + seq_lens_cpu_upper_bound=torch.tensor([sequence_end], dtype=torch.int32), + max_seq_len=sequence_end, + num_reqs=1, + num_actual_tokens=query_len, + max_query_len=query_len, + block_table_tensor=block_table, + slot_mapping=slots, + causal=True, + ) + assert runtime.attn_metadata_builder is not None + metadata = runtime.attn_metadata_builder.build(0, common) + attn_metadata = {layer_name: metadata for layer_name in runtime.kv_caches} + slot_mapping = {layer_name: slots for layer_name in runtime.kv_caches} + with ( + set_current_vllm_config(runtime.draft_vllm_config), + set_forward_context( + attn_metadata, + runtime.draft_vllm_config, + num_tokens=query_len, + skip_compiled=True, + slot_mapping=slot_mapping, + ), + ): + hidden = runtime.model(input_ids=input_ids, positions=positions) + token = int(runtime.model.compute_logits(hidden[1:2]).argmax(dim=-1).item()) + torch.accelerator.synchronize() + for cache in runtime.kv_caches.values(): + cache[1 : 1 + len(block_ids)].zero_() + return token + + +def _make_handler(status: RuntimeStatus, proposal_engine: Any | None = None): + class StatusHandler(BaseHTTPRequestHandler): + def do_GET(self) -> None: # noqa: N802 + if self.path not in ("/", "/healthz", "/readyz", "/v1/status"): + self.send_error(404) + return + payload = asdict(status) + if proposal_engine is not None: + payload["proposal_count"] = proposal_engine.proposal_count + payload["proposal_active_requests"] = ( + proposal_engine.allocator.active_requests + ) + payload["proposal_max_requests"] = ( + proposal_engine.allocator.max_requests + ) + payload["proposal_last_latency_ms"] = proposal_engine.last_latency_ms + payload["proposal_last_timing_ms"] = proposal_engine.last_timing_ms + payload["proposal_mean_timing_ms"] = proposal_engine.mean_timing_ms + payload["proposal_cold_bootstrap_count"] = ( + proposal_engine.cold_bootstrap_count + ) + payload["proposal_reconnect_count"] = proposal_engine.reconnect_count + payload["proposal_last_reconnect_latency_ms"] = ( + proposal_engine.last_reconnect_latency_ms + ) + payload["proposal_prefix_cache_tokens"] = ( + proposal_engine.prefix_cache_tokens + ) + payload["proposal_prefix_cache_host_gib"] = ( + proposal_engine.prefix_cache_host_bytes / 1024**3 + ) + payload["proposal_prefix_cache_gpu_gib"] = ( + proposal_engine.prefix_cache_device_bytes / 1024**3 + ) + payload["proposal_cuda_graph_enabled"] = ( + proposal_engine.cuda_graph_enabled + ) + payload["proposal_cuda_graph_shapes"] = ( + proposal_engine.cuda_graph_shapes + ) + payload["proposal_cuda_graph_capture_seconds"] = ( + proposal_engine.cuda_graph_capture_seconds + ) + payload["proposal_cuda_graph_memory_gib"] = ( + proposal_engine.cuda_graph_memory_gib + ) + payload["proposal_cuda_graph_replay_count"] = ( + proposal_engine.cuda_graph_replay_count + ) + payload["proposal_cuda_graph_eager_fallback_count"] = ( + proposal_engine.cuda_graph_eager_fallback_count + ) + body = json.dumps(payload, sort_keys=True).encode() + code = 200 if status.ready else 503 + self.send_response(code) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def log_message(self, format: str, *args: object) -> None: + logger.info("DSpark status HTTP: %s", format % args) + + return StatusHandler + + +def _serve_status( + args: argparse.Namespace, + status: RuntimeStatus, + stop: threading.Event, + proposal_engine: Any | None = None, +) -> None: + server = ThreadingHTTPServer( + (args.host, args.port), _make_handler(status, proposal_engine) + ) + server.timeout = 0.5 + + def request_stop(signum: int, _frame: object) -> None: + logger.info("Received signal %d; stopping DSpark status server", signum) + stop.set() + + signal.signal(signal.SIGINT, request_stop) + signal.signal(signal.SIGTERM, request_stop) + logger.info( + "DSpark status endpoint listening on http://%s:%d", args.host, args.port + ) + while not stop.is_set(): + server.handle_request() + server.server_close() + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--method", choices=("dspark", "dflash"), default="dspark") + parser.add_argument("--draft-model", type=Path, required=True) + parser.add_argument("--target-weights", type=Path, required=True) + parser.add_argument("--target-config", type=Path, required=True) + parser.add_argument("--device", type=int, default=0) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=8091) + parser.add_argument("--num-speculative-tokens", type=int, default=3) + parser.add_argument("--max-model-len", type=int, default=32768) + parser.add_argument("--max-num-batched-tokens", type=int, default=64) + parser.add_argument("--max-num-seqs", type=int, default=2) + parser.add_argument("--max-retained-requests", type=int) + parser.add_argument("--draft-kv-cache-gib", type=float, default=1.0) + parser.add_argument("--draft-kv-window", type=int, default=32768) + parser.add_argument("--proposal-address", default="tcp://127.0.0.1:8092") + parser.add_argument("--disable-proposal-transport", action="store_true") + parser.add_argument("--enable-cuda-graph", action="store_true") + parser.add_argument("--cuda-graph-warmups", type=int, default=2) + parser.add_argument("--skip-smoke-test", action="store_true") + parser.add_argument("--exit-after-load", action="store_true") + args = parser.parse_args() + if ( + args.max_retained_requests is not None + and args.max_retained_requests < args.max_num_seqs + ): + parser.error("--max-retained-requests must be >= --max-num-seqs") + return args + + +def main() -> None: + args = _parse_args() + status = RuntimeStatus( + draft_model=str(args.draft_model), + method=args.method, + target_weights=str(args.target_weights), + num_speculative_tokens=args.num_speculative_tokens, + max_model_len=args.max_model_len, + draft_kv_window=args.draft_kv_window, + ) + try: + runtime = _load_runtime(args, status) + if not args.skip_smoke_test: + status.phase = "eager_smoke_test" + status.smoke_token = _run_eager_smoke( + runtime, torch.device("cuda", args.device) + ) + status.draft_kv_cache_smoke = True + stop = threading.Event() + proposal_engine = None + proposal_server = None + if not args.disable_proposal_transport: + status.phase = "starting_proposal_transport" + from vllm.entrypoints.k3_dspark_rpc import ( + K3DSparkDraftEngine, + K3DSparkZMQServer, + ) + + proposal_engine = K3DSparkDraftEngine( + runtime, + max_requests=( + args.max_retained_requests + if args.max_retained_requests is not None + else args.max_num_seqs + ), + window_size=args.draft_kv_window, + device=torch.device("cuda", args.device), + ) + if args.enable_cuda_graph: + status.phase = f"capturing_{args.method}_cuda_graphs" + proposal_engine.capture_cuda_graphs(warmups=args.cuda_graph_warmups) + proposal_server = K3DSparkZMQServer( + proposal_engine, + address=args.proposal_address, + stop=stop, + ) + proposal_server.start() + status.proposal_transport_ready = True + status.proposal_address = args.proposal_address + status.phase = "ready" + else: + status.phase = "ready_without_transport" + status.ready = True + status.allocated_gib = torch.cuda.memory_allocated(args.device) / 1024**3 + status.reserved_gib = torch.cuda.memory_reserved(args.device) / 1024**3 + logger.info("Standalone K3 draft is loaded: %s", json.dumps(asdict(status))) + if args.exit_after_load: + print(json.dumps(asdict(status), sort_keys=True), flush=True) + return + _serve_status(args, status, stop, proposal_engine) + if proposal_server is not None: + proposal_server.join() + except Exception as exc: + status.phase = "failed" + status.error = f"{type(exc).__name__}: {exc}" + logger.exception("Standalone K3 draft startup failed") + raise + + +if __name__ == "__main__": + main() diff --git a/vllm/v1/worker/gpu/buffer_utils.py b/vllm/v1/worker/gpu/buffer_utils.py index f8336fa0749d..ff4cdde87bc7 100644 --- a/vllm/v1/worker/gpu/buffer_utils.py +++ b/vllm/v1/worker/gpu/buffer_utils.py @@ -152,6 +152,12 @@ def __init__( self.write_starts = new_buffer(self.num_rows, dtype=torch.int32) self.write_cu_lens = new_buffer(self.num_rows, dtype=torch.int32) + @property + def cpu(self) -> torch.Tensor | None: + """Return the host backing tensor when this tensor uses UVA.""" + uva_buf = getattr(self, "_uva_buf", None) + return None if uva_buf is None else uva_buf.cpu + def stage_write( self, index: int, start: int, x: Iterable[int] | Iterable[float] ) -> None: diff --git a/vllm/v1/worker/gpu/input_batch.py b/vllm/v1/worker/gpu/input_batch.py index 9bd83781583b..d1bedf75a79e 100644 --- a/vllm/v1/worker/gpu/input_batch.py +++ b/vllm/v1/worker/gpu/input_batch.py @@ -105,6 +105,11 @@ class InputBatch: max_req_tokens: int | None = None valid_num_draft_tokens_per_req: np.ndarray | None = None + # Optional host view of the request token table. Remote speculators use + # this to verify that a target prefix-cache hit belongs to retained draft + # state before reconnecting it. The tensor is shared, not copied. + all_token_ids_cpu: torch.Tensor | None = None + # When > 0, dummy batches carry seeded-random token ids instead of zeros. # All-zero ids embed identically, so every dummy token routes to the SAME # MoE experts — grouped expert reads collapse and the profiled cost of a diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 24e5bfbfe4c1..28f782620e2c 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1558,6 +1558,7 @@ def prepare_inputs( prompt_lens=prompt_lens, max_req_tokens=max_req_tokens, valid_num_draft_tokens_per_req=valid_num_draft_tokens_per_req, + all_token_ids_cpu=self.req_states.all_token_ids.cpu, ) # InputBuffers are reused across real, dummy, and captured batches. # Clear stale padding before a capacity manager optionally marks a diff --git a/vllm/v1/worker/gpu/spec_decode/__init__.py b/vllm/v1/worker/gpu/spec_decode/__init__.py index c70f169f7be6..767825145dba 100644 --- a/vllm/v1/worker/gpu/spec_decode/__init__.py +++ b/vllm/v1/worker/gpu/spec_decode/__init__.py @@ -1,5 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import os + import torch from vllm.config import VllmConfig @@ -9,12 +11,36 @@ def init_speculator(vllm_config: VllmConfig, device: torch.device): speculative_config = vllm_config.speculative_config assert speculative_config is not None if speculative_config.method == "dflash": + remote_address = os.environ.get("VLLM_K3_DRAFT_REMOTE_ADDRESS") + if remote_address: + from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import ( + RemoteK3DSparkSpeculator, + ) + + return RemoteK3DSparkSpeculator( + vllm_config, + device, + address=remote_address, + ) from vllm.v1.worker.gpu.spec_decode.dflash.speculator import ( DFlashSpeculator, ) return DFlashSpeculator(vllm_config, device) elif speculative_config.method == "dspark": + remote_address = os.environ.get( + "VLLM_K3_DRAFT_REMOTE_ADDRESS" + ) or os.environ.get("VLLM_K3_DSPARK_REMOTE_ADDRESS") + if remote_address: + from vllm.v1.worker.gpu.spec_decode.dspark.remote_speculator import ( + RemoteK3DSparkSpeculator, + ) + + return RemoteK3DSparkSpeculator( + vllm_config, + device, + address=remote_address, + ) from vllm.v1.worker.gpu.spec_decode.dspark.speculator import ( DSparkSpeculator, ) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py new file mode 100644 index 000000000000..354e4815fdd6 --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py @@ -0,0 +1,770 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Verifier-side proxy for a dedicated RTX 3090 K3 draft process.""" + +from __future__ import annotations + +import json +import os +import time +from dataclasses import dataclass +from typing import Any + +import torch +import zmq + +from vllm.config import VllmConfig +from vllm.distributed import get_tp_group +from vllm.logger import init_logger +from vllm.v1.worker.gpu.input_batch import InputBatch +from vllm.v1.worker.gpu.spec_decode.eagle.eagle3_utils import ( + get_eagle3_aux_layers_from_config, +) +from vllm.v1.worker.gpu.spec_decode.speculator import ( + BaseSpeculator, + CUDAGraphCapturePhase, +) + +logger = init_logger(__name__) + +PROTOCOL_VERSION = 2 + + +@dataclass +class _RetainedRequestPrefix: + token_ids: torch.Tensor + committed_end: int + context_start: int + serial: int + + +def _build_valid_context_plan( + input_batch: InputBatch, + rejected_counts: list[int], +) -> tuple[list[int], list[int]]: + """Return valid row indices and per-request counts.""" + if len(rejected_counts) != input_batch.num_reqs: + raise ValueError("Rejected-token count does not match the request batch") + gather_indices: list[int] = [] + valid_counts: list[int] = [] + offset = 0 + for request_idx, (scheduled, rejected) in enumerate( + zip(input_batch.num_scheduled_tokens.tolist(), rejected_counts) + ): + valid = int(scheduled) - int(rejected) + if not 0 <= valid <= int(scheduled): + raise ValueError( + f"Invalid valid-context length for request {request_idx}: " + f"scheduled={scheduled}, rejected={rejected}" + ) + gather_indices.extend(range(offset, offset + valid)) + valid_counts.append(valid) + offset += int(scheduled) + return gather_indices, valid_counts + + +def _anchor_positions_from_context( + context_counts: list[int], context_positions: torch.Tensor +) -> list[int]: + """Return the position immediately following each request's context.""" + anchors: list[int] = [] + offset = 0 + for count in context_counts: + if count <= 0: + raise ValueError("Every remote draft request requires context rows") + offset += count + anchors.append(int(context_positions[offset - 1]) + 1) + if offset != context_positions.numel(): + raise ValueError("Remote draft context counts do not match the position tensor") + return anchors + + +def _contiguous_draft_output( + draft_tokens: torch.Tensor, + num_reqs: int, + num_speculative_tokens: int, +) -> torch.Tensor: + """Return the active TP-broadcast region with a compact row stride.""" + return draft_tokens[:num_reqs, :num_speculative_tokens].contiguous() + + +class RemoteK3DSparkSpeculator(BaseSpeculator): + """Forward target auxiliary states to a standalone greedy draft server.""" + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + *, + address: str, + ) -> None: + self.vllm_config = vllm_config + self.device = device + self.speculative_config = vllm_config.speculative_config + assert self.speculative_config is not None + self.method = str(self.speculative_config.method) + if self.method not in ("dspark", "dflash"): + raise ValueError(f"Unsupported remote K3 draft method: {self.method}") + if self.speculative_config.draft_sample_method != "greedy": + raise ValueError("Remote K3 draft currently supports greedy drafting only") + if self.speculative_config.rejection_sample_method != "block": + raise ValueError( + "Remote K3 draft currently requires block rejection sampling" + ) + if vllm_config.model_config.dtype != torch.bfloat16: + raise ValueError("Remote K3 DSpark transport currently requires BF16") + self.num_speculative_steps = int(self.speculative_config.num_speculative_tokens) + self.max_num_reqs = int(vllm_config.scheduler_config.max_num_seqs) + self.max_num_tokens = int(vllm_config.scheduler_config.max_num_batched_tokens) + draft_hf_config = self.speculative_config.draft_model_config.hf_config + aux_layers = get_eagle3_aux_layers_from_config(self.speculative_config) + if not aux_layers: + raise ValueError( + f"Remote K3 {self.method} config does not declare auxiliary layers" + ) + self.num_aux_layers = len(aux_layers) + target_hidden_size = int( + getattr(draft_hf_config, "target_hidden_size", None) + or draft_hf_config.hidden_size + ) + self.raw_context_width = int(target_hidden_size * self.num_aux_layers) + self.address = address + self.timeout_ms = int( + os.environ.get( + "VLLM_K3_DRAFT_REMOTE_TIMEOUT_MS", + os.environ.get("VLLM_K3_DSPARK_REMOTE_TIMEOUT_MS", "30000"), + ) + ) + self.supports_mm_inputs = False + self.draft_logits: torch.Tensor | None = None + self.draft_tokens = torch.full( + (self.max_num_reqs, self.num_speculative_steps), + -1, + dtype=torch.int64, + device=device, + ) + self._known_requests: set[str] = set() + self._disabled_requests: set[str] = set() + self._active_requests: set[str] = set() + self._retained_prefixes: dict[str, _RetainedRequestPrefix] = {} + self._retained_serial = 0 + self._remote_max_requests = self.max_num_reqs + self._remote_block_size = 1 + self._remote_window_size = 0 + self._remote_prefix_cache_tokens = 0 + self._timing_log_interval = int( + os.environ.get("VLLM_K3_DRAFT_TIMING_LOG_INTERVAL", "0") + ) + if self._timing_log_interval < 0: + raise ValueError("VLLM_K3_DRAFT_TIMING_LOG_INTERVAL must be >= 0") + self._timing_count = 0 + self._timing_totals_ms: dict[str, float] = {} + + tp_group = get_tp_group() + self._tp_group = tp_group + self._tp_rank = int(tp_group.rank_in_group) + self._zmq_context: zmq.Context | None = None + self._socket: zmq.Socket | None = None + if self._tp_rank == 0: + # P2P is unavailable on the target host. Keep the mandatory D2H + # hop off pageable memory so it can be queued directly after the + # target forward on the current stream. + self._positions_staging = torch.empty( + self.max_num_tokens, + dtype=torch.int64, + pin_memory=True, + ) + self._context_staging = torch.empty( + (self.max_num_tokens, self.raw_context_width), + dtype=vllm_config.model_config.dtype, + pin_memory=True, + ) + self._rejected_staging = torch.empty( + self.max_num_reqs, + dtype=torch.int32, + pin_memory=True, + ) + self._anchor_staging = torch.empty( + self.max_num_reqs, + dtype=torch.int64, + pin_memory=True, + ) + self._zmq_context = zmq.Context() + self._connect() + response = self._rpc( + [json.dumps({"protocol": PROTOCOL_VERSION, "op": "PING"}).encode()] + ) + if response.get("op") != "PONG": + raise RuntimeError(f"Unexpected K3 draft health response: {response}") + if response.get("method") != self.method: + raise RuntimeError( + "Remote K3 draft method mismatch: " + f"target={self.method}, server={response.get('method')}" + ) + self._remote_max_requests = int( + response.get("max_requests", self.max_num_reqs) + ) + self._remote_block_size = int(response.get("block_size", 1)) + self._remote_window_size = int(response.get("window_size", 0)) + self._remote_prefix_cache_tokens = int( + response.get("prefix_cache_tokens", 0) + ) + if int(response.get("active_requests", 0)): + # A restarted verifier cannot safely identify state left by an + # older process, so establish a clean protocol epoch. + self._rpc( + [json.dumps({"protocol": PROTOCOL_VERSION, "op": "CLEAR"}).encode()] + ) + + logger.info( + "Remote K3 %s proxy initialized: address=%s, TP rank=%d, K=%d", + self.method, + address, + self._tp_rank, + self.num_speculative_steps, + ) + + def _connect(self) -> None: + assert self._zmq_context is not None + if self._socket is not None: + self._socket.close() + socket = self._zmq_context.socket(zmq.REQ) + socket.setsockopt(zmq.LINGER, 0) + socket.setsockopt(zmq.SNDTIMEO, self.timeout_ms) + socket.setsockopt(zmq.RCVTIMEO, self.timeout_ms) + socket.connect(self.address) + self._socket = socket + + def _rpc(self, frames: list[bytes]) -> dict[str, Any]: + assert self._socket is not None + try: + self._socket.send_multipart(frames) + response = self._socket.recv_json() + except Exception: + self._connect() + raise + if not isinstance(response, dict) or not response.get("ok", False): + raise RuntimeError(f"K3 DSpark RPC failed: {response}") + if int(response.get("protocol", -1)) != PROTOCOL_VERSION: + raise RuntimeError(f"K3 DSpark protocol mismatch: {response}") + return response + + def init_cudagraph_manager(self, cudagraph_mode=None) -> None: + """The standalone drafter owns its CUDA graph lifecycle.""" + + def capture(self, *, capture_phase: CUDAGraphCapturePhase) -> None: + """The verifier has no local draft graph to capture.""" + + def _free_remote_requests(self, request_ids: set[str] | list[str]) -> None: + remote_request_ids = sorted(set(request_ids) & self._known_requests) + if not remote_request_ids: + return + self._rpc( + [ + json.dumps( + { + "protocol": PROTOCOL_VERSION, + "op": "FREE", + "request_ids": remote_request_ids, + } + ).encode() + ] + ) + self._known_requests.difference_update(remote_request_ids) + for request_id in remote_request_ids: + self._retained_prefixes.pop(request_id, None) + + def _ensure_remote_capacity(self, current_request_ids: set[str]) -> None: + while len(self._known_requests) >= self._remote_max_requests: + candidates = self._known_requests - current_request_ids + if not candidates: + raise RuntimeError( + "Remote DSpark request capacity is exhausted by active requests" + ) + request_id = min( + candidates, + key=lambda req_id: ( + self._retained_prefixes[req_id].serial + if req_id in self._retained_prefixes + else -1 + ), + ) + self._free_remote_requests({request_id}) + + @staticmethod + def _token_prefix( + input_batch: InputBatch, + request_idx: int, + prefix_end: int, + ) -> torch.Tensor | None: + token_table = input_batch.all_token_ids_cpu + if token_table is None or prefix_end < 0: + return None + state_idx = int(input_batch.idx_mapping_np[request_idx]) + if state_idx < 0 or prefix_end > token_table.shape[1]: + return None + return token_table[state_idx, :prefix_end] + + def _can_restore_prefix( + self, + retained: _RetainedRequestPrefix, + prefix_end: int, + ) -> bool: + if ( + prefix_end <= 0 + or retained.committed_end < prefix_end + or self._remote_window_size <= 0 + or self._remote_prefix_cache_tokens < self._remote_window_size + ): + return False + restore_start = max(0, prefix_end - self._remote_window_size) + restore_start = ( + restore_start // self._remote_block_size * self._remote_block_size + ) + retained_start = max( + retained.context_start, + retained.committed_end - self._remote_prefix_cache_tokens, + ) + return restore_start >= retained_start + + def _find_reconnect_source( + self, + token_prefix: torch.Tensor, + prefix_end: int, + current_request_ids: set[str], + ) -> str | None: + candidates: list[tuple[int, int, str]] = [] + for request_id in self._known_requests - current_request_ids: + retained = self._retained_prefixes.get(request_id) + if retained is None or not self._can_restore_prefix(retained, prefix_end): + continue + if torch.equal(retained.token_ids[:prefix_end], token_prefix): + candidates.append( + ( + retained.committed_end - prefix_end, + -retained.serial, + request_id, + ) + ) + return min(candidates)[2] if candidates else None + + def _reconnect_request( + self, + source_request_id: str, + request_id: str, + prefix_end: int, + token_prefix: torch.Tensor, + ) -> bool: + try: + response = self._rpc( + [ + json.dumps( + { + "protocol": PROTOCOL_VERSION, + "op": "RECONNECT", + "source_request_id": source_request_id, + "request_id": request_id, + "prefix_end": prefix_end, + } + ).encode() + ] + ) + except Exception: + logger.exception( + "Remote K3 DSpark prefix reconnect failed: source=%s, " + "request=%s, prefix_end=%d", + source_request_id, + request_id, + prefix_end, + ) + return False + + self._retained_prefixes.pop(source_request_id) + self._known_requests.discard(source_request_id) + self._known_requests.add(request_id) + self._retained_serial += 1 + self._retained_prefixes[request_id] = _RetainedRequestPrefix( + token_ids=token_prefix.clone(), + committed_end=prefix_end, + context_start=int(response.get("restored_start", 0)), + serial=self._retained_serial, + ) + logger.info( + "Remote K3 DSpark prefix reconnected: source=%s, request=%s, " + "prefix_end=%d, restored_start=%s, latency_ms=%.1f", + source_request_id, + request_id, + prefix_end, + response.get("restored_start"), + float(response.get("latency_ms", 0.0)), + ) + return True + + def _remember_prefix( + self, + input_batch: InputBatch, + request_idx: int, + request_id: str, + committed_end: int, + context_start: int | None = None, + ) -> None: + token_prefix = self._token_prefix(input_batch, request_idx, committed_end) + if token_prefix is None: + return + previous = self._retained_prefixes.get(request_id) + if context_start is None: + context_start = previous.context_start if previous is not None else 0 + self._retained_serial += 1 + self._retained_prefixes[request_id] = _RetainedRequestPrefix( + token_ids=token_prefix.clone(), + committed_end=committed_end, + context_start=context_start, + serial=self._retained_serial, + ) + + def _copy_tokens_from_response( + self, + response: dict[str, Any], + active_indices: list[int], + num_speculative_tokens: int, + ) -> None: + tokens = response.get("tokens") + expected_shape = (len(active_indices), num_speculative_tokens) + if ( + not isinstance(tokens, list) + or len(tokens) != expected_shape[0] + or any( + not isinstance(row, list) or len(row) != expected_shape[1] + for row in tokens + ) + ): + raise ValueError( + f"Remote DSpark token response has the wrong shape; " + f"expected={expected_shape}, got={tokens!r}" + ) + remote_tokens = torch.tensor(tokens, dtype=torch.int64, device=self.device) + active_gpu = torch.tensor(active_indices, dtype=torch.int64, device=self.device) + # ``draft_tokens`` is allocated at the configured maximum depth, while + # adaptive speculation and the per-batch schedule can request a + # smaller depth for an individual step. Copy into the matching width + # instead of requiring every response to have the maximum width. + self.draft_tokens[:, :num_speculative_tokens].index_copy_( + 0, active_gpu, remote_tokens + ) + + def _record_timing(self, timing_ms: dict[str, float]) -> None: + if self._timing_log_interval <= 0: + return + self._timing_count += 1 + for key, value in timing_ms.items(): + self._timing_totals_ms[key] = self._timing_totals_ms.get(key, 0.0) + value + if self._timing_count < self._timing_log_interval: + return + means = { + key: value / self._timing_count + for key, value in self._timing_totals_ms.items() + } + logger.info( + "Remote K3 %s timing over %d proposals (ms): %s", + self.method, + self._timing_count, + ", ".join(f"{key}={value:.3f}" for key, value in means.items()), + ) + self._timing_count = 0 + self._timing_totals_ms.clear() + + def _rank0_propose( + self, + input_batch: InputBatch, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + num_speculative_tokens: int, + ) -> None: + started = time.perf_counter() + if aux_hidden_states is None or len(aux_hidden_states) != self.num_aux_layers: + raise ValueError( + f"Remote K3 {self.method} requires {self.num_aux_layers} configured " + "target auxiliary hidden states" + ) + + num_reqs = input_batch.num_reqs + current_request_ids = set(input_batch.req_ids) + previous_active_requests = self._active_requests + self._disabled_requests.intersection_update(current_request_ids) + idx_mapping = input_batch.idx_mapping[:num_reqs].long() + sampled_counts = num_sampled[:num_reqs] + sampled_anchors = last_sampled[idx_mapping, 0] + prefill_anchors = next_prefill_tokens[0, idx_mapping] + anchor_tokens = torch.where( + sampled_counts > 0, + sampled_anchors, + prefill_anchors, + ).to(torch.int64) + # Queue both small D2H copies and wait once. The same stream owns the + # preceding token-table update, so this synchronization also makes the + # UVA-backed table safe for prefix matching below. + rejected_staging = self._rejected_staging[:num_reqs] + anchor_staging = self._anchor_staging[:num_reqs] + rejected_staging.copy_(num_rejected[:num_reqs], non_blocking=True) + anchor_staging.copy_(anchor_tokens, non_blocking=True) + torch.cuda.current_stream(self.device).synchronize() + rejected_counts = rejected_staging.tolist() + anchor_tokens_cpu = anchor_staging.tolist() + gather_indices, valid_counts = _build_valid_context_plan( + input_batch, rejected_counts + ) + metadata_ready = time.perf_counter() + + active_indices: list[int] = [] + requests: list[dict[str, Any]] = [] + request_context_starts: list[int | None] = [] + selected_gather_indices: list[int] = [] + gather_offset = 0 + for request_idx, request_id in enumerate(input_batch.req_ids): + valid_count = valid_counts[request_idx] + request_gather = gather_indices[gather_offset : gather_offset + valid_count] + gather_offset += valid_count + if request_id in self._disabled_requests or valid_count <= 0: + continue + first_position = int(input_batch.num_computed_tokens_np[request_idx]) + is_continuing = ( + request_id in previous_active_requests + and request_id in self._known_requests + ) + reset = False + context_start: int | None = None + if not is_continuing: + if first_position == 0: + if request_id in self._known_requests: + self._free_remote_requests({request_id}) + self._ensure_remote_capacity(current_request_ids) + reset = True + context_start = 0 + else: + token_prefix = self._token_prefix( + input_batch, + request_idx, + first_position, + ) + source_request_id: str | None = None + if token_prefix is not None: + retained = self._retained_prefixes.get(request_id) + if ( + request_id in self._known_requests + and retained is not None + and self._can_restore_prefix(retained, first_position) + and torch.equal( + retained.token_ids[:first_position], token_prefix + ) + ): + source_request_id = request_id + else: + source_request_id = self._find_reconnect_source( + token_prefix, + first_position, + current_request_ids, + ) + if source_request_id is None or token_prefix is None: + if request_id in self._known_requests: + self._free_remote_requests({request_id}) + self._ensure_remote_capacity(current_request_ids) + reset = True + context_start = first_position + logger.warning( + "Remote K3 %s cold-bootstrapping cache-restored " + "request %s at position %d from %d fresh context " + "rows; target verification preserves correctness.", + self.method, + request_id, + first_position, + valid_count, + ) + elif not self._reconnect_request( + source_request_id, + request_id, + first_position, + token_prefix, + ): + self._disabled_requests.add(request_id) + continue + requests.append( + { + "request_id": request_id, + "reset": reset, + "reset_position": first_position if reset else 0, + "context_count": valid_count, + "anchor_token_id": int(anchor_tokens_cpu[request_idx]), + } + ) + self._known_requests.add(request_id) + active_indices.append(request_idx) + request_context_starts.append(context_start) + selected_gather_indices.extend(request_gather) + + self._active_requests = self._known_requests & current_request_ids + if not active_indices: + return + requests_ready = time.perf_counter() + + indices_gpu = torch.tensor( + selected_gather_indices, dtype=torch.int64, device=self.device + ) + positions = input_batch.positions.index_select(0, indices_gpu) + context = torch.cat( + [hidden.index_select(0, indices_gpu) for hidden in aux_hidden_states], + dim=-1, + ) + num_context_rows = int(context.shape[0]) + if num_context_rows > self.max_num_tokens: + raise ValueError( + f"Remote DSpark context has {num_context_rows} rows, max is " + f"{self.max_num_tokens}" + ) + if context.shape[1] != self.raw_context_width: + raise ValueError( + f"Remote DSpark context width is {context.shape[1]}, expected " + f"{self.raw_context_width}" + ) + context_ready = time.perf_counter() + positions_staging = self._positions_staging[:num_context_rows] + context_staging = self._context_staging[:num_context_rows] + positions_staging.copy_(positions, non_blocking=True) + context_staging.copy_(context, non_blocking=True) + torch.cuda.current_stream(self.device).synchronize() + context_copied = time.perf_counter() + anchor_positions = _anchor_positions_from_context( + [int(request["context_count"]) for request in requests], + positions_staging, + ) + for request, anchor_position in zip(requests, anchor_positions): + request["anchor_position"] = anchor_position + positions_frame = positions_staging.numpy().tobytes() + # NumPy has inconsistent bfloat16 support; preserve its exact bits as u16. + context_frame = context_staging.view(torch.uint16).numpy().tobytes() + serialized = time.perf_counter() + header = { + "protocol": PROTOCOL_VERSION, + "op": "PROPOSE", + "projected": False, + "num_speculative_tokens": num_speculative_tokens, + "requests": requests, + } + response = self._rpc( + [json.dumps(header).encode(), positions_frame, context_frame] + ) + rpc_done = time.perf_counter() + self._copy_tokens_from_response( + response, + active_indices, + num_speculative_tokens, + ) + output_copied = time.perf_counter() + timing_ms = { + "metadata_d2h": (metadata_ready - started) * 1000, + "request_plan": (requests_ready - metadata_ready) * 1000, + "context_gather": (context_ready - requests_ready) * 1000, + "context_d2h": (context_copied - context_ready) * 1000, + "serialize": (serialized - context_copied) * 1000, + "rpc_roundtrip": (rpc_done - serialized) * 1000, + "tokens_h2d": (output_copied - rpc_done) * 1000, + "client_total": (output_copied - started) * 1000, + } + server_timing = response.get("timing_ms") + if isinstance(server_timing, dict): + for key, value in server_timing.items(): + if isinstance(value, int | float): + timing_ms[f"server_{key}"] = float(value) + self._record_timing(timing_ms) + for request_idx, request, anchor_position, context_start in zip( + active_indices, + requests, + anchor_positions, + request_context_starts, + ): + self._remember_prefix( + input_batch, + request_idx, + str(request["request_id"]), + anchor_position, + context_start, + ) + + @torch.inference_mode() + def propose( + self, + input_batch: InputBatch, + attn_metadata: dict[str, Any], + slot_mappings: dict[str, torch.Tensor], + last_hidden_states: torch.Tensor, + aux_hidden_states: list[torch.Tensor] | None, + num_sampled: torch.Tensor, + num_rejected: torch.Tensor, + last_sampled: torch.Tensor, + next_prefill_tokens: torch.Tensor, + temperature: torch.Tensor, + seeds: torch.Tensor, + num_speculative_tokens: int | None = None, + num_tokens_across_dp: torch.Tensor | None = None, + dummy_run: bool = False, + skip_attn_for_dummy_run: bool = False, + mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, + is_profile: bool = False, + ) -> torch.Tensor: + del ( + attn_metadata, + slot_mappings, + last_hidden_states, + temperature, + seeds, + num_tokens_across_dp, + skip_attn_for_dummy_run, + mm_inputs, + ) + active_k = ( + int(num_speculative_tokens) + if num_speculative_tokens is not None + else self.num_speculative_steps + ) + if not 1 <= active_k <= self.num_speculative_steps: + raise ValueError( + f"Remote DSpark depth must be in [1, " + f"{self.num_speculative_steps}], got {active_k}" + ) + output = self.draft_tokens[: input_batch.num_reqs, :active_k] + output.fill_(-1) + if self._tp_rank == 0 and not (dummy_run or is_profile): + try: + self._rank0_propose( + input_batch, + aux_hidden_states, + num_sampled, + num_rejected, + last_sampled, + next_prefill_tokens, + active_k, + ) + except Exception: + output.fill_(-1) + # The verifier cannot know whether a timed-out request mutated + # remote KV. Fail closed for those requests until they leave + # the active batch; FREE remains safe even if the server never + # created the state. + self._disabled_requests.update(input_batch.req_ids) + logger.exception( + "Remote K3 DSpark proposal failed; drafting is disabled for " + "this step" + ) + # Slicing the active depth from the max-width persistent buffer leaves + # a larger row stride whenever adaptive K is below the configured + # maximum. NCCL broadcast requires a contiguous tensor. Materialize + # only the tiny [batch, K] result after rank 0 has populated it. + output = _contiguous_draft_output( + self.draft_tokens, + input_batch.num_reqs, + active_k, + ) + self._tp_group.broadcast(output, src=0) + return output From 23f9f27f575877ccdb60017148645da0c03da8fc Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Sun, 23 Aug 2026 05:56:47 +0900 Subject: [PATCH 35/52] fix(spec_decode): allow scheduler-selected zero draft depth Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com> --- .../test_acceptance_length_controller.py | 16 ++++++++++++ .../test_k3_dspark_remote_speculator.py | 26 +++++++++++++++++++ .../spec_decode/dspark/remote_speculator.py | 5 ++++ vllm/v1/worker/gpu/spec_decode/utils.py | 2 ++ 4 files changed, 49 insertions(+) diff --git a/tests/v1/spec_decode/test_acceptance_length_controller.py b/tests/v1/spec_decode/test_acceptance_length_controller.py index 9d0cc654a788..2e12b5a7c151 100644 --- a/tests/v1/spec_decode/test_acceptance_length_controller.py +++ b/tests/v1/spec_decode/test_acceptance_length_controller.py @@ -260,6 +260,22 @@ def test_runner_v2_limits_drafts_to_adaptive_depth(): ) +def test_runner_v2_allows_scheduler_to_disable_speculation(): + draft_tokens = torch.tensor([[1, 2, 3], [4, 5, 6]]) + + limited = limit_draft_tokens( + draft_tokens, + num_speculative_tokens=0, + max_num_speculative_tokens=3, + ) + + assert limited.shape == (2, 0) + assert ( + limited.untyped_storage().data_ptr() + == draft_tokens.untyped_storage().data_ptr() + ) + + def test_synthetic_scheduler_output_uses_default_speculative_depth(): output = SchedulerOutput.make_empty() diff --git a/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py b/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py index dcfe1da8fd1c..90328357cce7 100644 --- a/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py +++ b/tests/v1/spec_decode/test_k3_dspark_remote_speculator.py @@ -55,6 +55,32 @@ def test_remote_tokens_copy_supports_adaptive_depth(): ] +def test_remote_speculator_accepts_scheduler_selected_zero_depth(): + proxy = RemoteK3DSparkSpeculator.__new__(RemoteK3DSparkSpeculator) + proxy.num_speculative_steps = 3 + proxy.draft_tokens = torch.full((4, 3), -1, dtype=torch.int64) + batch = SimpleNamespace(num_reqs=2) + empty = torch.empty(0) + + output = proxy.propose( + batch, + {}, + {}, + empty, + None, + empty, + empty, + empty, + empty, + empty, + empty, + num_speculative_tokens=0, + ) + + assert output.shape == (2, 0) + assert output.is_contiguous() + + def test_adaptive_depth_output_is_contiguous_for_tp_broadcast(): draft_tokens = torch.arange(24, dtype=torch.int64).view(3, 8) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py index 354e4815fdd6..876d6ad9eef7 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/remote_speculator.py @@ -728,6 +728,11 @@ def propose( if num_speculative_tokens is not None else self.num_speculative_steps ) + # A scheduler can intentionally disable speculation for one step (for + # example, a request with max_tokens=1). ModelRunner treats an empty + # second dimension as a normal non-speculative step. + if active_k == 0: + return self.draft_tokens[: input_batch.num_reqs, :0].contiguous() if not 1 <= active_k <= self.num_speculative_steps: raise ValueError( f"Remote DSpark depth must be in [1, " diff --git a/vllm/v1/worker/gpu/spec_decode/utils.py b/vllm/v1/worker/gpu/spec_decode/utils.py index 51484005a8ed..9ad96093268d 100644 --- a/vllm/v1/worker/gpu/spec_decode/utils.py +++ b/vllm/v1/worker/gpu/spec_decode/utils.py @@ -34,6 +34,8 @@ def limit_draft_tokens( "Speculator returned unsupported draft shape " f"{tuple(draft_tokens.shape)}; expected a 2D tensor." ) + if num_speculative_tokens == 0: + return draft_tokens[:, :0] if not 1 <= num_speculative_tokens <= max_num_speculative_tokens: raise RuntimeError( "Scheduler selected an invalid speculative-token count " From f49f6c30621bd68781cfd0d8fd5385b592cee1af Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:00:10 +0900 Subject: [PATCH 36/52] fix(cache): avoid target EAGLE drop for disaggregated DSpark Keep DSpark and DFlash scheduling lookahead semantics while applying EAGLE's last-hash target-cache drop only when an actual target KV group is marked as EAGLE. This preserves fine target APC tails for remote/disaggregated drafts and retains the legacy fallback for classic EAGLE. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com> --- .../core/test_dspark_prefix_cache_policy.py | 35 +++++++++++++++++++ vllm/v1/core/sched/scheduler.py | 27 ++++++++++++-- 2 files changed, 60 insertions(+), 2 deletions(-) create mode 100644 tests/v1/core/test_dspark_prefix_cache_policy.py diff --git a/tests/v1/core/test_dspark_prefix_cache_policy.py b/tests/v1/core/test_dspark_prefix_cache_policy.py new file mode 100644 index 000000000000..efda0b307207 --- /dev/null +++ b/tests/v1/core/test_dspark_prefix_cache_policy.py @@ -0,0 +1,35 @@ +from types import SimpleNamespace + +from vllm.v1.core.sched import scheduler as scheduler_module + + +def _spec(method: str, use_eagle: bool = True): + return SimpleNamespace(method=method, use_eagle=lambda: use_eagle) + + +def _groups(*flags: bool): + return [SimpleNamespace(is_eagle_group=flag) for flag in flags] + + +def test_dspark_without_target_eagle_group_does_not_drop_target_cache_tail(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("dspark"), _groups(False, False)) is False + + +def test_explicit_target_eagle_group_keeps_cache_drop(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("dspark"), _groups(False, True)) is True + + +def test_classic_eagle_keeps_legacy_unannotated_fallback(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("eagle3"), _groups(False, False)) is True + + +def test_non_eagle_speculation_never_drops_target_cache_tail(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("ngram", use_eagle=False), _groups(True)) is False diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 56b0ac63a7b2..2fc0a09dc4d8 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -70,6 +70,24 @@ logger = init_logger(__name__) +def use_eagle_for_target_cache( + speculative_config: Any | None, + kv_cache_groups: Iterable[Any], +) -> bool: + """Whether target KV groups require EAGLE's last-hash drop. + + DSpark and DFlash use EAGLE-style scheduling/lookahead but can keep their + draft KV outside the target cache groups. In that layout, falling back from + an empty ``is_eagle_group`` set to all target groups drops one valid target + prefix hash. Classic EAGLE keeps the historical unannotated fallback. + """ + if speculative_config is None or not speculative_config.use_eagle(): + return False + if any(group.is_eagle_group for group in kv_cache_groups): + return True + return speculative_config.method not in ("dspark", "dflash") + + class Scheduler(SchedulerInterface): def __init__( self, @@ -270,6 +288,11 @@ def __init__( ) self.use_eagle = speculative_config.use_eagle() + self.use_eagle_for_target_cache = use_eagle_for_target_cache( + speculative_config, + kv_cache_config.kv_cache_groups, + ) + # Create the KV cache manager. if hash_block_size is None: hash_block_size = block_size @@ -279,7 +302,7 @@ def __init__( max_model_len=self.max_model_len, max_in_flight_tokens=vllm_config.max_in_flight_tokens, enable_caching=self.cache_config.enable_prefix_caching, - use_eagle=self.use_eagle, + use_eagle=self.use_eagle_for_target_cache, log_stats=self.log_stats, enable_kv_cache_events=self.enable_kv_cache_events, dcp_world_size=self.dcp_world_size, @@ -406,7 +429,7 @@ def _mamba_block_aligned_split( # Eagle, FullAttn prunes the last matching block, so back off one # block to avoid a Mamba cache miss. last_cache_position = request.num_tokens - request.num_tokens % block_size - if self.use_eagle: + if getattr(self, "use_eagle_for_target_cache", self.use_eagle): # EAGLE drops the last complete draft-attention block. Convert the # resulting maximum reusable prefix to the recurrent-state grid. # Older configurations without an annotated EAGLE group retain From dbf86def22900a4a2b8f58437d50c50a366f7266 Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Sat, 22 Aug 2026 14:54:15 +0900 Subject: [PATCH 37/52] test(cache): cover DFlash target-cache policy Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com> --- tests/v1/core/test_dspark_prefix_cache_policy.py | 6 ++++++ vllm/v1/core/sched/scheduler.py | 7 +++++++ 2 files changed, 13 insertions(+) diff --git a/tests/v1/core/test_dspark_prefix_cache_policy.py b/tests/v1/core/test_dspark_prefix_cache_policy.py index efda0b307207..482e5ed722ec 100644 --- a/tests/v1/core/test_dspark_prefix_cache_policy.py +++ b/tests/v1/core/test_dspark_prefix_cache_policy.py @@ -17,6 +17,12 @@ def test_dspark_without_target_eagle_group_does_not_drop_target_cache_tail(): assert selector(_spec("dspark"), _groups(False, False)) is False +def test_dflash_without_target_eagle_group_does_not_drop_target_cache_tail(): + selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) + assert selector is not None, "target-cache EAGLE policy selector is missing" + assert selector(_spec("dflash"), _groups(False, False)) is False + + def test_explicit_target_eagle_group_keeps_cache_drop(): selector = getattr(scheduler_module, "use_eagle_for_target_cache", None) assert selector is not None, "target-cache EAGLE policy selector is missing" diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 2fc0a09dc4d8..1095c77f0a18 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -80,6 +80,13 @@ def use_eagle_for_target_cache( draft KV outside the target cache groups. In that layout, falling back from an empty ``is_eagle_group`` set to all target groups drops one valid target prefix hash. Classic EAGLE keeps the historical unannotated fallback. + + Args: + speculative_config: Active speculative-decoding configuration, if any. + kv_cache_groups: Target-cache groups inspected for explicit EAGLE state. + + Returns: + Whether target-cache prefix matching must drop EAGLE's trailing hash. """ if speculative_config is None or not speculative_config.use_eagle(): return False From 1bed2daa1cc5a57345b9757c58aab46094d4ba77 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Wed, 19 Aug 2026 01:15:26 -0400 Subject: [PATCH 38/52] scheduler: skip speculative decoding when all scheduled requests need <=1 output token When a batch contains only new requests (no running ones) and every one has max_tokens <= 1, set num_spec_tokens_to_schedule = 0. Speculative decoding cannot help a 1-token output, so the draft pass and verification are pure overhead. This is the shape of every max_tokens=1 API call, every prefill-throughput benchmark, and every embedding/classification-style request. Measured on RTX 5090 (31.4 GiB), Qwen3.8-27B EXL3, MTP=6: 1-token request latency 141 ms -> 127 ms 2051-token prefill bench 7445 -> 7635 tok/s (+2.5%) TG on normal requests 189.8 tok/s (unchanged) The guard is conservative: it requires scheduled_running_reqs to be empty, so an in-flight multi-token generation can never lose its draft tokens. Signed-off-by: Michel Belleau --- vllm/v1/core/sched/scheduler.py | 13 +++++++++++++ 1 file changed, 13 insertions(+) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index 1095c77f0a18..b0807c231efc 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1292,6 +1292,19 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: scheduled_encoder_inputs ) + # Skip speculative decoding when every scheduled request needs at most + # one output token. Drafting cannot pay off for a 1-token output, so the + # draft pass plus verification is pure added latency. Only fires when no + # running request is scheduled, so an in-flight multi-token generation + # can never lose its draft tokens. + if ( + num_spec_tokens_to_schedule > 0 + and scheduled_new_reqs + and not scheduled_running_reqs + and all(req.max_tokens <= 1 for req in scheduled_new_reqs) + ): + num_spec_tokens_to_schedule = 0 + scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, scheduled_cached_reqs=cached_reqs_data, From 27e71cae6d44f9a9f48b09783ae5954c4077dd12 Mon Sep 17 00:00:00 2001 From: Michel Belleau Date: Thu, 20 Aug 2026 07:18:28 -0400 Subject: [PATCH 39/52] =?UTF-8?q?fix:=20address=20CodeRabbit=20review=20?= =?UTF-8?q?=E2=80=94=20document=20lookahead=20reservation=20is=20harmless?= =?UTF-8?q?=20for=20single-token=20requests?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- vllm/v1/core/sched/scheduler.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index b0807c231efc..e7596ebdd955 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -1304,6 +1304,17 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: and all(req.max_tokens <= 1 for req in scheduled_new_reqs) ): num_spec_tokens_to_schedule = 0 + # NOTE: allocate_slots (called above for each request) already + # reserved KV blocks for self.num_lookahead_tokens speculative + # slots. We cannot retroactively shrink that reservation here + # because it was made per-request during the scheduling loop, + # before this aggregate skip condition is evaluated. The + # reservation is harmless: these are single-token requests + # (max_tokens <= 1) that finish immediately after one decode + # step, so the extra blocks are released at the next scheduling + # iteration. Zeroing num_spec_tokens_to_schedule prevents the + # draft model from running, which is the source of the wasted + # latency we are avoiding. scheduler_output = SchedulerOutput( scheduled_new_reqs=new_reqs_data, From 74ec7af2d0725173ab7f86abec971173a6eb9e28 Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Sat, 22 Aug 2026 20:31:29 +0900 Subject: [PATCH 40/52] fix(comm): prewarm FlashInfer PCIe graph shapes Call the finalized FlashInfer workspace prepare API during vLLM graph warmup so autotune and cache lookup complete before CUDA graph capture. Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com> --- tests/distributed/test_flashinfer_pcie_all_reduce.py | 5 +++++ .../device_communicators/flashinfer_pcie_all_reduce.py | 6 +++++- 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/tests/distributed/test_flashinfer_pcie_all_reduce.py b/tests/distributed/test_flashinfer_pcie_all_reduce.py index 84ee37c2d155..85c3be6a66ef 100644 --- a/tests/distributed/test_flashinfer_pcie_all_reduce.py +++ b/tests/distributed/test_flashinfer_pcie_all_reduce.py @@ -21,6 +21,7 @@ def __init__(self, **kwargs: Any) -> None: self.kwargs = kwargs self.destroyed = False self.last_input: torch.Tensor | None = None + self.prepared: list[tuple[list[tuple[int, ...]], torch.dtype]] = [] FakeWorkspace.instances.append(self) def supports(self, inp: torch.Tensor) -> bool: @@ -35,6 +36,9 @@ def all_reduce( out.copy_(inp) return out + def prepare(self, shapes, *, dtype) -> None: + self.prepared.append((list(shapes), dtype)) + def destroy(self) -> None: self.destroyed = True @@ -101,6 +105,7 @@ def test_capture_routes_graph_calls_without_reusing_eager_state() -> None: assert torch.equal(actual, inp) assert len(FakeWorkspace.instances) == 1 assert FakeWorkspace.instances[0].last_input is inp + assert FakeWorkspace.instances[0].prepared == [([(1, 4)], torch.float32)] pool.close() diff --git a/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py b/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py index 1d3b43082c02..53b7c0a86220 100644 --- a/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py +++ b/vllm/distributed/device_communicators/flashinfer_pcie_all_reduce.py @@ -169,11 +169,15 @@ def prepare_graph_all_reduce( ) -> None: del stream channel_id = self._active_channel_id or self._EAGER_CHANNEL_ID - if not self._workspace(channel_id).supports(inp): + workspace = self._workspace(channel_id) + if not workspace.supports(inp): raise ValueError( "FlashInfer PCIe IPC graph warmup received an unsupported " f"shape {tuple(inp.shape)}" ) + hidden = inp.shape[-1] + batch = inp.numel() // hidden + workspace.prepare([(batch, hidden)], dtype=inp.dtype) def all_reduce( self, From 93917b397ec50e73f53a6aaf95246bcd536c9022 Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Sun, 30 Aug 2026 06:58:43 +0900 Subject: [PATCH 41/52] feat(kquant): opt into W4A8 for coupled QSRT prefill --- tests/quantization/test_kquant_hybrid.py | 11 + .../layers/quantization/kquant_hybrid.py | 253 ++++++++++++++++-- 2 files changed, 245 insertions(+), 19 deletions(-) diff --git a/tests/quantization/test_kquant_hybrid.py b/tests/quantization/test_kquant_hybrid.py index b300408fddcf..187b790c927c 100644 --- a/tests/quantization/test_kquant_hybrid.py +++ b/tests/quantization/test_kquant_hybrid.py @@ -9,6 +9,7 @@ from vllm.model_executor.layers.quantization import get_quantization_config from vllm.model_executor.layers.quantization.kquant_hybrid import ( KQuantHybridConfig, + _qsrt_atoms_v2_w4a8_prefill_enabled, _b12x_tiles_for_geometry, _is_dense_layer_ignored, _read_hybrid_keys, @@ -158,6 +159,16 @@ def test_config_accepts_coupled_pure_k2_atoms_v2_profile() -> None: assert config.qsrt == descriptor +def test_qsrt_w4a8_prefill_is_opt_in_and_pure_k2_only(monkeypatch) -> None: + monkeypatch.delenv("VLLM_KQUANT_QSRT_W4A8_PREFILL", raising=False) + assert not _qsrt_atoms_v2_w4a8_prefill_enabled(pure_k2=True) + + monkeypatch.setenv("VLLM_KQUANT_QSRT_W4A8_PREFILL", "1") + assert _qsrt_atoms_v2_w4a8_prefill_enabled(pure_k2=True) + with pytest.raises(ValueError, match="uniform coupled pure-K2"): + _qsrt_atoms_v2_w4a8_prefill_enabled(pure_k2=False) + + def test_atoms_v2_profiles_have_canonical_row_strides() -> None: assert _atom_slot_stride_for_profile(PROFILE) == 119005184 assert _atom_slot_stride_for_profile(PURE_K2_PROFILE) == 77242368 diff --git a/vllm/model_executor/layers/quantization/kquant_hybrid.py b/vllm/model_executor/layers/quantization/kquant_hybrid.py index bfc365cdb727..5580dbfac052 100644 --- a/vllm/model_executor/layers/quantization/kquant_hybrid.py +++ b/vllm/model_executor/layers/quantization/kquant_hybrid.py @@ -5,8 +5,9 @@ The high-quality expert tier is stored as NVFP4 or MXFP4. The secondary tier is either generic ``exl3_3`` trellis tensors or TP-independent -``qsrt_sqg_e4m3`` atoms. Both secondary formats execute through the B12X -W4A16 trellis path; QSRT's exact endpoint is X4T. +``qsrt_sqg_e4m3`` atoms. Generic EXL3 executes through B12X W4A16; the +uniform coupled-K2 atoms-v2 profile can optionally use B12X W4A8-MX for +prefill while retaining W4A16 for the small decode band. """ import dataclasses @@ -64,6 +65,18 @@ _QSRT_ATOMS_V2_PROFILE_COUPLED_H308 = "k3x22_k4x2_coupled_h512_h128" +def _qsrt_atoms_v2_w4a8_prefill_enabled(*, pure_k2: bool) -> bool: + """Enable W4A8 prefill only for the one profile it supports safely.""" + requested = os.getenv("VLLM_KQUANT_QSRT_W4A8_PREFILL", "0").strip().lower() + enabled = requested in {"1", "true", "yes", "on"} + if enabled and not pure_k2: + raise ValueError( + "VLLM_KQUANT_QSRT_W4A8_PREFILL requires the uniform coupled pure-K2 " + f"profile {_QSRT_ATOMS_V2_PROFILE_COUPLED_K2!r}" + ) + return enabled + + def _stack_exl3_intermediate_rotations( w13_svh: torch.Tensor, w2_suh: torch.Tensor, @@ -178,7 +191,7 @@ def _b12x_tiles_for_geometry( class _HybridSharedRuntime: - """Process-wide b12x W4A16 runtime shared by every hybrid MoE layer. + """Process-wide B12X runtime shared by every hybrid MoE layer. One preplanned-launch cache and one scratch/route buffer set serve all layers: launches on a single stream never overlap and every @@ -197,7 +210,12 @@ def __init__(self) -> None: # H128(h * down_suh); this stable buffer receives its inverse transform. self.kquant_logical_mid: torch.Tensor | None = None self.trellis_scratch: torch.Tensor | None = None + self.trellis_prefill_scratch: torch.Tensor | None = None self.trellis_output: torch.Tensor | None = None + self.trellis_prefill_input: torch.Tensor | None = None + self.trellis_prefill_output: torch.Tensor | None = None + self.trellis_prefill_topk_ids: torch.Tensor | None = None + self.trellis_prefill_topk_weights: torch.Tensor | None = None # X4T expands only routed scale rows immediately before W4A16. The # output grids are shared across layers on the same CUDA stream. self.x4t_w13_scale_scratch: torch.Tensor | None = None @@ -248,6 +266,9 @@ def __init__( self.uses_qsrt_atoms = False self.trellis_weights: Any = None self.trellis_plan: Any = None + self.trellis_prefill_weights: Any = None + self.trellis_prefill_plan: Any = None + self.trellis_use_w4a8_prefill = False self.runtime_ready = False @@ -1017,10 +1038,12 @@ def manifest_file(field: str) -> Path: "QSRT atoms-v2 profile disagrees with the model descriptor" ) pure_k2 = metadata_v2.profile == _QSRT_ATOMS_V2_PROFILE_COUPLED_K2 + state.trellis_use_w4a8_prefill = ( + _qsrt_atoms_v2_w4a8_prefill_enabled(pure_k2=pure_k2) + ) tp_size = get_tensor_model_parallel_world_size() tp_rank = get_tensor_model_parallel_rank() - plan = fused_moe.plan_weights( - quant_modes="w4a16", + weight_plan_kwargs: dict[str, Any] = dict( source_format="qsrt_sqg_e4m3", activation=self.moe.activation.value, params_dtype=self.moe.in_dtype, @@ -1033,6 +1056,10 @@ def manifest_file(field: str) -> Path: qsrt_storage_format="qsrt_atoms_v2", qsrt_profile=metadata_v2.profile, ) + plan = fused_moe.plan_weights( + quant_modes="w4a16", + **weight_plan_kwargs, + ) with open_qsrt_atom_v2_extent( metadata_v2, shard_count=tp_size, @@ -1054,15 +1081,39 @@ def manifest_file(field: str) -> Path: down_svh=metadata_v2.down_svh.unsqueeze(0).to(device), qsrt_rotation_draws=metadata_v2.rotation_draws, ) + if state.trellis_use_w4a8_prefill: + # QSRT preparation currently accepts one quant mode per plan, + # but both execution modes consume the exact same immutable + # trellis-native tensors. Build a second validated package + # that aliases those allocations instead of preparing a + # second ~model-sized representation. + w4a8_plan = fused_moe.plan_weights( + quant_modes="w4a8_mx", + **weight_plan_kwargs, + ) + representation = state.trellis_weights.representation + if representation is None: + raise RuntimeError("QSRT W4A16 representation was not prepared") + state.trellis_prefill_weights = dataclasses.replace( + state.trellis_weights, + plan=w4a8_plan, + representation=dataclasses.replace( + representation, + quant_mode="w4a8_mx", + ), + ) layer.qsrt_atom_placeholder.data = ( layer.qsrt_atom_placeholder.data.new_empty((0,)) ) logger.info( - "Loaded QSRT layer %d atoms-v2 shard %d/%d: %d QSRT experts", + "Loaded QSRT layer %d atoms-v2 shard %d/%d: %d QSRT experts%s", layer_index, tp_rank, tp_size, state.num_secondary, + " using W4A16 decode + W4A8 prefill" + if state.trellis_use_w4a8_prefill + else "", ) return @@ -1487,6 +1538,11 @@ def _ensure_runtime(self, layer: "RoutedExperts", m: int, topk: int) -> None: if getattr(state, "trellis_weights", None) is not None: from b12x.moe import fused_moe + w4a16_max_m = ( + _B12X_DECODE_M + if state.trellis_prefill_weights is not None + else runtime.max_m + ) key = ( "trellis", state.num_secondary, @@ -1496,12 +1552,12 @@ def _ensure_runtime(self, layer: "RoutedExperts", m: int, topk: int) -> None: self.moe.activation.value, state.tiles, runtime.topk, - runtime.max_m, + w4a16_max_m, ) plan = runtime.launches.get(key) if plan is None: caps = fused_moe.Caps( - max_tokens=runtime.max_m, + max_tokens=w4a16_max_m, num_topk=runtime.topk, device=torch.accelerator.current_device_index(), weight_plan=state.trellis_weights.plan, @@ -1525,7 +1581,7 @@ def _ensure_runtime(self, layer: "RoutedExperts", m: int, topk: int) -> None: dtype=spec.dtype, device=spec.device, ) - output_shape = (runtime.max_m, state.hidden_size) + output_shape = (w4a16_max_m, state.hidden_size) if ( runtime.trellis_output is None or tuple(runtime.trellis_output.shape) != output_shape @@ -1536,6 +1592,97 @@ def _ensure_runtime(self, layer: "RoutedExperts", m: int, topk: int) -> None: dtype=torch.float32, device=spec.device, ) + if state.trellis_prefill_weights is not None: + prefill_output_shape = (runtime.max_m, state.hidden_size) + prefill_key = ( + "trellis-w4a8-prefill", + state.num_secondary, + state.hidden_size, + state.intermediate_size, + self.moe.activation.value, + state.tiles, + runtime.topk, + runtime.max_m, + ) + prefill_plan = runtime.launches.get(prefill_key) + if prefill_plan is None: + prefill_plan = fused_moe.plan( + fused_moe.Caps( + max_tokens=runtime.max_m, + num_topk=runtime.topk, + device=torch.accelerator.current_device_index(), + weight_plan=state.trellis_prefill_weights.plan, + quant_mode="w4a8_mx", + route_num_experts=0, + ) + ) + if prefill_plan.launch_plan.implementation != "dynamic": + raise RuntimeError( + "kquant_hybrid: W4A8 prefill plan is not dynamic" + ) + runtime.launches[prefill_key] = prefill_plan + state.trellis_prefill_plan = prefill_plan + prefill_spec = prefill_plan.scratch_specs()[0] + prefill_need = int(torch.Size(prefill_spec.shape).numel()) + prefill_scratch = runtime.trellis_prefill_scratch + if prefill_scratch is None or ( + prefill_scratch.numel() < prefill_need + or prefill_scratch.dtype != prefill_spec.dtype + ): + runtime.trellis_prefill_scratch = torch.empty( + prefill_spec.shape, + dtype=prefill_spec.dtype, + device=prefill_spec.device, + ) + if ( + runtime.trellis_prefill_input is None + or tuple(runtime.trellis_prefill_input.shape) + != prefill_output_shape + or runtime.trellis_prefill_input.device != prefill_spec.device + or runtime.trellis_prefill_input.dtype != self.moe.in_dtype + ): + runtime.trellis_prefill_input = torch.empty( + prefill_output_shape, + dtype=self.moe.in_dtype, + device=prefill_spec.device, + ) + if ( + runtime.trellis_prefill_output is None + or tuple(runtime.trellis_prefill_output.shape) + != prefill_output_shape + or runtime.trellis_prefill_output.device != prefill_spec.device + or runtime.trellis_prefill_output.dtype != self.moe.in_dtype + ): + runtime.trellis_prefill_output = torch.empty( + prefill_output_shape, + dtype=self.moe.in_dtype, + device=prefill_spec.device, + ) + routing_shape = (runtime.max_m, runtime.topk) + if ( + runtime.trellis_prefill_topk_ids is None + or tuple(runtime.trellis_prefill_topk_ids.shape) + != routing_shape + or runtime.trellis_prefill_topk_ids.device + != prefill_spec.device + ): + runtime.trellis_prefill_topk_ids = torch.empty( + routing_shape, + dtype=torch.int32, + device=prefill_spec.device, + ) + if ( + runtime.trellis_prefill_topk_weights is None + or tuple(runtime.trellis_prefill_topk_weights.shape) + != routing_shape + or runtime.trellis_prefill_topk_weights.device + != prefill_spec.device + ): + runtime.trellis_prefill_topk_weights = torch.empty( + routing_shape, + dtype=torch.float32, + device=prefill_spec.device, + ) if ( os.getenv("VLLM_KQUANT_CAPTURE_DIR") and os.getenv("VLLM_KQUANT_CAPTURE_PROFILE", "sampled_hessian") @@ -1770,19 +1917,87 @@ def _apply_once( tids = tids.contiguous() if runtime.trellis_output is None: raise RuntimeError("QSRT trellis output was not allocated eagerly") - binding = fused_moe.bind( - state.trellis_plan, - scratch=runtime.trellis_scratch, - a=x if x.is_contiguous() else x.contiguous(), - experts=state.trellis_weights, + use_w4a8_prefill = ( + state.trellis_prefill_weights is not None and not decode + ) + trellis_plan = ( + state.trellis_prefill_plan + if use_w4a8_prefill + else state.trellis_plan + ) + trellis_weights = ( + state.trellis_prefill_weights + if use_w4a8_prefill + else state.trellis_weights + ) + trellis_scratch = ( + runtime.trellis_prefill_scratch + if use_w4a8_prefill + else runtime.trellis_scratch + ) + trellis_output = ( + runtime.trellis_prefill_output + if use_w4a8_prefill + else runtime.trellis_output + ) + if trellis_scratch is None: + raise RuntimeError("QSRT trellis scratch was not allocated eagerly") + if ( + trellis_plan is None + or trellis_weights is None + or trellis_output is None + ): + raise RuntimeError("QSRT trellis execution resources are incomplete") + launch_input = x if x.is_contiguous() else x.contiguous() + if use_w4a8_prefill: + from b12x.moe.fused_moe._impl import ( + run_w4a8_coupled_outer_transform, + sanitize_w4a8_routing, + ) + + prefill_input = runtime.trellis_prefill_input + safe_ids = runtime.trellis_prefill_topk_ids + safe_weights = runtime.trellis_prefill_topk_weights + prepared_value = trellis_weights.representation.value + if prefill_input is None or safe_ids is None or safe_weights is None: + raise RuntimeError( + "QSRT W4A8 prefill buffers were not allocated eagerly" + ) + launch_input = run_w4a8_coupled_outer_transform( + launch_input, + prefill_input[:m], + prepared_value.gate_suh, + output_transform=False, + ) + tids, weights = sanitize_w4a8_routing( + tids, + weights, + safe_ids[:m], + safe_weights[:m], + num_experts=state.num_secondary, + ) + bind_kwargs: dict[str, Any] = dict( + scratch=trellis_scratch, + a=launch_input, + experts=trellis_weights, topk_weights=weights, topk_ids=tids, - route_expert_map=state.emap_secondary, - output=runtime.trellis_output[:m], + output=trellis_output[:m], ) - # The unified full-rotation top-k sum emits fp32; downstream - # layers expect the model dtype. - out_trellis = fused_moe.run(binding=binding)[:m].to(x.dtype) + if not use_w4a8_prefill: + bind_kwargs["route_expert_map"] = state.emap_secondary + binding = fused_moe.bind(trellis_plan, **bind_kwargs) + # W4A16 decode emits fp32 while W4A8 prefill emits model dtype; + # normalize both contracts for downstream layers. + out_trellis = fused_moe.run(binding=binding)[:m] + if use_w4a8_prefill: + out_trellis = run_w4a8_coupled_outer_transform( + out_trellis, + out_trellis, + prepared_value.down_svh, + output_transform=True, + ) + out_trellis = out_trellis.to(x.dtype) if ( os.getenv("VLLM_KQUANT_CAPTURE_DIR") and os.getenv("VLLM_KQUANT_CAPTURE_PROFILE", "sampled_hessian") From c7fd9c8d200bb5503e6bff0f9d9a76ac128bcf2d Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Tue, 1 Sep 2026 16:55:48 +0900 Subject: [PATCH 42/52] perf(mla): fuse K3 DCP verification queries Share KV loads across fixed K=3 verification rows, select capacity-specific graph plans, and add guarded q-rep and sparse policies. Assisted-by: OpenAI Codex Signed-off-by: myshytf <9619163+myshytf@users.noreply.github.com> --- tests/models/kimi_k3/test_mla_padding.py | 111 ++++++ tests/v1/attention/test_b12x_mla.py | 157 +++++++- vllm/envs.py | 28 ++ vllm/models/kimi_k3/nvidia/mla.py | 114 +++++- vllm/v1/attention/backends/mla/b12x_mla.py | 430 ++++++++++++++------- 5 files changed, 700 insertions(+), 140 deletions(-) diff --git a/tests/models/kimi_k3/test_mla_padding.py b/tests/models/kimi_k3/test_mla_padding.py index 007556801152..cf137217e41a 100644 --- a/tests/models/kimi_k3/test_mla_padding.py +++ b/tests/models/kimi_k3/test_mla_padding.py @@ -88,6 +88,117 @@ def safe_bmm(q, weight, output, *, use_safe_op): torch.testing.assert_close(result, expected.transpose(0, 1)) +def test_kimi_mla_decode_query_uses_replicated_absorbed_weight(monkeypatch): + from vllm.models.kimi_k3.nvidia import mla + + attention = object.__new__(mla.MultiHeadLatentAttention) + torch.nn.Module.__init__(attention) + attention.kv_lora_rank = 5 + attention.dcp_q_replicate = True + attention.W_UK_T = torch.nn.Parameter( + torch.randn((2, 3, 5), dtype=torch.bfloat16), requires_grad=False + ) + attention.W_UK_T_dcp_qrep = torch.randn((4, 3, 5), dtype=torch.bfloat16) + query = torch.randn((3, 4, 3), dtype=torch.bfloat16) + captured = {} + + def safe_bmm(q, weight, output, *, use_safe_op): + captured["weight"] = weight + torch.bmm(q.contiguous(), weight, out=output) + + monkeypatch.setattr(mla, "_run_mla_query_bmm", safe_bmm) + + result = attention._absorb_decode_query(query) + + assert captured["weight"] is attention.W_UK_T_dcp_qrep + expected = torch.bmm( + query.transpose(0, 1).contiguous(), + attention.W_UK_T_dcp_qrep, + ) + torch.testing.assert_close(result, expected.transpose(0, 1)) + + +def test_kimi_mla_qrep_layer_allowlist(monkeypatch): + from vllm.models.kimi_k3.nvidia import mla + + config = SimpleNamespace( + parallel_config=SimpleNamespace( + decode_context_parallel_size=8, + prefill_context_parallel_size=1, + ) + ) + monkeypatch.setattr(mla.envs, "VLLM_DCP_Q_REPLICATE", True) + monkeypatch.setattr( + mla.envs, + "VLLM_K3_DCP_Q_REPLICATE_LAYERS", + "4,12-16,92", + ) + + assert mla._k3_dcp_qrep_enabled("model.layers.4.self_attn", config) + assert mla._k3_dcp_qrep_enabled("model.layers.14.self_attn", config) + assert not mla._k3_dcp_qrep_enabled("model.layers.20.self_attn", config) + assert mla._k3_dcp_qrep_enabled("model.layers.92.self_attn", config) + + +def test_kimi_mla_qrep_requires_explicit_memory_policy(monkeypatch): + from vllm.models.kimi_k3.nvidia import mla + + config = SimpleNamespace( + parallel_config=SimpleNamespace( + decode_context_parallel_size=8, + prefill_context_parallel_size=1, + ) + ) + monkeypatch.setattr(mla.envs, "VLLM_DCP_Q_REPLICATE", True) + monkeypatch.setattr(mla.envs, "VLLM_K3_DCP_Q_REPLICATE_LAYERS", None) + + with pytest.raises(ValueError, match="explicit layer list"): + mla._k3_dcp_qrep_enabled("model.layers.4.self_attn", config) + + +def test_kimi_mla_qrep_explicit_all(monkeypatch): + from vllm.models.kimi_k3.nvidia import mla + + config = SimpleNamespace( + parallel_config=SimpleNamespace( + decode_context_parallel_size=8, + prefill_context_parallel_size=1, + ) + ) + monkeypatch.setattr(mla.envs, "VLLM_DCP_Q_REPLICATE", True) + monkeypatch.setattr(mla.envs, "VLLM_K3_DCP_Q_REPLICATE_LAYERS", "all") + + assert mla._k3_dcp_qrep_enabled("model.layers.4.self_attn", config) + + +def test_kimi_mla_qrep_uses_dcp_group_head_width(): + from vllm.models.kimi_k3.nvidia import mla + + assert mla._k3_projected_query_heads(12, 8, True) == 96 + assert mla._k3_projected_query_heads(6, 8, True) == 48 + assert mla._k3_projected_query_heads(6, 8, False) == 6 + + +def test_kimi_mla_qrep_rejects_invalid_layer_range(monkeypatch): + from vllm.models.kimi_k3.nvidia import mla + + config = SimpleNamespace( + parallel_config=SimpleNamespace( + decode_context_parallel_size=8, + prefill_context_parallel_size=1, + ) + ) + monkeypatch.setattr(mla.envs, "VLLM_DCP_Q_REPLICATE", True) + monkeypatch.setattr( + mla.envs, + "VLLM_K3_DCP_Q_REPLICATE_LAYERS", + "16-12", + ) + + with pytest.raises(ValueError, match="Invalid K3 qrep layer range"): + mla._k3_dcp_qrep_enabled("model.layers.12.self_attn", config) + + def test_kimi_mla_defines_graph_padding_before_output_projection(monkeypatch): from vllm.models.kimi_k3.nvidia import mla diff --git a/tests/v1/attention/test_b12x_mla.py b/tests/v1/attention/test_b12x_mla.py index 8d2dc5d28b74..8f41cac1fbf9 100644 --- a/tests/v1/attention/test_b12x_mla.py +++ b/tests/v1/attention/test_b12x_mla.py @@ -63,6 +63,20 @@ def test_b12x_mla_limits_active_cache_splits( assert b12x_mla._active_dense_mla_splits(plan, max_seq_len) == expected +def test_b12x_mla_row_caps_cover_full_graph_shapes() -> None: + assert b12x_mla._dense_mla_plan_row_caps(28) == (1, 2, 4, 8, 16, 28) + + +def test_b12x_mla_selects_smallest_covering_plan() -> None: + plans = {1: "b1", 2: "b2", 4: "b4", 8: "b8"} + + assert b12x_mla._select_dense_mla_plan(plans, 1) == "b1" + assert b12x_mla._select_dense_mla_plan(plans, 3) == "b4" + assert b12x_mla._select_dense_mla_plan(plans, 8) == "b8" + with pytest.raises(ValueError, match="exceed the planned capacities"): + b12x_mla._select_dense_mla_plan(plans, 9) + + def test_b12x_mla_plans_local_interleaved_dcp_cache() -> None: config = SimpleNamespace( parallel_config=SimpleNamespace( @@ -76,15 +90,15 @@ def test_b12x_mla_plans_local_interleaved_dcp_cache() -> None: def test_mla_uses_one_kv_shard_for_replicated_dcp_cache() -> None: - replicated = SimpleNamespace(dcp_replicated=True, dcp_kv_shard_count=None) - sharded = SimpleNamespace(dcp_replicated=False, dcp_kv_shard_count=None) + replicated = SimpleNamespace(get_num_dcp_kv_shards=lambda _: 1) + sharded = SimpleNamespace(get_num_dcp_kv_shards=lambda dcp_size: dcp_size) assert mla_attention._get_mla_kv_dcp_world_size(replicated, 16) == 1 assert mla_attention._get_mla_kv_dcp_world_size(sharded, 16) == 16 def test_mla_rejects_partial_dcp_cache_without_matching_subgroup() -> None: - partial = SimpleNamespace(dcp_replicated=False, dcp_kv_shard_count=4) + partial = SimpleNamespace(get_num_dcp_kv_shards=lambda _: 4) with pytest.raises(NotImplementedError, match="partial DCP KV sharding"): mla_attention._get_mla_kv_dcp_world_size(partial, 16) @@ -195,6 +209,7 @@ def _fake_impl(monkeypatch, *, num_heads: int = 8) -> tuple[B12xMLAImpl, _FakeDe impl.dcp_world_size = 1 impl._dcp_comm_backend = "a2a" impl._dcp_max_batch_size = 16 + impl.dcp_q_replicate = False impl._compiled_bindings = set() dense_mla = _FakeDenseMLA() impl._dense_mla = dense_mla @@ -450,6 +465,59 @@ def test_b12x_mla_builder_maps_causal_verification_lengths_to_dcp_rank( ) +def test_b12x_mla_builder_preserves_tiled_q4_dcp_verification( + monkeypatch, +) -> None: + builder = object.__new__(B12xMLAMetadataBuilder) + builder._dense_mla_plan = _FakePlan() + builder._dense_mla_plans = {4: _FakePlan()} + verify_plan = SimpleNamespace(caps=SimpleNamespace(max_page_table_width=4)) + builder._dense_mla_verify_plans = {1: verify_plan} + builder._dense_mla_scratch = torch.empty(256, dtype=torch.uint8) + builder._dense_mla_padded_q = None + builder._dense_mla_padded_output = None + builder._max_dense_mla_rows = 8 + builder._dense_mla_flat_block_table = torch.zeros(8, 4, dtype=torch.int32) + builder._dense_mla_flat_seq_lens = torch.empty(8, dtype=torch.int32) + builder._dense_mla_flat_query_start_loc = torch.arange(9, dtype=torch.int32) + builder._dense_mla_causal_offsets = torch.arange(-3, 1, dtype=torch.int32) + builder._dense_mla_flat_global_seq_lens = torch.empty(8, dtype=torch.int32) + builder._dense_mla_flat_dcp_remainder = torch.empty(8, dtype=torch.int32) + builder.dcp_world_size = 4 + builder._dcp_rank = 3 + builder.cp_kv_cache_interleave_size = 2 + + source_table = torch.tensor([[3, 4, 5, 6, 90]], dtype=torch.int32) + metadata = SimpleNamespace( + causal=True, + num_decodes=1, + num_decode_tokens=4, + decode=SimpleNamespace( + block_table=source_table, + seq_lens=torch.tensor([4], dtype=torch.int32), + dcp_tot_seq_lens=torch.tensor([17], dtype=torch.int32), + ), + ) + monkeypatch.setattr( + b12x_mla.MLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata, + ) + + result = builder.build(0, SimpleNamespace()) + + assert result.dense_mla_plan is verify_plan + torch.testing.assert_close( + result.dense_mla_verify_block_table, + source_table[:, :4], + ) + torch.testing.assert_close( + result.dense_mla_query_cache_seq_lens, + torch.tensor([2, 3, 4, 4], dtype=torch.int32), + ) + assert getattr(result, "dense_mla_flat_block_table", None) is None + + def test_b12x_mla_builder_bounds_single_token_draft_table(monkeypatch) -> None: builder = object.__new__(B12xMLAMetadataBuilder) builder._dense_mla_plan = _FakePlan() @@ -533,6 +601,42 @@ def test_b12x_mla_adapter_uses_flattened_non_causal_rows(monkeypatch) -> None: assert lse is not None and lse.shape == (query_rows, 6) +def test_b12x_mla_adapter_uses_tiled_query_visibility(monkeypatch) -> None: + impl, dense_mla = _fake_impl(monkeypatch) + query_rows = 4 + q = torch.randn(query_rows, 8, 576, dtype=torch.bfloat16) + cache = torch.randn(4, 16, 576, dtype=torch.bfloat16) + source_table = torch.tensor([[0, 1, 2, 3, 90]], dtype=torch.int32) + verify_table = source_table[:, :4].contiguous() + query_cache_seq_lens = torch.tensor([29, 30, 31, 32], dtype=torch.int32) + query_start_loc = torch.tensor([0, 4], dtype=torch.int32) + metadata = SimpleNamespace( + dense_mla_plan=_FakePlan(), + dense_mla_scratch=torch.empty(256, dtype=torch.uint8), + dense_mla_verify_block_table=verify_table, + dense_mla_query_cache_seq_lens=query_cache_seq_lens, + query_start_loc=query_start_loc, + decode=SimpleNamespace( + block_table=source_table, + seq_lens=torch.tensor([32], dtype=torch.int32), + ), + ) + + output, lse = impl.forward_mqa( + q, + cache, + metadata, + SimpleNamespace(_q_scale=None, _k_scale=None), + ) + + binding = dense_mla.bindings[0] + assert binding.page_table is verify_table + assert binding.cu_seqlens_q.data_ptr() == query_start_loc.data_ptr() + assert binding.query_cache_seqlens is query_cache_seq_lens + assert output.shape == (query_rows, 8, 512) + assert lse is not None and lse.shape == (query_rows, 8) + + def test_b12x_mla_adapter_passes_fp8_scales(monkeypatch) -> None: impl, dense_mla = _fake_impl(monkeypatch) q = torch.empty(1, 8, 576, dtype=torch.float8_e4m3fn) @@ -616,6 +720,53 @@ def reduce(output, lse, actual_group, **kwargs): assert lse is None +def test_b12x_mla_adapter_skips_query_gather_for_qrep(monkeypatch) -> None: + impl, dense_mla = _fake_impl(monkeypatch, num_heads=6) + impl.dcp_world_size = 8 + impl.dcp_q_replicate = True + batch = 2 + q = torch.randn(batch, 48, 576, dtype=torch.bfloat16) + cache = torch.randn(4, 16, 576, dtype=torch.bfloat16) + group = SimpleNamespace(world_size=8) + calls: list[str] = [] + + monkeypatch.setattr(b12x_mla, "get_dcp_group", lambda: group) + + def unexpected_gather(*args, **kwargs): + raise AssertionError("qrep must skip the query all-gather") + + def reduce(output, lse, actual_group, **kwargs): + assert actual_group is group + calls.append("reduce") + return output[:, :6] + + monkeypatch.setattr(b12x_mla, "dcp_b12x_all_gather_heads", unexpected_gather) + monkeypatch.setattr(b12x_mla, "dcp_a2a_lse_reduce", reduce) + metadata = SimpleNamespace( + dense_mla_plan=_FakePlan(), + dense_mla_scratch=torch.empty(256, dtype=torch.uint8), + dense_mla_padded_q=torch.empty(batch, 48, 576, dtype=torch.bfloat16), + dense_mla_padded_output=torch.zeros(batch, 48, 512, dtype=torch.bfloat16), + query_start_loc=torch.tensor([0, 1, 2], dtype=torch.int32), + decode=SimpleNamespace( + block_table=torch.tensor([[0, 1], [2, 3]], dtype=torch.int32), + seq_lens=torch.tensor([17, 31], dtype=torch.int32), + ), + ) + + output, lse = impl.forward_mqa( + q, + cache, + metadata, + SimpleNamespace(_q_scale=None, _k_scale=None), + ) + + assert calls == ["reduce"] + assert dense_mla.bindings[0].q.data_ptr() == q.data_ptr() + assert output.shape == (batch, 6, 512) + assert lse is None + + def test_b12x_mla_adapter_skips_dcp_for_replicated_cache(monkeypatch) -> None: impl, dense_mla = _fake_impl(monkeypatch, num_heads=6) impl.dcp_world_size = 8 diff --git a/vllm/envs.py b/vllm/envs.py index 92d8c3e0098a..0ca61813913c 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -236,6 +236,12 @@ VLLM_USE_DEEP_GEMM_E8M0: bool = True VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True VLLM_DCP_Q_REPLICATE: bool = False + VLLM_K3_DCP_Q_REPLICATE_LAYERS: str | None = None + VLLM_K3_DYNAMIC_SPARSE_STRIDE: int = 1 + VLLM_K3_DYNAMIC_SPARSE_MIN_TOKENS: int = 0 + VLLM_K3_DYNAMIC_SPARSE_SINK_TOKENS: int = 4096 + VLLM_K3_DYNAMIC_SPARSE_RECENT_TOKENS: int = 32768 + VLLM_K3_DYNAMIC_SPARSE_REFRESH_INTERVAL: int = 128 VLLM_USE_DIRECT_DCP_A2A: bool | None = None VLLM_USE_DIRECT_DCP_Q_GATHER: bool | None = None VLLM_USE_DIRECT_DCP_KV_GATHER: bool | None = None @@ -1805,6 +1811,28 @@ def _resolve_rust_cli_path() -> str | None: ), # Opt-in MLA DCP query replication: skip the decode query all-gather. "VLLM_DCP_Q_REPLICATE": lambda: bool(int(os.getenv("VLLM_DCP_Q_REPLICATE", "0"))), + # Required Kimi-K3 layer allow-list when DCP query replication is enabled. + # Use "all" only after explicitly qualifying the persistent VRAM cost. + "VLLM_K3_DCP_Q_REPLICATE_LAYERS": lambda: os.getenv( + "VLLM_K3_DCP_Q_REPLICATE_LAYERS" + ), + # Experimental Kimi-K3 dense-MLA sparsity. Stride 1 is exact and leaves + # the production path unchanged. + "VLLM_K3_DYNAMIC_SPARSE_STRIDE": lambda: int( + os.getenv("VLLM_K3_DYNAMIC_SPARSE_STRIDE", "1") + ), + "VLLM_K3_DYNAMIC_SPARSE_MIN_TOKENS": lambda: int( + os.getenv("VLLM_K3_DYNAMIC_SPARSE_MIN_TOKENS", "0") + ), + "VLLM_K3_DYNAMIC_SPARSE_SINK_TOKENS": lambda: int( + os.getenv("VLLM_K3_DYNAMIC_SPARSE_SINK_TOKENS", "4096") + ), + "VLLM_K3_DYNAMIC_SPARSE_RECENT_TOKENS": lambda: int( + os.getenv("VLLM_K3_DYNAMIC_SPARSE_RECENT_TOKENS", "32768") + ), + "VLLM_K3_DYNAMIC_SPARSE_REFRESH_INTERVAL": lambda: int( + os.getenv("VLLM_K3_DYNAMIC_SPARSE_REFRESH_INTERVAL", "128") + ), # DeepGemm JITs the kernels on-demand. The warmup attempts to make DeepGemm # JIT all the required kernels before model execution so there is no # JIT'ing in the hot-path. However, this warmup increases the engine diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 766bd1570bfc..330ebfdf2bb6 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -29,6 +29,7 @@ """ import math +import re from typing import TYPE_CHECKING, cast import torch @@ -43,6 +44,7 @@ get_current_vllm_config, ) from vllm.distributed import ( + get_dcp_group, get_tensor_model_parallel_world_size, ) from vllm.forward_context import get_forward_context @@ -60,6 +62,7 @@ from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import ( ColumnParallelLinear, + DCPGroupColumnParallelLinear, MergedColumnParallelLinear, ReplicatedLinear, RowParallelLinear, @@ -117,6 +120,64 @@ _MLA_CALLER_OUTPUT_MIN_TOKENS = 1024 +def _parse_k3_qrep_layers(spec: str) -> frozenset[int] | None: + if spec.strip().lower() == "all": + return None + layers: set[int] = set() + for item in spec.split(","): + item = item.strip() + if not item: + continue + if "-" in item: + start_text, end_text = item.split("-", 1) + start = int(start_text) + end = int(end_text) + if start < 0 or end < start: + raise ValueError(f"Invalid K3 qrep layer range: {item!r}") + layers.update(range(start, end + 1)) + else: + layer = int(item) + if layer < 0: + raise ValueError(f"Invalid K3 qrep layer: {item!r}") + layers.add(layer) + return frozenset(layers) + + +def _k3_dcp_qrep_enabled(prefix: str, vllm_config: VllmConfig) -> bool: + parallel_config = vllm_config.parallel_config + if ( + not envs.VLLM_DCP_Q_REPLICATE + or parallel_config.decode_context_parallel_size <= 1 + or parallel_config.prefill_context_parallel_size > 1 + ): + return False + layer_spec = envs.VLLM_K3_DCP_Q_REPLICATE_LAYERS + if layer_spec is None: + raise ValueError( + "Kimi-K3 DCP query replication duplicates query and absorbed " + "projection weights. Set VLLM_K3_DCP_Q_REPLICATE_LAYERS to an " + "explicit layer list/range, or to 'all' after verifying the VRAM " + "budget." + ) + match = re.search(r"(?:^|\.)layers\.(\d+)(?:\.|$)", prefix) + if match is None: + raise ValueError( + "VLLM_K3_DCP_Q_REPLICATE_LAYERS requires a layer-qualified prefix, " + f"got {prefix!r}" + ) + layers = _parse_k3_qrep_layers(layer_spec) + return layers is None or int(match.group(1)) in layers + + +def _k3_projected_query_heads( + num_local_heads: int, + dcp_world_size: int, + dcp_q_replicate: bool, +) -> int: + """Return the DCP-group head width emitted by the query projection.""" + return int(num_local_heads) * (int(dcp_world_size) if dcp_q_replicate else 1) + + @torch.compile(backend=current_platform.simple_compile_backend) def _gate_sigmoid_mul(attn_out: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: """Apply the sigmoid output gate to a precomputed ``g_proj`` projection.""" @@ -294,6 +355,13 @@ def __init__( assert num_heads % tp_size == 0 self.num_heads = num_heads self.num_local_heads = num_heads // tp_size + vllm_config = get_current_vllm_config() + self.dcp_q_replicate = _k3_dcp_qrep_enabled(prefix, vllm_config) + q_proj_cls = ( + DCPGroupColumnParallelLinear + if self.dcp_q_replicate + else ColumnParallelLinear + ) # ---- Pre-attention projections (fusable front-end) ---- # Two query variants: a low-rank q-LoRA path (Kimi-K3) fused with the @@ -325,7 +393,7 @@ def __init__( disable_tp=True, ) self.q_a_layernorm = RMSNorm(self.q_lora_rank, eps=config.rms_norm_eps) - self.q_b_proj = ColumnParallelLinear( + self.q_b_proj = q_proj_cls( self.q_lora_rank, self.num_heads * self.qk_head_dim, bias=False, @@ -335,7 +403,7 @@ def __init__( else: # Uncompressed query: full-rank q_proj (TP-split over heads) plus a # replicated kv-down projection (shared latent across TP ranks). - self.q_proj = ColumnParallelLinear( + self.q_proj = q_proj_cls( self.hidden_size, self.num_heads * self.qk_head_dim, bias=False, @@ -440,7 +508,6 @@ def __init__( self.impl.dcp_rank = 0 self.q_pad_num_heads = getattr(self.impl, "q_pad_num_heads", None) - vllm_config = get_current_vllm_config() parallel_config = vllm_config.parallel_config assert parallel_config.prefill_context_parallel_size == 1, ( "Kimi-K3 MultiHeadLatentAttention does not support prefill context " @@ -450,6 +517,13 @@ def __init__( self.backend_owns_decode_dcp = _backend_owns_decode_dcp( self.impl, self.dcp_world_size ) + if self.dcp_q_replicate: + if not self.backend_owns_decode_dcp: + raise NotImplementedError( + "Kimi-K3 DCP query replication requires a backend-owned " + "decode DCP path." + ) + self.impl.dcp_q_replicate = True assert ( self.dcp_world_size <= 1 or self.rotary_emb is None @@ -639,6 +713,12 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: pre_w_uk_t.copy_(w_uk_t) w_uk_t = pre_w_uk_t replace_parameter(self, "W_UK_T", w_uk_t, prefer_copy=True) + self.W_UK_T_dcp_qrep: torch.Tensor | None = None + if self.dcp_q_replicate: + self.W_UK_T_dcp_qrep = get_dcp_group().all_gather( + self.W_UK_T.contiguous(), + dim=0, + ) quant_method = ( self.quant_config.get_quant_method(self, prefix=self.layer_name) @@ -678,9 +758,15 @@ def _absorb_decode_query(self, q_nope: torch.Tensor) -> torch.Tensor: """ query = q_nope.transpose(0, 1).contiguous() output = query.new_empty((query.shape[0], query.shape[1], self.kv_lora_rank)) + weight = ( + self.W_UK_T_dcp_qrep + if getattr(self, "dcp_q_replicate", False) + else self.W_UK_T + ) + assert weight is not None _run_mla_query_bmm( query, - self.W_UK_T, + weight, output, use_safe_op=True, ) @@ -730,13 +816,21 @@ def _forward_attn( self.kv_a_layernorm.weight.data, self.rms_norm_eps, ) - q = self.q_b_proj(q_c)[0].view(-1, self.num_local_heads, self.qk_head_dim) + q_heads = _k3_projected_query_heads( + self.num_local_heads, + self.dcp_world_size, + self.dcp_q_replicate, + ) + q = self.q_b_proj(q_c)[0].view(-1, q_heads, self.qk_head_dim) else: # Uncompressed query: project directly (no q-LoRA, no q norm) and # normalize only the kv latent. - q = self.q_proj(hidden_states)[0].view( - -1, self.num_local_heads, self.qk_head_dim + q_heads = _k3_projected_query_heads( + self.num_local_heads, + self.dcp_world_size, + self.dcp_q_replicate, ) + q = self.q_proj(hidden_states)[0].view(-1, q_heads, self.qk_head_dim) kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0] kv_c, k_pe = kv_lora.split( [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 @@ -855,8 +949,12 @@ def _attention( # ---- Prefill: fused key-concat + cache-insert + attention ---- if num_mha_tokens > 0: + prefill_q = q[num_mqa_tokens:] + if getattr(self, "dcp_q_replicate", False): + q_proj = self.q_b_proj if self.q_lora_rank is not None else self.q_proj + prefill_q = q_proj._local_view(prefill_q) self._forward_prefill_fused( - q[num_mqa_tokens:], + prefill_q, kv_c_normed[num_mqa_tokens:], k_pe[num_mqa_tokens:], rope_positions[num_mqa_tokens:] if rope_positions is not None else None, diff --git a/vllm/v1/attention/backends/mla/b12x_mla.py b/vllm/v1/attention/backends/mla/b12x_mla.py index b22c4c253673..98c017fd5afc 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla.py +++ b/vllm/v1/attention/backends/mla/b12x_mla.py @@ -4,11 +4,13 @@ from __future__ import annotations +from bisect import bisect_left from dataclasses import dataclass from typing import Any, ClassVar, cast import torch +from vllm import envs from vllm.config import VllmConfig, get_current_vllm_config from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import get_dcp_group @@ -141,6 +143,34 @@ def _active_dense_mla_splits(plan: Any, max_seq_len: int | None) -> int: ) +def _dense_mla_plan_row_caps(max_rows: int) -> tuple[int, ...]: + """Return CUDA-graph-friendly row capacities through ``max_rows``.""" + if max_rows <= 0: + raise ValueError("dense MLA row capacity must be positive") + caps: list[int] = [] + row_cap = 1 + while row_cap < max_rows: + caps.append(row_cap) + row_cap *= 2 + caps.append(max_rows) + return tuple(caps) + + +def _select_dense_mla_plan( + plans: dict[int, Any], + total_rows: int, +) -> Any: + """Select the smallest launch plan that covers the live query rows.""" + row_caps = tuple(sorted(plans)) + index = bisect_left(row_caps, total_rows) + if total_rows <= 0 or index >= len(row_caps): + raise ValueError( + "B12X_MLA query rows exceed the planned capacities: " + f"rows={total_rows}, capacities={row_caps}" + ) + return plans[row_caps[index]] + + def _create_dense_mla_plan( vllm_config: VllmConfig, device: torch.device, @@ -148,6 +178,9 @@ def _create_dense_mla_plan( page_size: int, num_q_heads: int, max_total_q: int | None = None, + max_batch: int | None = None, + mode: str = "decode", + uses_query_cache_seqlens: bool = False, dcp_size: int | None = None, max_cache_tokens: int | None = None, ) -> Any: @@ -162,6 +195,40 @@ def _create_dense_mla_plan( if max_cache_tokens is not None else _max_dcp_local_cache_tokens(vllm_config, dcp_size=dcp_size) ) + max_batch = int(max_total_q if max_batch is None else max_batch) + dcp_size = int( + vllm_config.parallel_config.decode_context_parallel_size + if dcp_size is None + else dcp_size + ) + + def local_tokens(global_tokens: int) -> int: + return (max(int(global_tokens), 0) + dcp_size - 1) // dcp_size + + sparse_stride = int(envs.VLLM_K3_DYNAMIC_SPARSE_STRIDE) + if sparse_stride > 1: + logger.warning_once( + "Kimi-K3 dynamic sparse MLA is enabled with stride=%d. This " + "changes attention semantics and requires model-quality " + "qualification; stride=1 is the exact production default.", + sparse_stride, + ) + sparse_min_tokens = local_tokens(envs.VLLM_K3_DYNAMIC_SPARSE_MIN_TOKENS) + sparse_sink_chunks = ( + local_tokens(envs.VLLM_K3_DYNAMIC_SPARSE_SINK_TOKENS) + 63 + ) // 64 + sparse_recent_chunks = ( + local_tokens(envs.VLLM_K3_DYNAMIC_SPARSE_RECENT_TOKENS) + 63 + ) // 64 + # A local DCP shard length can differ from its peers at an interleave + # boundary. Until the kernel accepts a global refresh clock, disabling the + # periodic dense refresh under DCP avoids mixing dense and sparse shards in + # one exact LSE reduction. The sparse policy itself remains rank-consistent. + sparse_refresh_interval = ( + local_tokens(envs.VLLM_K3_DYNAMIC_SPARSE_REFRESH_INTERVAL) + if dcp_size == 1 + else 0 + ) if max_total_q > _MAX_B12X_QUERY_ROWS: raise ValueError( "B12X_MLA supports at most " @@ -175,17 +242,23 @@ def _create_dense_mla_plan( caps = dense_mla.Caps( device=device, - mode="decode", + mode=mode, dtype=torch.bfloat16, kv_dtype=_planned_kv_dtype(vllm_config), num_q_heads=num_q_heads, page_size=page_size, max_total_q=max_total_q, - max_batch=max_total_q, + max_batch=max_batch, max_cache_tokens=max_cache_tokens, max_page_table_width=_page_table_width(max_cache_tokens, page_size), num_cache_pages=_MAX_I32, use_cuda_graph=True, + uses_query_cache_seqlens=uses_query_cache_seqlens, + sparse_stride=sparse_stride, + sparse_min_tokens=sparse_min_tokens, + sparse_sink_chunks=sparse_sink_chunks, + sparse_recent_chunks=sparse_recent_chunks, + sparse_refresh_interval=sparse_refresh_interval, ) return dense_mla.plan(caps) @@ -201,6 +274,8 @@ class B12xMLAMetadata(MLACommonMetadata): dense_mla_flat_block_table: torch.Tensor | None = None dense_mla_flat_seq_lens: torch.Tensor | None = None dense_mla_flat_query_start_loc: torch.Tensor | None = None + dense_mla_verify_block_table: torch.Tensor | None = None + dense_mla_query_cache_seq_lens: torch.Tensor | None = None dense_mla_dcp_world_size: int = 1 @@ -245,19 +320,55 @@ def __init__( sliding_window = getattr(kv_cache_spec, "sliding_window", None) if sliding_window is not None: max_cache_tokens = min(max_cache_tokens, int(sliding_window)) - self._dense_mla_plan = _create_dense_mla_plan( - vllm_config, - device, - page_size=self.page_size, - num_q_heads=self._kernel_heads, - max_total_q=max_dense_mla_rows, - dcp_size=self.dcp_world_size, - max_cache_tokens=max_cache_tokens, - ) - self._workspace_specs = self._dense_mla_plan.shapes_and_dtypes() - if len(self._workspace_specs) != 1: - raise RuntimeError("B12X_MLA expected exactly one scratch buffer.") - scratch_shape, scratch_dtype = self._workspace_specs[0] + self._dense_mla_plans = { + rows: _create_dense_mla_plan( + vllm_config, + device, + page_size=self.page_size, + num_q_heads=self._kernel_heads, + max_total_q=rows, + dcp_size=self.dcp_world_size, + max_cache_tokens=max_cache_tokens, + ) + for rows in _dense_mla_plan_row_caps(max_dense_mla_rows) + } + self._dense_mla_verify_plans: dict[int, Any] = {} + if _planned_kv_dtype(vllm_config) == torch.float8_e4m3fn: + self._dense_mla_verify_plans = { + batch: _create_dense_mla_plan( + vllm_config, + device, + page_size=self.page_size, + num_q_heads=self._kernel_heads, + max_total_q=batch * 4, + max_batch=batch, + mode="verify", + uses_query_cache_seqlens=True, + dcp_size=self.dcp_world_size, + max_cache_tokens=max_cache_tokens, + ) + for batch in range( + 1, + int(vllm_config.scheduler_config.max_num_seqs) + 1, + ) + } + self._dense_mla_plan = self._dense_mla_plans[max_dense_mla_rows] + workspace_specs = [ + plan.shapes_and_dtypes() + for plan in ( + *self._dense_mla_plans.values(), + *self._dense_mla_verify_plans.values(), + ) + ] + if any(len(specs) != 1 for specs in workspace_specs): + raise RuntimeError("B12X_MLA expected exactly one scratch buffer per plan.") + scratch_dtype = workspace_specs[0][0][1] + if any(specs[0][1] != scratch_dtype for specs in workspace_specs): + raise RuntimeError("B12X_MLA plan scratch dtypes do not match.") + scratch_shape = max( + (specs[0][0] for specs in workspace_specs), + key=lambda shape: shape[0], + ) # Every attention layer represented by this builder executes serially # on the model stream. One builder-owned buffer therefore gives each # eager bind a stable caller-owned address without a backend workspace @@ -321,18 +432,83 @@ def __init__( else None ) logger.info_once( - "B12X dense K3 MLA plan: local_heads=%d, effective_heads=%d, " + "B12X dense K3 MLA plans: local_heads=%d, effective_heads=%d, " "kernel_heads=%d, page_size=%d, " - "max_decode_rows=%d, max_cache_tokens=%d, splits=%d", + "max_decode_rows=%d, max_cache_tokens=%d, rows/splits=%s, " + "verify_batch/splits=%s", self.num_heads, self._effective_heads, self._kernel_heads, self.page_size, max_dense_mla_rows, max_cache_tokens, - self._dense_mla_plan.num_splits, + ",".join( + f"{rows}/{plan.num_splits}" + for rows, plan in self._dense_mla_plans.items() + ), + ",".join( + f"{batch}/{plan.num_splits}" + for batch, plan in self._dense_mla_verify_plans.items() + ) + or "disabled", ) + def _materialize_query_cache_seq_lens( + self, + metadata: B12xMLAMetadata, + decode_metadata: Any, + *, + query_len: int, + total_q: int, + ) -> torch.Tensor: + flat_lens = self._dense_mla_flat_seq_lens[:total_q] + if not metadata.causal: + flat_lens.copy_( + decode_metadata.seq_lens[:, None].expand(-1, query_len).reshape(total_q) + ) + return flat_lens + + offsets = self._dense_mla_causal_offsets[-query_len:] + if self.dcp_world_size == 1: + torch.add( + decode_metadata.seq_lens[:, None], + offsets, + out=flat_lens.view(metadata.num_decodes, query_len), + ) + return flat_lens + + global_source_lens = decode_metadata.dcp_tot_seq_lens + if global_source_lens is None: + raise RuntimeError( + "B12X_MLA causal DCP verification requires global decode " + "sequence lengths." + ) + assert self._dense_mla_flat_global_seq_lens is not None + assert self._dense_mla_flat_dcp_remainder is not None + global_flat_lens = self._dense_mla_flat_global_seq_lens[:total_q] + torch.add( + global_source_lens[:, None], + offsets, + out=global_flat_lens.view(metadata.num_decodes, query_len), + ) + virtual_block = self.dcp_world_size * self.cp_kv_cache_interleave_size + torch.div( + global_flat_lens, + virtual_block, + rounding_mode="floor", + out=flat_lens, + ) + flat_lens.mul_(self.cp_kv_cache_interleave_size) + remainder = self._dense_mla_flat_dcp_remainder[:total_q] + torch.remainder(global_flat_lens, virtual_block, out=remainder) + remainder.sub_(self._dcp_rank * self.cp_kv_cache_interleave_size) + remainder.clamp_( + min=0, + max=self.cp_kv_cache_interleave_size, + ) + flat_lens.add_(remainder) + return flat_lens + def build( self, common_prefix_len: int, @@ -347,95 +523,74 @@ def build( fast_build=fast_build, ), ) - metadata.dense_mla_plan = self._dense_mla_plan + live_rows = max(1, int(metadata.num_decode_tokens)) + plans = getattr( + self, + "_dense_mla_plans", + {self._max_dense_mla_rows: self._dense_mla_plan}, + ) + metadata.dense_mla_plan = _select_dense_mla_plan(plans, live_rows) metadata.dense_mla_scratch = self._dense_mla_scratch metadata.dense_mla_padded_q = self._dense_mla_padded_q metadata.dense_mla_padded_output = self._dense_mla_padded_output metadata.dense_mla_dcp_world_size = self.dcp_world_size decode_metadata = metadata.decode - flatten_decode = False - if decode_metadata is not None and metadata.num_decodes > 0: - flatten_decode = metadata.num_decode_tokens > metadata.num_decodes or int( - decode_metadata.block_table.shape[1] - ) > int(self._dense_mla_plan.caps.max_page_table_width) - if flatten_decode: - assert decode_metadata is not None - total_q = int(metadata.num_decode_tokens) - if total_q > self._max_dense_mla_rows: - raise ValueError( - "B12X_MLA query block exceeds its flattened capacity: " - f"rows={total_q}, capacity={self._max_dense_mla_rows}." - ) - if total_q % metadata.num_decodes: - raise ValueError( - "B12X_MLA requires a uniform query block, got " - f"tokens={total_q}, requests={metadata.num_decodes}." - ) - query_len = total_q // metadata.num_decodes - source_table = decode_metadata.block_table - flat_table = self._dense_mla_flat_block_table[:total_q] - # A bounded speculative cache can retain a position-indexed worker - # table wider than the resident cache. Sequence lengths make the - # omitted suffix unreachable by the dense-MLA kernel. - source_width = min(int(source_table.shape[1]), int(flat_table.shape[1])) - flat_table[:, :source_width].copy_( - source_table[:, None, :source_width] - .expand(-1, query_len, -1) - .reshape(total_q, source_width) - ) - flat_lens = self._dense_mla_flat_seq_lens[:total_q] - if metadata.causal: - offsets = self._dense_mla_causal_offsets[-query_len:] - if self.dcp_world_size > 1: - global_source_lens = decode_metadata.dcp_tot_seq_lens - if global_source_lens is None: - raise RuntimeError( - "B12X_MLA causal DCP verification requires global " - "decode sequence lengths." - ) - assert self._dense_mla_flat_global_seq_lens is not None - assert self._dense_mla_flat_dcp_remainder is not None - global_flat_lens = self._dense_mla_flat_global_seq_lens[:total_q] - torch.add( - global_source_lens[:, None], - offsets, - out=global_flat_lens.view(metadata.num_decodes, query_len), - ) - virtual_block = ( - self.dcp_world_size * self.cp_kv_cache_interleave_size - ) - torch.div( - global_flat_lens, - virtual_block, - rounding_mode="floor", - out=flat_lens, - ) - flat_lens.mul_(self.cp_kv_cache_interleave_size) - remainder = self._dense_mla_flat_dcp_remainder[:total_q] - torch.remainder(global_flat_lens, virtual_block, out=remainder) - remainder.sub_(self._dcp_rank * self.cp_kv_cache_interleave_size) - remainder.clamp_( - min=0, - max=self.cp_kv_cache_interleave_size, - ) - flat_lens.add_(remainder) - else: - torch.add( - decode_metadata.seq_lens[:, None], - offsets, - out=flat_lens.view(metadata.num_decodes, query_len), - ) - else: - flat_lens.copy_( - decode_metadata.seq_lens[:, None] - .expand(-1, query_len) - .reshape(total_q) - ) - metadata.dense_mla_flat_block_table = flat_table - metadata.dense_mla_flat_seq_lens = flat_lens - metadata.dense_mla_flat_query_start_loc = ( - self._dense_mla_flat_query_start_loc[: total_q + 1] + if decode_metadata is None or metadata.num_decodes <= 0: + return metadata + multi_query = metadata.num_decode_tokens > metadata.num_decodes + table_too_wide = int(decode_metadata.block_table.shape[1]) > int( + self._dense_mla_plan.caps.max_page_table_width + ) + if not (multi_query or table_too_wide): + return metadata + + total_q = int(metadata.num_decode_tokens) + if total_q > self._max_dense_mla_rows: + raise ValueError( + "B12X_MLA query block exceeds its flattened capacity: " + f"rows={total_q}, capacity={self._max_dense_mla_rows}." + ) + if total_q % metadata.num_decodes: + raise ValueError( + "B12X_MLA requires a uniform query block, got " + f"tokens={total_q}, requests={metadata.num_decodes}." ) + query_len = total_q // metadata.num_decodes + source_table = decode_metadata.block_table + flat_lens = self._materialize_query_cache_seq_lens( + metadata, + decode_metadata, + query_len=query_len, + total_q=total_q, + ) + verify_plans = getattr(self, "_dense_mla_verify_plans", {}) + tiled_verify = ( + metadata.causal and query_len == 4 and metadata.num_decodes in verify_plans + ) + if tiled_verify: + verify_table = self._dense_mla_flat_block_table[: metadata.num_decodes] + source_width = min( + int(source_table.shape[1]), + int(verify_table.shape[1]), + ) + verify_table[:, :source_width].copy_(source_table[:, :source_width]) + metadata.dense_mla_plan = verify_plans[metadata.num_decodes] + metadata.dense_mla_verify_block_table = verify_table + metadata.dense_mla_query_cache_seq_lens = flat_lens + return metadata + + flat_table = self._dense_mla_flat_block_table[:total_q] + source_width = min(int(source_table.shape[1]), int(flat_table.shape[1])) + flat_table[:, :source_width].copy_( + source_table[:, None, :source_width] + .expand(-1, query_len, -1) + .reshape(total_q, source_width) + ) + metadata.dense_mla_flat_block_table = flat_table + metadata.dense_mla_flat_seq_lens = flat_lens + metadata.dense_mla_flat_query_start_loc = self._dense_mla_flat_query_start_loc[ + : total_q + 1 + ] return metadata @@ -638,6 +793,7 @@ def __init__( self._dense_mla = _load_dense_mla() self._dcp_comm_backend = vllm_config.parallel_config.dcp_comm_backend self._dcp_max_batch_size = vllm_config.scheduler_config.max_num_batched_tokens + self.dcp_q_replicate = False self._compiled_bindings: set[tuple[object, ...]] = set() def forward_mqa( @@ -663,6 +819,18 @@ def forward_mqa( block_table = attn_metadata.decode.block_table seq_lens = attn_metadata.decode.seq_lens query_start_loc = attn_metadata.query_start_loc + query_cache_seq_lens = getattr( + attn_metadata, + "dense_mla_query_cache_seq_lens", + None, + ) + verify_block_table = getattr( + attn_metadata, + "dense_mla_verify_block_table", + None, + ) + if verify_block_table is not None: + block_table = verify_block_table flat_block_table = getattr(attn_metadata, "dense_mla_flat_block_table", None) if flat_block_table is not None: block_table = flat_block_table @@ -677,16 +845,11 @@ def forward_mqa( batch = int(seq_lens.shape[0]) total_q = int(q.shape[0]) - if total_q != batch: + if query_cache_seq_lens is None and total_q != batch: raise ValueError( "B12X_MLA requires one query row per prepared decode sequence, " f"got {total_q} rows for {batch} sequences." ) - if int(q.shape[1]) != self.num_heads: - raise ValueError( - f"B12X_MLA expected {self.num_heads} query heads, got {q.shape[1]}." - ) - metadata_dcp_world_size = int( getattr(attn_metadata, "dense_mla_dcp_world_size", self.dcp_world_size) ) @@ -697,33 +860,41 @@ def forward_mqa( ) effective_heads = self.num_heads * metadata_dcp_world_size kernel_heads = _kernel_query_heads(self.num_heads, metadata_dcp_world_size) + qrep_decode = self.dcp_q_replicate and metadata_dcp_world_size > 1 + expected_input_heads = effective_heads if qrep_decode else self.num_heads + if int(q.shape[1]) != expected_input_heads: + raise ValueError( + f"B12X_MLA expected {expected_input_heads} query heads, " + f"got {q.shape[1]}." + ) dcp_group = None if metadata_dcp_world_size > 1: dcp_group = get_dcp_group() - gathered_q = getattr(attn_metadata, "dense_mla_padded_q", None) - if gathered_q is None: - raise RuntimeError( - "B12X_MLA DCP metadata is missing caller-owned query storage." - ) - if int(gathered_q.shape[0]) < total_q: - raise ValueError( - "B12X_MLA DCP query capacity is smaller than the decode " - f"batch: capacity={gathered_q.shape[0]}, required={total_q}." - ) - if gathered_q.dtype != q.dtype: - raise TypeError( - "B12X_MLA DCP query storage does not match the live query: " - f"buffer={gathered_q.dtype}, query={q.dtype}." + if not qrep_decode: + gathered_q = getattr(attn_metadata, "dense_mla_padded_q", None) + if gathered_q is None: + raise RuntimeError( + "B12X_MLA DCP metadata is missing caller-owned query storage." + ) + if int(gathered_q.shape[0]) < total_q: + raise ValueError( + "B12X_MLA DCP query capacity is smaller than the decode " + f"batch: capacity={gathered_q.shape[0]}, required={total_q}." + ) + if gathered_q.dtype != q.dtype: + raise TypeError( + "B12X_MLA DCP query storage does not match the live query: " + f"buffer={gathered_q.dtype}, query={q.dtype}." + ) + gathered_q = gathered_q[:total_q, :effective_heads] + q = dcp_b12x_all_gather_heads( + q, + dcp_group, + max_batch_size=self._dcp_max_batch_size, + output_head_dim=self.kv_lora_rank, + out=gathered_q, ) - gathered_q = gathered_q[:total_q, :effective_heads] - q = dcp_b12x_all_gather_heads( - q, - dcp_group, - max_batch_size=self._dcp_max_batch_size, - output_head_dim=self.kv_lora_rank, - out=gathered_q, - ) actual_heads = int(q.shape[1]) if actual_heads != effective_heads: @@ -790,6 +961,7 @@ def forward_mqa( output=output, page_table=block_table, cache_seqlens=seq_lens, + query_cache_seqlens=query_cache_seq_lens, cu_seqlens_q=query_start_loc[: batch + 1], q_scale=layer._q_scale if quantized else None, kv_scale=layer._k_scale if quantized else None, From 9d505242216e0c96b73f8448e6ea4a3677e3d9f6 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 11 Aug 2026 02:35:59 +0000 Subject: [PATCH 43/52] [Attention][MLA] Fuse Kimi-K3 chunked-context K/V packing The K3 MLA layer delegated chunked-context prefill to `impl._compute_prefill_context`, which per chunk casts `kv_nope` to fp8, casts `k_pe`, concatenates `[k_nope | k_pe]`, and re-quantizes a query the fused new-token epilogue had already quantized. Give the layer its own context loop so that tail collapses into one kernel per chunk: `fused_kimi_k3_mla_kv_concat{,_quant_fp8}` reads the strided `kv_b_proj` output and the gathered `k_pe` in place -- the latter still in its fp8 cache layout -- and writes a contiguous key plus, on the fp8 path, a contiguous fp8 V. Only the gather and `kv_b_proj` remain. Casts to E4M3 use the native pairwise converters, so the result is bit-identical to `.to(torch.float8_e4m3fn)`. The `kv_b_proj` input cast becomes a one-time contract check at weight load. Also wire an `out` tensor through `run_prefill_context_chunk`, reusing the existing `supports_out()` capability. A backend honoring it writes each chunk's partial straight into the accumulating context partial, removing the per-chunk output copy for every non-continuation chunk; a continuation still has to merge with the partial already in place, so it keeps its own buffer. Decode context parallelism is unchanged and still uses `impl._context_parallel_compute_prefill_context`. Tests: pytest tests/models/kimi_k3/test_mla_prefill_context.py \ tests/v1/attention/test_mla_prefill_registry.py \ tests/v1/attention/test_mla_context_chunks.py \ tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py \ tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py -> 63 passed The new `tests/models/kimi_k3/test_mla_prefill_context.py` asserts the fused loop hands the prefill backend the same (q, k, v) per chunk, and returns the same merged partial, as the generic impl -- across bf16/fp8 caches and both `supports_out()` modes. No model eval was run: the change is not intended to alter model semantics, and the fp8 packing is verified exact against torch's cast. AI assistance was used. Signed-off-by: Yongye Zhu Co-Authored-By: Claude Opus 5 (1M context) --- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 281 ++++++++++++++++- csrc/libtorch_stable/ops.h | 10 + csrc/libtorch_stable/torch_bindings.cpp | 11 + .../test_kimi_k3_mla_fused_epilogue.py | 70 +++- .../kimi_k3/test_mla_prefill_context.py | 298 ++++++++++++++++++ .../v1/attention/test_mla_prefill_registry.py | 4 +- .../layers/attention/mla_attention.py | 38 ++- vllm/models/kimi_k3/nvidia/mla.py | 210 +++++++++++- .../ops/fused_mla_key_concat_kv_cache.py | 56 ++++ .../backends/mla/prefill/aiter_flash_attn.py | 5 + .../v1/attention/backends/mla/prefill/base.py | 14 +- .../backends/mla/prefill/flash_attn.py | 2 + .../backends/mla/prefill/flashinfer.py | 2 + .../backends/mla/prefill/tokenspeed_mla.py | 5 +- .../backends/mla/prefill/trtllm_ragged.py | 16 +- 15 files changed, 987 insertions(+), 35 deletions(-) create mode 100644 tests/models/kimi_k3/test_mla_prefill_context.py diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index b0d52e8b816b..db4805dde0fd 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -25,6 +25,14 @@ * matching MLA's _q_scale / _k_scale / _v_scale (the cache latent uses a * separate cache scale). Per-tensor E4M3. * + * K/V pack (fused_kimi_k3_mla_kv_concat{,_quant_fp8}): + * - k_out[t, h] = [k_nope[t, h] | k_pe[t]], cast to fp8 by the _quant_fp8 + * variant, which also writes v_fp8[t, h] = cast(v[t, h]) + * The chunked-context epilogue: no cache insert, no q, no RoPE (the cached + * k_pe is already rotated) and no scaling. `k_nope`/`v` are strided views + * of one kv_b_proj output and `k_pe` may already be fp8 (a plain fp8 cache is + * gathered without dequantizing); the outputs are contiguous. + * * fp8_ds_mla (fused_kimi_k3_mla_key_concat_ds_mla_insert): * - full key concat (bf16), and cache insert in DeepSeek's 656-byte * block-scaled layout (NoPE fp8 in 4 tiles of 128 with per-tile dynamic @@ -183,6 +191,36 @@ __device__ __forceinline__ void copyChunk8(void* dst, const scalar_t* src, } } +// Cast 8 source elements (one uint4 of bf16/fp16) to E4M3 with no scaling, +// using the native pairwise converters so the result is bit-identical to +// `src.to(torch.float8_e4m3fn)`. +template +__device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst, + const scalar_t* src) { +#ifndef USE_ROCM + uint4 const input = *reinterpret_cast(src); + using Converter = vllm::_typeConvert; + auto const* input2 = + reinterpret_cast(&input); + uint2 output; + auto* output2 = reinterpret_cast<__nv_fp8x2_storage_t*>(&output); + #pragma unroll + for (int i = 0; i < 4; ++i) { + if constexpr (std::is_same_v) { + output2[i] = __nv_cvt_bfloat16raw2_to_fp8x2( + static_cast<__nv_bfloat162_raw>(input2[i]), __NV_SATFINITE, + __NV_E4M3); + } else { + output2[i] = __nv_cvt_halfraw2_to_fp8x2( + static_cast<__half2_raw>(input2[i]), __NV_SATFINITE, __NV_E4M3); + } + } + *reinterpret_cast(dst) = output; +#else + copyChunk8(dst, src, 1.0f); +#endif +} + // Concat + store one head's full key: dst[e] = [k_nope | k_pe], e in [0, 192). // FP8 dst is byte-addressed; bf16 dst is scalar_t-addressed (dst_elem_size). template @@ -202,6 +240,38 @@ __device__ __forceinline__ void writeFullKey(void* dst, const scalar_t* k_nope, } } +// Unscaled full-key writer for the chunked-context pack, driven by a half warp +// (16 lanes x 8 elems covers NoPE 128; 8 of them cover RoPE 64). Unlike +// `writeFullKey` there is no scale and no RoPE (a gathered k_pe is already +// rotated), and KPE_FP8 handles a k_pe that the gather left in the fp8 cache +// layout -- so neither dtype needs an intermediate cast or a broadcast copy. +template +__device__ __forceinline__ void writeFullKeyPack(void* dst, + const scalar_t* k_nope, + const void* k_pe, int laneId) { + static_assert(!KPE_FP8 || FP8, "an fp8 k_pe requires an fp8 key output"); + constexpr int kElemSize = FP8 ? 1 : sizeof(scalar_t); + auto* d = reinterpret_cast(dst); + for (int e = laneId * kVecElems; e < kQkHeadDim; e += 16 * kVecElems) { + void* out = d + e * kElemSize; + const scalar_t* src = k_nope + e; + if (e >= kQkNopeHeadDim) { + int const rope_e = e - kQkNopeHeadDim; + if constexpr (KPE_FP8) { + *reinterpret_cast(out) = *reinterpret_cast( + reinterpret_cast(k_pe) + rope_e); + continue; + } + src = reinterpret_cast(k_pe) + rope_e; + } + if constexpr (FP8) { + copyChunk8UnitFp8(reinterpret_cast(out), src); + } else { + *reinterpret_cast(out) = *reinterpret_cast(src); + } + } +} + // Store a prefill query, rotating only q[..., 128:192]. For bf16 dst may alias // q (in-place); fp8 writes the quantized query directly to its output. template @@ -457,6 +527,72 @@ __global__ void fusedKimiK3MLAQKVQuantKVCacheFp8Kernel( #endif } +// ──────────────────────────────────────────────────────────────────────────── +// K/V pack variant (chunked context): concatenate the strided kv_b_proj output +// with k_pe into a contiguous row-major key, casting to E4M3 (and casting V +// into its own contiguous output) when FP8. Unlike the other kernels this one +// has no per-token cache slot, so a warp handles two (token, head) rows -- 16 +// lanes each -- and the grid is capped so a long context becomes a grid-stride +// loop instead of a huge launch. +// ──────────────────────────────────────────────────────────────────────────── +template +__global__ void fusedKimiK3MLAKVConcatPackKernel( + const scalar_t* __restrict__ k_nope, int64_t const kn_tok_stride, + int64_t const kn_head_stride, const void* __restrict__ k_pe, + int64_t const k_pe_tok_stride, const scalar_t* __restrict__ v, + int64_t const v_tok_stride, int64_t const v_head_stride, + void* __restrict__ k_out, int64_t const ko_tok_stride, + int64_t const ko_head_stride, uint8_t* __restrict__ v_fp8, + int64_t const vo_tok_stride, int64_t const vo_head_stride, + int const num_tokens, int const num_heads) { + int const warpsPerBlock = blockDim.x / 32; + int const laneId = threadIdx.x % 16; + int const rowInWarp = (threadIdx.x % 32) / 16; + int64_t const globalWarpIdx = + static_cast(blockIdx.x) * warpsPerBlock + threadIdx.x / 32; + int64_t const totalRows = static_cast(num_tokens) * num_heads; + if (globalWarpIdx * 2 >= totalRows) return; + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaGridDependencySynchronize(); +#endif + + constexpr int kOutElemSize = FP8 ? 1 : sizeof(scalar_t); + int64_t const gridStride = + static_cast(gridDim.x) * warpsPerBlock * 2; + for (int64_t row = globalWarpIdx * 2 + rowInWarp; row < totalRows; + row += gridStride) { + int const tokenIdx = static_cast(row / num_heads); + int const headIdx = static_cast(row % num_heads); + const void* pe; + if constexpr (KPE_FP8) { + pe = reinterpret_cast(k_pe) + tokenIdx * k_pe_tok_stride; + } else { + pe = reinterpret_cast(k_pe) + tokenIdx * k_pe_tok_stride; + } + writeFullKeyPack( + reinterpret_cast(k_out) + + (tokenIdx * ko_tok_stride + headIdx * ko_head_stride) * + kOutElemSize, + k_nope + tokenIdx * kn_tok_stride + headIdx * kn_head_stride, pe, + laneId); + + // bf16 V needs no cast, so it stays a strided view of the kv_b_proj output. + if constexpr (FP8) { + const scalar_t* vh = + v + tokenIdx * v_tok_stride + headIdx * v_head_stride; + uint8_t* vo = v_fp8 + tokenIdx * vo_tok_stride + headIdx * vo_head_stride; + for (int e = laneId * kVecElems; e < kVHeadDim; e += 16 * kVecElems) { + copyChunk8UnitFp8(vo + e, vh + e); + } + } + } + +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900) + cudaTriggerProgrammaticLaunchCompletion(); +#endif +} + // ──────────────────────────────────────────────────────────────────────────── // ds_mla variant: concat full key (bf16) + fp8_ds_mla latent cache insert // @@ -648,16 +784,20 @@ __global__ void fusedKimiK3MLADecodeQConcatDsMlaKernel( #endif } -// PDL-aware launch of a (token, num_heads + 1)-warp grid. +// PDL-aware launch of a (token, slots_per_token)-row grid, `rows_per_warp` rows +// to a warp. `max_blocks > 0` caps the grid, leaving the kernel's grid-stride +// loop to cover the remaining rows. template -static void launchPdl(KernelT kernel, int num_tokens, int num_heads, - cudaStream_t stream, Args... args) { +static void launchPdlSlots(KernelT kernel, int num_tokens, int slots_per_token, + int rows_per_warp, int max_blocks, + cudaStream_t stream, Args... args) { constexpr int kBlockSize = 256; constexpr int kWarpsPerBlock = kBlockSize / 32; - int64_t const total_warps = - static_cast(num_tokens) * (num_heads + 1); - int const grid = + int64_t const total_rows = static_cast(num_tokens) * slots_per_token; + int64_t const total_warps = (total_rows + rows_per_warp - 1) / rows_per_warp; + int grid = static_cast((total_warps + kWarpsPerBlock - 1) / kWarpsPerBlock); + if (max_blocks > 0 && grid > max_blocks) grid = max_blocks; #ifndef USE_ROCM static int const sm_version = getSMVersion(); cudaLaunchConfig_t config; @@ -681,6 +821,13 @@ static void launchPdl(KernelT kernel, int num_tokens, int num_heads, #endif } +// Cache-inserting kernels use one extra warp per token for the cache slot. +template +static void launchPdl(KernelT kernel, int num_tokens, int num_heads, + cudaStream_t stream, Args... args) { + launchPdlSlots(kernel, num_tokens, num_heads + 1, 1, 0, stream, args...); +} + void checkBfloat16Support(torch::headeronly::ScalarType dtype) { #ifndef USE_ROCM if (dtype == torch::headeronly::ScalarType::BFloat16) { @@ -913,6 +1060,128 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert( }); } +namespace { +// Shared checks for the chunked-context K/V pack ops. `k_nope` / `v` are +// strided views of one kv_b_proj output and `k_pe` is a strided view of the +// gathered context workspace, so only a unit last-dim stride (not full +// contiguity) is required of the inputs. Returns true when `k_pe` is still in +// the fp8 cache layout. +bool check_kv_concat_inputs(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& k_out, + torch::headeronly::ScalarType out_dtype) { + using torch::headeronly::ScalarType; + ScalarType const dt = k_nope.scalar_type(); + STD_TORCH_CHECK(k_nope.device().is_cuda() && k_nope.dim() == 3 && + k_nope.size(2) == 128 && k_nope.stride(2) == 1, + "k_nope must be CUDA [T, H, 128] with unit last stride"); + bool const k_pe_is_fp8 = k_pe.scalar_type() == ScalarType::Float8_e4m3fn; + STD_TORCH_CHECK( + k_pe.device() == k_nope.device() && k_pe.dim() == 2 && + k_pe.size(1) == 64 && k_pe.stride(1) == 1 && + (k_pe.scalar_type() == dt || + (k_pe_is_fp8 && out_dtype == ScalarType::Float8_e4m3fn)), + "k_pe must be CUDA [T, 64] with unit last-dim stride and use the input " + "dtype, or float8_e4m3fn when the key output is fp8"); + STD_TORCH_CHECK(k_out.device() == k_nope.device() && k_out.is_contiguous() && + k_out.scalar_type() == out_dtype && k_out.dim() == 3 && + k_out.size(2) == 192 && k_out.size(0) == k_nope.size(0) && + k_out.size(1) == k_nope.size(1) && + k_pe.size(0) == k_nope.size(0), + "k_out must be a contiguous CUDA [T, H, 192] tensor of the " + "output dtype matching k_nope's token and head dimensions"); + vllm::kimi_k3_fused_ops::checkBfloat16Support(dt); + return k_pe_is_fp8; +} +} // namespace + +void fused_kimi_k3_mla_kv_concat( + torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16 + torch::stable::Tensor const& k_pe, // [T, 64] same dtype as k_nope + torch::stable::Tensor& k_out) { // [T, H, 192] same dtype, written + namespace kk3 = vllm::kimi_k3_fused_ops; + torch::headeronly::ScalarType const dt = k_nope.scalar_type(); + check_kv_concat_inputs(k_nope, k_pe, k_out, dt); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES(dt, "fused_kimi_k3_mla_kv_concat", [&] { + kk3::launchPdlSlots( + kk3::fusedKimiK3MLAKVConcatPackKernel, + num_tokens, num_heads, 2, get_device_prop()->multiProcessorCount * 8, + stream, reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(), + k_pe.stride(0), static_cast(nullptr), int64_t{0}, + int64_t{0}, k_out.mutable_data_ptr(), k_out.stride(0), k_out.stride(1), + static_cast(nullptr), int64_t{0}, int64_t{0}, num_tokens, + num_heads); + }); +} + +void fused_kimi_k3_mla_kv_concat_quant_fp8( + torch::stable::Tensor const& k_nope, // [T, H, 128] bf16/fp16 + torch::stable::Tensor const& k_pe, // [T, 64] bf16/fp16/fp8 + torch::stable::Tensor const& v, // [T, H, 128] bf16/fp16 + torch::stable::Tensor& k_fp8, // [T, H, 192] fp8, written + torch::stable::Tensor& v_fp8) { // [T, H, 128] fp8, written + using torch::headeronly::ScalarType; + namespace kk3 = vllm::kimi_k3_fused_ops; + ScalarType const dt = k_nope.scalar_type(); + bool const k_pe_is_fp8 = + check_kv_concat_inputs(k_nope, k_pe, k_fp8, ScalarType::Float8_e4m3fn); + STD_TORCH_CHECK(v.device() == k_nope.device() && v.scalar_type() == dt && + v.dim() == 3 && v.size(2) == 128 && v.stride(2) == 1 && + v.size(0) == k_nope.size(0) && + v.size(1) == k_nope.size(1), + "v must match k_nope's dtype and token/head dimensions and " + "have shape [T, H, 128] with unit last stride"); + STD_TORCH_CHECK(v_fp8.device() == k_nope.device() && v_fp8.is_contiguous() && + v_fp8.scalar_type() == ScalarType::Float8_e4m3fn && + v_fp8.dim() == 3 && v_fp8.size(2) == 128 && + v_fp8.size(0) == k_nope.size(0) && + v_fp8.size(1) == k_nope.size(1), + "v_fp8 must be a contiguous CUDA [T, H, 128] fp8 tensor " + "matching k_nope's token and head dimensions"); + + int const num_tokens = static_cast(k_nope.size(0)); + int const num_heads = static_cast(k_nope.size(1)); + if (num_tokens == 0) return; + + const torch::stable::accelerator::DeviceGuard device_guard( + k_nope.get_device_index()); + const cudaStream_t stream = + get_current_cuda_stream(k_nope.get_device_index()); + + VLLM_STABLE_DISPATCH_HALF_TYPES( + dt, "fused_kimi_k3_mla_kv_concat_quant_fp8", [&] { + auto launch = [&](auto kernel) { + kk3::launchPdlSlots( + kernel, num_tokens, num_heads, 2, + get_device_prop()->multiProcessorCount * 8, stream, + reinterpret_cast(k_nope.const_data_ptr()), + k_nope.stride(0), k_nope.stride(1), k_pe.const_data_ptr(), + k_pe.stride(0), + reinterpret_cast(v.const_data_ptr()), + v.stride(0), v.stride(1), k_fp8.mutable_data_ptr(), + k_fp8.stride(0), k_fp8.stride(1), + reinterpret_cast(v_fp8.mutable_data_ptr()), + v_fp8.stride(0), v_fp8.stride(1), num_tokens, num_heads); + }; + if (k_pe_is_fp8) { + launch(kk3::fusedKimiK3MLAKVConcatPackKernel); + } else { + launch(kk3::fusedKimiK3MLAKVConcatPackKernel); + } + }); +} + void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( torch::stable::Tensor const& q, // [Tp, H, 192] bf16 torch::stable::Tensor const& k_nope, // [Tp, H, 128] bf16 diff --git a/csrc/libtorch_stable/ops.h b/csrc/libtorch_stable/ops.h index 31a61bb13604..3e0a22622071 100644 --- a/csrc/libtorch_stable/ops.h +++ b/csrc/libtorch_stable/ops.h @@ -306,6 +306,16 @@ void fused_kimi_k3_mla_key_concat_ds_mla_insert( std::optional position_ids, std::optional cos_sin_cache); +void fused_kimi_k3_mla_kv_concat(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor& k_out); + +void fused_kimi_k3_mla_kv_concat_quant_fp8(torch::stable::Tensor const& k_nope, + torch::stable::Tensor const& k_pe, + torch::stable::Tensor const& v, + torch::stable::Tensor& k_fp8, + torch::stable::Tensor& v_fp8); + void fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert( torch::stable::Tensor const& q, torch::stable::Tensor const& k_nope, torch::stable::Tensor const& k_pe, torch::stable::Tensor const& kv_c_normed, diff --git a/csrc/libtorch_stable/torch_bindings.cpp b/csrc/libtorch_stable/torch_bindings.cpp index b44ba109729a..1e1896b35007 100644 --- a/csrc/libtorch_stable/torch_bindings.cpp +++ b/csrc/libtorch_stable/torch_bindings.cpp @@ -471,6 +471,13 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "Tensor! k_out, Tensor! k_cache, Tensor slot_mapping, " "int cache_block_size, Tensor? position_ids=None, " "Tensor? cos_sin_cache=None) -> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat(Tensor k_nope, Tensor k_pe, Tensor! k_out) " + "-> ()"); + ops.def( + "fused_kimi_k3_mla_kv_concat_quant_fp8(" + "Tensor k_nope, Tensor k_pe, Tensor v, Tensor! k_fp8, Tensor! v_fp8) " + "-> ()"); ops.def( "fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert(" "Tensor q, Tensor k_nope, Tensor k_pe, Tensor kv_c_normed, Tensor v, " @@ -774,6 +781,10 @@ STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, ops) { TORCH_BOX(&fused_kimi_k3_mla_key_concat_kv_cache_insert)); ops.impl("fused_kimi_k3_mla_key_concat_ds_mla_insert", TORCH_BOX(&fused_kimi_k3_mla_key_concat_ds_mla_insert)); + ops.impl("fused_kimi_k3_mla_kv_concat", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat)); + ops.impl("fused_kimi_k3_mla_kv_concat_quant_fp8", + TORCH_BOX(&fused_kimi_k3_mla_kv_concat_quant_fp8)); ops.impl("fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert", TORCH_BOX(&fused_kimi_k3_mla_qkv_quant_kv_cache_fp8_insert)); ops.impl("fused_kimi_k3_mla_decode_q_concat_kv_cache_insert", diff --git a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py index 0e66e0da1e0d..d7cb7ca7eb45 100644 --- a/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py +++ b/tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py @@ -9,6 +9,8 @@ fused_mla_decode_q_concat_kv_cache_insert, fused_mla_key_concat_ds_mla_insert, fused_mla_key_concat_kv_cache_insert, + fused_mla_kv_concat, + fused_mla_kv_concat_quant_fp8, fused_mla_qkv_quant_kv_cache_fp8_insert, ) from vllm.platforms import current_platform @@ -25,8 +27,8 @@ _SLOTS = (0, 3, 9) -def _randn(*shape: int) -> torch.Tensor: - return torch.randn(*shape, device="cuda", dtype=_DTYPE) * 0.2 +def _randn(*shape: int, dtype: torch.dtype = _DTYPE) -> torch.Tensor: + return torch.randn(*shape, device="cuda", dtype=dtype) * 0.2 def _rope_cache(max_position: int = 32) -> torch.Tensor: @@ -66,6 +68,70 @@ def _assert_fp8_close(actual: torch.Tensor, expected: torch.Tensor) -> None: ) +def _strided_context_inputs( + num_tokens: int, num_heads: int, dtype: torch.dtype, k_pe_fp8: bool +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """The real prefill-context layouts: ``k_nope``/``v`` as the two strided + halves of one kv_b_proj output and ``k_pe`` as a column slice of the gather + workspace (left in the fp8 cache layout for a plain fp8 cache).""" + kv_nope = _randn(num_tokens, num_heads, 256, dtype=dtype) + k_nope, v = kv_nope.split((128, 128), dim=-1) + workspace = _randn(num_tokens, 576, dtype=dtype) + if k_pe_fp8: + workspace = workspace.to(torch.float8_e4m3fn) + return k_nope, workspace[:, 512:].unsqueeze(1), v + + +@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS]) +@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS]) +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_context_kv_concat_accepts_strided_inputs( + num_tokens: int, num_heads: int, input_dtype: torch.dtype +) -> None: + k_nope, k_pe, _ = _strided_context_inputs( + num_tokens, num_heads, input_dtype, k_pe_fp8=False + ) + + k = fused_mla_kv_concat(k_nope, k_pe) + + assert k.is_contiguous() + assert k.dtype == input_dtype + torch.testing.assert_close(k[..., :128], k_nope, atol=0, rtol=0) + torch.testing.assert_close( + k[..., 128:], k_pe.expand(-1, num_heads, -1), atol=0, rtol=0 + ) + + +@pytest.mark.parametrize("k_pe_dtype", ["input", "fp8"]) +@pytest.mark.parametrize("num_tokens", [0, _NUM_TOKENS]) +@pytest.mark.parametrize("num_heads", [3, _NUM_HEADS]) +@pytest.mark.parametrize("input_dtype", [torch.bfloat16, torch.float16]) +@torch.inference_mode() +def test_context_kv_pack_quantizes_strided_inputs( + k_pe_dtype: str, + num_tokens: int, + num_heads: int, + input_dtype: torch.dtype, +) -> None: + k_nope, k_pe, v = _strided_context_inputs( + num_tokens, num_heads, input_dtype, k_pe_fp8=k_pe_dtype == "fp8" + ) + + k_actual, v_actual = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v) + + fp8 = torch.float8_e4m3fn + k_expected = torch.empty_like(k_actual) + k_expected[..., :128] = k_nope.to(fp8) + k_expected[..., 128:] = k_pe.to(fp8) + assert k_actual.is_contiguous() + assert v_actual.is_contiguous() + # The pack casts with the native pairwise converters, so it must be + # bit-identical to torch's `.to(fp8)` rather than merely close. + torch.testing.assert_close(k_actual.float(), k_expected.float(), atol=0, rtol=0) + torch.testing.assert_close(v_actual.float(), v.to(fp8).float(), atol=0, rtol=0) + + @pytest.mark.parametrize("cache_kind", ["bf16", "fp8", "fp8_ds_mla"]) @torch.inference_mode() def test_prefill_epilogue_fuses_gptj_rope(cache_kind: str) -> None: diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py new file mode 100644 index 000000000000..b0166c2a31d0 --- /dev/null +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -0,0 +1,298 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""K3's fused chunked-context prefill must match the generic MLA impl. + +The layer owns its own context loop so it can fuse the per-chunk K/V pack and +skip re-quantizing an already-quantized query. That is only safe if it feeds the +prefill backend exactly what ``MLACommonBaseImpl._compute_prefill_context`` +would, chunk for chunk. +""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm.model_executor.layers.attention.mla_attention import ( + MLACommonBaseImpl, + MLACommonPrefillMetadata, + build_mla_chunked_context_metadata, +) +from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.platforms import current_platform + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA requires CUDA" +) + +_KV_LORA_RANK = 512 +_QK_NOPE = 128 +_QK_ROPE = 64 +_V_HEAD_DIM = 128 +_ENTRY = _KV_LORA_RANK + _QK_ROPE +_NUM_HEADS = 2 +_BLOCK_SIZE = 16 +_WORKSPACE_TOKENS = 128 +# Splits one long request across chunks, packs short ones together, and leaves +# the last request without any context. +_CONTEXT_LENS = [200, 48, 32, 0] +_QUERY_LENS = [8, 4, 6, 5] + + +class _RecordingPrefillBackend: + """Records what each chunk is asked to attend over. + + ``honors_out`` mimics the two backend families: one that writes into a + caller-provided ``out`` (trtllm_ragged, flashinfer, tokenspeed) and one that + always returns its own buffer (flash_attn with a padded V, aiter). + """ + + def __init__(self, honors_out: bool = False) -> None: + self.calls: list[tuple[torch.Tensor, torch.Tensor, torch.Tensor]] = [] + self.out_destinations: list[torch.Tensor | None] = [] + self._honors_out = honors_out + + @staticmethod + def get_name() -> str: + return "recording" + + def supports_out(self) -> bool: + return self._honors_out + + def run_prefill_context_chunk(self, *, chunk, q, k, v, out=None): + self.calls.append((q.float().clone(), k.float().clone(), v.float().clone())) + self.out_destinations.append(out) + assert out is None or self._honors_out + # Fold K/V into the partial so any packing difference shows up in the + # merged context output, not just in the recorded calls. + digest = (k.float().mean() + v.float().mean()).item() + num_q = q.shape[0] + if out is None: + out = torch.empty( + (num_q, _NUM_HEADS, _V_HEAD_DIM), device=q.device, dtype=torch.bfloat16 + ) + else: + assert out.shape == (num_q, _NUM_HEADS, _V_HEAD_DIM) + out.fill_(digest) + lse = torch.full( + (_NUM_HEADS, num_q), + 1.0 + chunk.index, + device=q.device, + dtype=torch.float32, + ) + return out, lse + + +class _KVBProj(torch.nn.Module): + """Stand-in for the layer's ``kv_b_proj`` (returns an (out, bias) tuple). + + With a plain fp8 cache the gathered latent arrives fp8 and K3 feeds it in + without a cast, so ``kv_b_proj`` must consume it directly -- mimicked here by + an fp8 weight whose apply dequantizes, like the fp8 linear methods. Either + weight dtype produces a bf16 output. + """ + + def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: + super().__init__() + weight = ( + torch.randn( + _NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM), + _KV_LORA_RANK, + device=device, + dtype=torch.bfloat16, + ) + * 0.05 + ) + self.register_buffer("weight", weight.to(weight_dtype)) + + def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: + return torch.nn.functional.linear( + x.to(torch.bfloat16), self.weight.to(torch.bfloat16) + ), None + + +class _FusedLayer: + """Only the attributes K3's context loop reads.""" + + _compute_prefill_context = MultiHeadLatentAttention._compute_prefill_context + _gather_context_latent = MultiHeadLatentAttention._gather_context_latent + _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache + + def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: + self.kv_b_proj = kv_b_proj + self.kv_cache = kv_cache + self.kv_cache_dtype = kv_cache_dtype + self._k_scale = k_scale + self.kv_lora_rank = _KV_LORA_RANK + self.num_local_heads = _NUM_HEADS + self.qk_nope_head_dim = _QK_NOPE + self.v_head_dim = _V_HEAD_DIM + + +class _ReferenceImpl: + """Only the attributes the generic context loop reads.""" + + _compute_prefill_context = MLACommonBaseImpl._compute_prefill_context + _concat_k_nope_k_pe = MLACommonBaseImpl._concat_k_nope_k_pe + _use_flashinfer_concat_mla_k = False + + def __init__(self, kv_b_proj, kv_cache_dtype) -> None: + self.kv_b_proj = kv_b_proj + self.kv_cache_dtype = kv_cache_dtype + self.kv_lora_rank = _KV_LORA_RANK + self.num_heads = _NUM_HEADS + self.qk_nope_head_dim = _QK_NOPE + self.qk_rope_head_dim = _QK_ROPE + self.v_head_dim = _V_HEAD_DIM + + +def _build_prefill_metadata( + device: torch.device, + workspace_dtype: torch.dtype, + q_data_type: torch.dtype, + backend: _RecordingPrefillBackend, +) -> MLACommonPrefillMetadata: + query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32) + query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0) + workspace = torch.empty( + (_WORKSPACE_TOKENS, _ENTRY), dtype=workspace_dtype, device=device + ) + chunked_context = build_mla_chunked_context_metadata( + context_lens_cpu=torch.tensor(_CONTEXT_LENS, dtype=torch.int32), + prefill_query_start_loc_cpu=query_start_loc_cpu, + chunked_prefill_workspace=workspace, + chunked_prefill_workspace_size=_WORKSPACE_TOKENS, + block_size=_BLOCK_SIZE, + align_chunk_to_block=True, + device=device, + dcp_world_size=1, + dcp_local_block_size=1, + dcp_virtual_block_size=1, + ) + assert chunked_context is not None + assert len(chunked_context.chunks) > 1, "the batch must exercise accumulation" + + max_blocks = (max(_CONTEXT_LENS) + max(_QUERY_LENS)) // _BLOCK_SIZE + 1 + num_blocks = max_blocks * len(_CONTEXT_LENS) + block_table = torch.arange(num_blocks, dtype=torch.int32, device=device).view( + len(_CONTEXT_LENS), max_blocks + ) + return MLACommonPrefillMetadata( + block_table=block_table, + query_start_loc=query_start_loc_cpu.to(device), + max_query_len=max(_QUERY_LENS), + chunked_context=chunked_context, + q_data_type=q_data_type, + output_dtype=torch.bfloat16, + prefill_backend=backend, + ), num_blocks + + +@pytest.mark.parametrize("honors_out", [False, True], ids=["copy_out", "writes_out"]) +@pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@torch.inference_mode() +def test_fused_context_matches_generic_impl( + kv_cache_dtype: str, honors_out: bool +) -> None: + torch.manual_seed(0) + device = torch.device("cuda") + fp8 = current_platform.fp8_dtype() + quantized = kv_cache_dtype == "fp8" + # An fp8 cache is read by an fp8 query, so the workspace keeps the fp8 + # layout; a bf16 cache dequantizes into a bf16 workspace. + q_data_type = fp8 if quantized else torch.bfloat16 + workspace_dtype = q_data_type + + kv_b_proj = _KVBProj(device, weight_dtype=q_data_type) + k_scale = torch.ones(1, dtype=torch.float32, device=device) + + backend_fused = _RecordingPrefillBackend(honors_out=honors_out) + # The reference impl never passes `out`, so it always allocates its own. + backend_ref = _RecordingPrefillBackend() + prefill_fused, num_blocks = _build_prefill_metadata( + device, workspace_dtype, q_data_type, backend_fused + ) + prefill_ref, _ = _build_prefill_metadata( + device, workspace_dtype, q_data_type, backend_ref + ) + + cache = torch.randn( + (num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16 + ) + kv_cache = cache.to(fp8) if quantized else cache + q = ( + torch.randn( + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), + device=device, + dtype=torch.bfloat16, + ) + * 0.2 + ).to(q_data_type) + + layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale) + fused_out, fused_lse = layer._compute_prefill_context( + q, SimpleNamespace(prefill=prefill_fused) + ) + + impl = _ReferenceImpl(kv_b_proj, kv_cache_dtype) + ref_out, ref_lse = impl._compute_prefill_context( + q, kv_cache, SimpleNamespace(prefill=prefill_ref), k_scale + ) + + assert len(backend_fused.calls) == len(backend_ref.calls) + for chunk_idx, (fused_call, ref_call) in enumerate( + zip(backend_fused.calls, backend_ref.calls, strict=True) + ): + for name, fused_t, ref_t in zip( + ("q", "k", "v"), fused_call, ref_call, strict=True + ): + torch.testing.assert_close( + fused_t, + ref_t, + atol=0, + rtol=0, + msg=lambda m, n=name, i=chunk_idx: f"chunk {i} {n} differs: {m}", + ) + torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0) + torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0) + + # Every chunk but the continuation should have been written in place, i.e. + # straight into the returned accumulator, with no intermediate copy. + wrote_in_place = [ + out is not None and out.data_ptr() == fused_out[chunk.token_slice].data_ptr() + for out, chunk in zip( + backend_fused.out_destinations, + prefill_fused.chunked_context.chunks, + strict=True, + ) + ] + continuations = [c.is_continuation for c in prefill_fused.chunked_context.chunks] + assert any(continuations), "the batch must exercise a continuation chunk" + if honors_out: + assert wrote_in_place == [not c for c in continuations] + else: + assert not any(wrote_in_place) + + +@torch.inference_mode() +def test_fused_context_rejects_an_unquantized_query() -> None: + """The fp8 query is produced by the new-token epilogue, not re-cast here.""" + torch.manual_seed(0) + device = torch.device("cuda") + fp8 = current_platform.fp8_dtype() + prefill, num_blocks = _build_prefill_metadata( + device, fp8, fp8, _RecordingPrefillBackend() + ) + layer = _FusedLayer( + _KVBProj(device, weight_dtype=fp8), + torch.zeros((num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=fp8), + "fp8", + torch.ones(1, dtype=torch.float32, device=device), + ) + q = torch.zeros( + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), + device=device, + dtype=torch.bfloat16, + ) + with pytest.raises(AssertionError, match="new-token epilogue"): + layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) diff --git a/tests/v1/attention/test_mla_prefill_registry.py b/tests/v1/attention/test_mla_prefill_registry.py index 05dcb77838ce..fc5efe3a5be0 100644 --- a/tests/v1/attention/test_mla_prefill_registry.py +++ b/tests/v1/attention/test_mla_prefill_registry.py @@ -24,7 +24,7 @@ def get_name() -> str: def run_prefill_new_tokens(self, q, k, v, return_softmax_lse): raise NotImplementedError - def run_prefill_context_chunk(self, chunk, q, k, v): + def run_prefill_context_chunk(self, chunk, q, k, v, out=None): raise NotImplementedError @@ -112,7 +112,7 @@ def get_name() -> str: def run_prefill_new_tokens(self, q, k, v, return_softmax_lse): raise NotImplementedError - def run_prefill_context_chunk(self, chunk, q, k, v): + def run_prefill_context_chunk(self, chunk, q, k, v, out=None): raise NotImplementedError assert MLAPrefillBackendEnum.CUSTOM.is_overridden() diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 2fe1154aa691..83ea816e59b3 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -3540,6 +3540,22 @@ def _gather_dcp_context_kv( raise RuntimeError("MLA DCP chunked prefill has no configured KV gather path.") +def neutralize_empty_context_partials( + chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata", + output: torch.Tensor, + output_lse: torch.Tensor, +) -> None: + """Neutralize the partial of every prefill that no chunk covers. + + A prefill without context is never gathered, so nothing would write its rows; + a zero output with an ``-inf`` lse carries no weight into the final merge + against the suffix partial. + """ + for token_slice in chunked_context.empty_token_slices: + output[token_slice].zero_() + output_lse[:, token_slice].fill_(float("-inf")) + + def init_mla_context_partial( chunked_context: "MLACommonPrefillMetadata.ChunkedContextMetadata", attn_output: torch.Tensor, @@ -3549,7 +3565,9 @@ def init_mla_context_partial( """Allocate the running context partial over all prefill tokens. Laid out like the chunk partials so the final whole-batch merge against the - suffix partial sees matching head strides. + suffix partial sees matching head strides. Callers whose backend honors an + ``out`` tensor already know that layout and can allocate directly, pairing it + with ``neutralize_empty_context_partials``. """ output = torch.empty( (num_tokens, *attn_output.shape[1:]), @@ -3561,10 +3579,7 @@ def init_mla_context_partial( dtype=attn_softmax_lse.dtype, device=attn_softmax_lse.device, ) - # No chunk covers a prefill without context, so neutralize its partial. - for token_slice in chunked_context.empty_token_slices: - output[token_slice].zero_() - output_lse[:, token_slice].fill_(float("-inf")) + neutralize_empty_context_partials(chunked_context, output, output_lse) return output, output_lse @@ -3574,16 +3589,26 @@ def accumulate_mla_context_chunk( attn_softmax_lse: torch.Tensor, output: torch.Tensor, output_lse: torch.Tensor, + output_written: bool = False, ) -> None: """Fold one chunk's partial into the running context partial. Only the first request may be a continuation; its tokens are merged and the remaining token range is initialized. + + Args: + output_written: The chunk's attention output already landed in + ``output[chunk.token_slice]`` because the backend was handed it as + ``out``, leaving only the lse to fold. Invalid for a continuation + chunk, whose leading tokens must be merged rather than overwritten. """ token_start = chunk.token_slice.start token_end = chunk.token_slice.stop init_start = token_start if chunk.is_continuation: + assert not output_written, ( + "a continuation chunk must not write over the partial it merges with" + ) init_start = chunk.continuation_token_end num_merged = init_start - token_start prefix_output, suffix_output = _match_merge_strides( @@ -3599,7 +3624,8 @@ def accumulate_mla_context_chunk( ) if init_start < token_end: written = init_start - token_start - output[init_start:token_end].copy_(attn_output[written:]) + if not output_written: + output[init_start:token_end].copy_(attn_output[written:]) output_lse[:, init_start:token_end].copy_(attn_softmax_lse[:, written:]) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 330ebfdf2bb6..cd0996bb068b 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -8,8 +8,9 @@ -> fused pre-attention ops (fused_qkv_a_proj / norms / q_b_proj) -> explicit prefill / decode split prefill: fused key-concat + cache-insert kernel -> run_prefill_new_tokens - (+ chunked-context merge); dispatched by cache dtype - (bf16 / plain fp8 / fp8_ds_mla) + (+ chunked-context merge, whose per-chunk gather -> kv_b_proj + -> fused K/V pack loop this layer owns); dispatched by cache + dtype (bf16 / plain fp8 / fp8_ds_mla) decode : W_UK absorb (BMM1) -> fused q-concat + cache-insert kernel -> impl.forward_mqa -> W_UV up-proj (MQA) -> optional output gate @@ -55,8 +56,12 @@ should_load_quant_weights, ) from vllm.model_executor.layers.attention.mla_attention import ( + _get_kv_b_proj_input_dtype, _preallocate_absorbed_mla_weights, _run_mla_query_bmm, + accumulate_mla_context_chunk, + init_mla_context_partial, + neutralize_empty_context_partials, ) from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.layernorm import RMSNorm @@ -79,6 +84,8 @@ fused_mla_decode_q_concat_kv_cache_insert, fused_mla_key_concat_ds_mla_insert, fused_mla_key_concat_kv_cache_insert, + fused_mla_kv_concat, + fused_mla_kv_concat_quant_fp8, fused_mla_qkv_quant_kv_cache_fp8_insert, ) from vllm.models.kimi_k3.nvidia.tp_projection import ( @@ -720,6 +727,25 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: dim=0, ) + # `_compute_prefill_context` hands the gathered context latent straight to + # `kv_b_proj`. With a plain fp8 cache the gather leaves it fp8, so + # `kv_b_proj` has to take an fp8 input; check that here rather than + # casting per chunk (the generic impl casts instead). fp8_ds_mla is + # exempt: it up-converts into a bf16 workspace. + if ( + is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype != "fp8_ds_mla" + ): + assert _get_kv_b_proj_input_dtype(self.kv_b_proj, use_fp8_prefill=True) in ( + None, + current_platform.fp8_dtype(), + ), ( + "Kimi-K3 with a plain fp8 KV cache needs a kv_b_proj that " + "consumes the fp8 gathered latent directly; this checkpoint's " + "kv_b_proj wants " + f"{_get_kv_b_proj_input_dtype(self.kv_b_proj, use_fp8_prefill=True)}." + ) + quant_method = ( self.quant_config.get_quant_method(self, prefix=self.layer_name) if self.quant_config @@ -1059,6 +1085,177 @@ def _decode_concat_cache( cos_sin_cache=cos_sin_cache, ) + def _compute_prefill_context( + self, + q: torch.Tensor, + attn_metadata: "MLACommonMetadata", + ) -> tuple[torch.Tensor, torch.Tensor]: + """Chunked-context prefill, K3-fused. Replaces the impl's version. + + Per chunk the impl gathers the paged latent, up-projects it, then casts + and concatenates K (and casts V) in two or three more launches. Here that + tail is one fused kernel per chunk -- ``fused_mla_kv_concat`` for a bf16 + query, ``fused_mla_kv_concat_quant_fp8`` when the query is fp8 -- reading + the strided ``kv_b_proj`` output in place and writing a contiguous key, so + only the gather and ``kv_b_proj`` remain. + + The impl's two input casts are gone as well: ``q`` already carries + ``prefill.q_data_type`` because the new-token epilogue quantized it, and + the gathered latent goes into ``kv_b_proj`` in whatever layout the gather + left it -- ``process_weights_after_loading`` checks once that + ``kv_b_proj`` accepts it, and its output is bf16 either way. + + The gathered ``k_pe`` is likewise used as-is (fp8 for a plain fp8 cache) + and needs no RoPE: it was rotated on the way in. + + Chunk partials are written straight into the accumulating context partial + when the prefill backend honors ``out``, so only the (64x smaller) lse is + copied per chunk. + + Decode context parallelism keeps using + ``impl._context_parallel_compute_prefill_context``; its extra allgather + and reorg are not fused here. + """ + prefill = attn_metadata.prefill + assert prefill is not None + prefill_backend = prefill.prefill_backend + assert prefill_backend is not None + chunked_context = prefill.chunked_context + assert chunked_context is not None + assert q.dtype == prefill.q_data_type, ( + "Kimi-K3 chunked context expects the new-token epilogue to have " + f"produced a {prefill.q_data_type} query; got {q.dtype}." + ) + + fp8_prefill = q.dtype == current_platform.fp8_dtype() + workspace = chunked_context.workspace + kv_cache = self._attn_read_kv_cache() + + def run_chunk( + chunk, out: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + self._gather_context_latent(chunk, kv_cache, prefill, fp8_prefill) + gathered = workspace[: chunk.num_context_tokens] + kv_c_normed = gathered[..., : self.kv_lora_rank] + kv_nope = self.kv_b_proj(kv_c_normed)[0].view( + -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim + ) + k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) + k_pe = gathered[..., self.kv_lora_rank :] + if fp8_prefill: + k, v = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v) + else: + k = fused_mla_kv_concat(k_nope, k_pe) + attn_output, attn_lse = prefill_backend.run_prefill_context_chunk( + chunk=chunk, q=q[chunk.token_slice], k=k, v=v, out=out + ) + assert out is None or attn_output.data_ptr() == out.data_ptr(), ( + f"{prefill_backend.get_name()} reports supports_out() but did not " + "write the context chunk into the `out` it was given." + ) + return attn_output, attn_lse + + chunks = chunked_context.chunks + if len(chunks) == 1 and not chunked_context.empty_token_slices: + # One chunk covering every prefill token: its partial *is* the context + # partial, so it needs neither an accumulator nor a copy. + return run_chunk(chunks[0]) + + # A backend honoring `out` writes each chunk's partial straight into the + # accumulator, so the per-chunk output copy disappears -- and because that + # contract fixes the trailing shape, the accumulator can be sized before + # any chunk runs. Otherwise the shape is only knowable from a real partial, + # so the first chunk runs ahead of the loop and is copied in. + writes_out = prefill_backend.supports_out() + if writes_out: + assert prefill.output_dtype is not None + output = torch.empty( + (q.shape[0], self.num_local_heads, self.v_head_dim), + dtype=prefill.output_dtype, + device=q.device, + ) + output_lse = torch.empty( + (self.num_local_heads, q.shape[0]), + dtype=torch.float32, + device=q.device, + ) + neutralize_empty_context_partials(chunked_context, output, output_lse) + else: + attn_output, attn_lse = run_chunk(chunks[0]) + output, output_lse = init_mla_context_partial( + chunked_context, attn_output, attn_lse, num_tokens=q.shape[0] + ) + accumulate_mla_context_chunk( + chunks[0], attn_output, attn_lse, output, output_lse + ) + chunks = chunks[1:] + + for chunk in chunks: + # A continuation chunk's leading tokens have to be merged with the + # partial already sitting there, so it cannot write in place. + out = ( + output[chunk.token_slice] + if writes_out and not chunk.is_continuation + else None + ) + attn_output, attn_lse = run_chunk(chunk, out=out) + accumulate_mla_context_chunk( + chunk, + attn_output, + attn_lse, + output, + output_lse, + output_written=out is not None, + ) + return output, output_lse + + def _gather_context_latent( + self, + chunk, + kv_cache: torch.Tensor, + prefill, + fp8_prefill: bool, + ) -> None: + """Gather one chunk's paged context latent into the workspace. + + Dispatched exactly as in ``impl._compute_prefill_context``: an fp8 query + reads the plain fp8 cache in its stored layout, anything else lands in the + workspace as the model dtype. + """ + workspace = prefill.chunked_context.workspace + toks = chunk.num_context_tokens + block_table = prefill.block_table[chunk.request_slice] + if self.kv_cache_dtype == "fp8_ds_mla": + ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache=kv_cache, + dst=workspace[:toks], + block_table=block_table, + workspace_starts=chunk.cu_seq_lens, + batch_size=chunk.num_requests, + seq_starts=chunk.starts, + ) + elif not fp8_prefill: + ops.gather_and_maybe_dequant_cache( + src_cache=kv_cache, + dst=workspace, + block_table=block_table, + cu_seq_lens=chunk.cu_seq_lens, + token_to_seq=chunk.token_to_seq, + num_tokens=toks, + kv_cache_dtype=self.kv_cache_dtype, + scale=self._k_scale, + seq_starts=chunk.starts, + ) + else: + ops.cp_gather_cache( + src_cache=kv_cache, + dst=workspace[:toks], + block_table=block_table, + cu_seq_lens=chunk.cu_seq_lens, + batch_size=chunk.num_requests, + seq_starts=chunk.starts, + ) + def _forward_prefill_fused( self, q: torch.Tensor, @@ -1073,8 +1270,9 @@ def _forward_prefill_fused( """Prefill using the fused key-concat + cache-insert kernel. Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one - fused kernel launch, dispatched by cache dtype. The chunked context - gather + online-softmax merge are delegated to the impl. + fused kernel launch, dispatched by cache dtype. Chunked context runs + through this layer's ``_compute_prefill_context``, except under DCP where + it is delegated to the impl. Supported configs (K3 fp8 policy): - bf16 cache -> bf16 prefill query @@ -1203,8 +1401,8 @@ def _forward_prefill_fused( ) ) else: - context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] - q, self._attn_read_kv_cache(), attn_metadata, self._k_scale + context_output, context_lse = self._compute_prefill_context( + q, attn_metadata ) compact_context_output = _reuse_consumed_query_for_context_output(q, out) compact_context_output.copy_(context_output[..., : self.v_head_dim]) diff --git a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py index 816da37153f5..24ae1a1a820c 100644 --- a/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py +++ b/vllm/models/kimi_k3/nvidia/ops/fused_mla_key_concat_kv_cache.py @@ -12,6 +12,9 @@ - ``fused_mla_qkv_quant_kv_cache_fp8_insert`` (fp8): additionally quantize ``q``/``k``/``v`` to E4M3 with ``q_scale`` / ``k_scale`` / ``v_scale`` (the cache shares ``k_scale``, as in ``concat_and_cache_mla``). +- ``fused_mla_kv_concat`` / ``fused_mla_kv_concat_quant_fp8`` (chunked context): + the same K concat (plus the fp8 K/V cast) without the cache insert, for context + chunks whose latent was gathered back out of the paged cache. The optional ``positions`` / ``cos_sin_cache`` pair enables GPT-J-style RoPE inside the epilogue. Omitting both keeps the K3 NoPE fast path. The kernels use @@ -156,6 +159,59 @@ def fused_mla_qkv_quant_kv_cache_fp8_insert( return q_fp8, k_fp8, v_fp8 +def _empty_full_key( + k_nope: torch.Tensor, k_pe: torch.Tensor, dtype: torch.dtype +) -> torch.Tensor: + num_tokens, num_heads, qk_nope_head_dim = k_nope.shape + return torch.empty( + (num_tokens, num_heads, qk_nope_head_dim + k_pe.shape[1]), + dtype=dtype, + device=k_nope.device, + ) + + +def fused_mla_kv_concat( + k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided + k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], same dtype as k_nope +) -> torch.Tensor: + """Concat ``k = [k_nope | k_pe]`` into a contiguous key, in one launch. + + The chunked-context counterpart of ``fused_mla_key_concat_kv_cache_insert``: + no query, no cache insert and no RoPE (the gathered ``k_pe`` is already + rotated). ``k_nope`` is a strided half of one ``kv_b_proj`` output and + ``k_pe`` a strided view of the gather workspace, so neither has to be made + contiguous first. + """ + k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1]) + k = _empty_full_key(k_nope, k_pe, k_nope.dtype) + if k.shape[0]: + torch.ops._C.fused_kimi_k3_mla_kv_concat(k_nope, k_pe, k) + return k + + +def fused_mla_kv_concat_quant_fp8( + k_nope: torch.Tensor, # [T, H, qk_nope_head_dim], may be strided + k_pe: torch.Tensor, # [T, rope] or [T, 1, rope], k_nope's dtype or fp8 + v: torch.Tensor, # [T, H, v_head_dim], may be strided +) -> tuple[torch.Tensor, torch.Tensor]: + """``fused_mla_kv_concat`` plus an fp8 cast of the key and of ``v``. + + ``k_pe`` may already be fp8: a plain fp8 cache is gathered without + dequantizing, and those bytes are copied through as-is. + + Returns contiguous ``(k_fp8, v_fp8)``. + """ + k_pe = k_pe.reshape(k_pe.shape[0], k_pe.shape[-1]) + fp8 = torch.float8_e4m3fn + k_fp8 = _empty_full_key(k_nope, k_pe, fp8) + v_fp8 = torch.empty(v.shape, dtype=fp8, device=v.device) + if k_fp8.shape[0]: + torch.ops._C.fused_kimi_k3_mla_kv_concat_quant_fp8( + k_nope, k_pe, v, k_fp8, v_fp8 + ) + return k_fp8, v_fp8 + + def fused_mla_decode_q_concat_kv_cache_insert( ql_nope: torch.Tensor, # [B, H, kv_lora_rank] (BMM1 output, absorbed q) q_pe: torch.Tensor, # [B, H, qk_rope_head_dim] diff --git a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py index 09726ec8cc03..3ba5a5a5fad3 100644 --- a/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/aiter_flash_attn.py @@ -106,7 +106,12 @@ def run_prefill_context_chunk( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: + assert out is None, ( + "AiterFlashAttnPrefillBackend does not report supports_out(), so it " + "is never given a context-chunk `out` to write into." + ) out, lse = self.flash_attn_varlen_func( q=q, k=k, diff --git a/vllm/v1/attention/backends/mla/prefill/base.py b/vllm/v1/attention/backends/mla/prefill/base.py index 986d0ca347cb..1b589cb295f4 100644 --- a/vllm/v1/attention/backends/mla/prefill/base.py +++ b/vllm/v1/attention/backends/mla/prefill/base.py @@ -74,13 +74,16 @@ def supports_quant_output(self, quant_key: "QuantKey") -> bool: return False def supports_out(self) -> bool: - """Whether `run_prefill_new_tokens` honors a caller-provided `out` - tensor of shape `[num_tokens, num_heads, v_head_dim]`, writing the - final result into it in place. + """Whether `run_prefill_new_tokens` and `run_prefill_context_chunk` honor + a caller-provided `out` tensor of shape + `[num_tokens, num_heads, v_head_dim]`, writing the result into it in place + and returning it. When True, callers may pass `out` and skip the post-hoc - slice/flatten/copy. False for backends that ignore `out` or emit a - padded (`qk_head_dim`) output. Overridden by backends that support it. + slice/flatten/copy -- and, for context chunks, size the accumulating + partial before running any chunk. False for backends that ignore `out` or + emit a padded (`qk_head_dim`) output. Overridden by backends that support + it. """ return False @@ -179,5 +182,6 @@ def run_prefill_context_chunk( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: raise NotImplementedError diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 97cb57c18a1e..8c65edd64cdc 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -457,6 +457,7 @@ def run_prefill_context_chunk( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: return self._flash_attn_varlen_diff_headdims( q=q, @@ -469,4 +470,5 @@ def run_prefill_context_chunk( softmax_scale=self.scale, causal=False, # Context is unmasked return_softmax_lse=True, + out=out, ) diff --git a/vllm/v1/attention/backends/mla/prefill/flashinfer.py b/vllm/v1/attention/backends/mla/prefill/flashinfer.py index 902aebbdde97..f2707aa970b4 100644 --- a/vllm/v1/attention/backends/mla/prefill/flashinfer.py +++ b/vllm/v1/attention/backends/mla/prefill/flashinfer.py @@ -228,11 +228,13 @@ def run_prefill_context_chunk( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: attn_out, lse = self._prefill_chunks[chunk.index].run( q=q, k=k, v=v, + out=out, return_lse=True, ) diff --git a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py index 7452300a033b..6985b0a70c24 100644 --- a/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/prefill/tokenspeed_mla.py @@ -166,11 +166,13 @@ def run_prefill_context_chunk( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: from tokenspeed_mla import tokenspeed_mla_prefill # See note in run_prefill_new_tokens — `v` is a split-view of `kv_nope` - # in `_compute_prefill_context` and arrives non-contiguous. + # in the generic `_compute_prefill_context` and arrives non-contiguous. + # (K3's fused fp8 K/V pack already hands over a contiguous `v`.) v = v.contiguous() attn_out, lse = tokenspeed_mla_prefill( @@ -187,6 +189,7 @@ def run_prefill_context_chunk( cum_seq_lens_q=chunk.query_start_loc, max_seq_len_q=chunk.max_query_len, enable_pdl=False, + out=out, ) # Convert from (q_len, num_heads) to (num_heads, q_len) diff --git a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py index b4f4181bdb8d..144a75aded70 100644 --- a/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py +++ b/vllm/v1/attention/backends/mla/prefill/trtllm_ragged.py @@ -147,16 +147,18 @@ def run_prefill_context_chunk( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, + out: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: from flashinfer.prefill import trtllm_ragged_attention_deepseek - out = torch.empty( - q.shape[0], - q.shape[1], - v.shape[2], - device=q.device, - dtype=self._prefill_metadata.output_dtype, - ) + if out is None: + out = torch.empty( + q.shape[0], + q.shape[1], + v.shape[2], + device=q.device, + dtype=self._prefill_metadata.output_dtype, + ) attn_out, lse = trtllm_ragged_attention_deepseek( query=q, From ee429fb0197b1317dbda812195f92c3d1369a266 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Tue, 11 Aug 2026 06:23:41 +0000 Subject: [PATCH 44/52] [Bugfix][MLA] Cast the gathered latent for a bf16 Kimi-K3 kv_b_proj The fused chunked-context loop replaced the impl's per-chunk cast of the gathered latent with a load-time check that `kv_b_proj` consumes fp8 directly. No stock K3 checkpoint satisfies that -- `kv_b_proj` is bf16 -- so every rank died at weight load under `--kv-cache-dtype fp8`: AssertionError: Kimi-K3 with a plain fp8 KV cache needs a kv_b_proj that consumes the fp8 gathered latent directly; this checkpoint's kv_b_proj wants torch.bfloat16. Cache dtype and `kv_b_proj` dtype are independent: a plain fp8 cache leaves the gather in fp8, but only an fp8-quantized `kv_b_proj` can take that. Restore the impl's cast in `run_chunk` and drop the check. It is a no-op `.to` when `kv_b_proj` already consumes the fp8 latent, so the fused path keeps its win; the fusion itself is untouched. The existing test missed this because `_KVBProj` was always built with `weight_dtype=q_data_type`, testing only the diagonal of the two dtypes, and its stand-in cast its own input so it accepted anything. It now rejects a non-bf16 input for a bf16 weight, as `F.linear` does, and the two dtypes are parametrized independently. Tests: pytest tests/models/kimi_k3/test_mla_prefill_context.py \ tests/v1/attention/test_mla_prefill_registry.py \ tests/v1/attention/test_mla_context_chunks.py \ tests/kernels/attention/test_kimi_k3_mla_fused_epilogue.py \ tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py -> 67 passed (63 before; +4 from the new parametrization) The two new bf16-kv_b_proj + fp8-cache cases fail without this change and pass with it. Model eval, Kimi-K3 2P1D PD disaggregation (prefill 2x TEP8, decode DEP16/TP1, EAGLE3-3, kv-cache-dtype fp8) on 8x4 B300 -- the configuration that could not load before: GSM8K 5-shot 0.9575 flexible-extract / 0.9575 strict-match (1319, 0 errors) OCRBench 0.892 +- 0.010 (1000, thinking effort high) AI assistance was used. Signed-off-by: Yongye Zhu Co-Authored-By: Claude Opus 5 (1M context) --- .../kimi_k3/test_mla_prefill_context.py | 33 +++++++++++++++---- vllm/models/kimi_k3/nvidia/mla.py | 33 +++++-------------- 2 files changed, 36 insertions(+), 30 deletions(-) diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py index b0166c2a31d0..6c6e4bf20e87 100644 --- a/tests/models/kimi_k3/test_mla_prefill_context.py +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -86,10 +86,15 @@ def run_prefill_context_chunk(self, *, chunk, q, k, v, out=None): class _KVBProj(torch.nn.Module): """Stand-in for the layer's ``kv_b_proj`` (returns an (out, bias) tuple). - With a plain fp8 cache the gathered latent arrives fp8 and K3 feeds it in - without a cast, so ``kv_b_proj`` must consume it directly -- mimicked here by - an fp8 weight whose apply dequantizes, like the fp8 linear methods. Either - weight dtype produces a bf16 output. + Enforces the same input contract as the real linear methods, because that is + what decides whether the gathered latent needs a cast: + + * an fp8 weight consumes the fp8 latent directly and dequantizes internally, + so it takes fp8 or bf16; + * a bf16 weight -- what a stock K3 checkpoint carries -- is a plain + ``F.linear`` and rejects anything but bf16, exactly as torch does. + + Either weight dtype produces a bf16 output. """ def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: @@ -106,6 +111,11 @@ def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: self.register_buffer("weight", weight.to(weight_dtype)) def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: + if self.weight.dtype == torch.bfloat16 and x.dtype != torch.bfloat16: + raise RuntimeError( + "a bfloat16 kv_b_proj cannot consume the gathered latent as " + f"{x.dtype}; it must be cast first" + ) return torch.nn.functional.linear( x.to(torch.bfloat16), self.weight.to(torch.bfloat16) ), None @@ -190,10 +200,19 @@ def _build_prefill_metadata( @pytest.mark.parametrize("honors_out", [False, True], ids=["copy_out", "writes_out"]) @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8"]) +@pytest.mark.parametrize( + "kv_b_proj_quantized", [True, False], ids=["fp8_kv_b_proj", "bf16_kv_b_proj"] +) @torch.inference_mode() def test_fused_context_matches_generic_impl( - kv_cache_dtype: str, honors_out: bool + kv_b_proj_quantized: bool, kv_cache_dtype: str, honors_out: bool ) -> None: + """Cache dtype and ``kv_b_proj`` dtype vary independently. + + A stock K3 checkpoint pairs a bf16 ``kv_b_proj`` with an fp8 cache, so the + fused loop cannot assume the gathered latent is already in the dtype + ``kv_b_proj`` accepts. + """ torch.manual_seed(0) device = torch.device("cuda") fp8 = current_platform.fp8_dtype() @@ -203,7 +222,9 @@ def test_fused_context_matches_generic_impl( q_data_type = fp8 if quantized else torch.bfloat16 workspace_dtype = q_data_type - kv_b_proj = _KVBProj(device, weight_dtype=q_data_type) + kv_b_proj = _KVBProj( + device, weight_dtype=fp8 if kv_b_proj_quantized else torch.bfloat16 + ) k_scale = torch.ones(1, dtype=torch.float32, device=device) backend_fused = _RecordingPrefillBackend(honors_out=honors_out) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index cd0996bb068b..0d8243950f61 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -727,25 +727,6 @@ def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: dim=0, ) - # `_compute_prefill_context` hands the gathered context latent straight to - # `kv_b_proj`. With a plain fp8 cache the gather leaves it fp8, so - # `kv_b_proj` has to take an fp8 input; check that here rather than - # casting per chunk (the generic impl casts instead). fp8_ds_mla is - # exempt: it up-converts into a bf16 workspace. - if ( - is_quantized_kv_cache(self.kv_cache_dtype) - and self.kv_cache_dtype != "fp8_ds_mla" - ): - assert _get_kv_b_proj_input_dtype(self.kv_b_proj, use_fp8_prefill=True) in ( - None, - current_platform.fp8_dtype(), - ), ( - "Kimi-K3 with a plain fp8 KV cache needs a kv_b_proj that " - "consumes the fp8 gathered latent directly; this checkpoint's " - "kv_b_proj wants " - f"{_get_kv_b_proj_input_dtype(self.kv_b_proj, use_fp8_prefill=True)}." - ) - quant_method = ( self.quant_config.get_quant_method(self, prefix=self.layer_name) if self.quant_config @@ -1099,11 +1080,12 @@ def _compute_prefill_context( the strided ``kv_b_proj`` output in place and writing a contiguous key, so only the gather and ``kv_b_proj`` remain. - The impl's two input casts are gone as well: ``q`` already carries - ``prefill.q_data_type`` because the new-token epilogue quantized it, and - the gathered latent goes into ``kv_b_proj`` in whatever layout the gather - left it -- ``process_weights_after_loading`` checks once that - ``kv_b_proj`` accepts it, and its output is bf16 either way. + The impl's query cast is gone as well: ``q`` already carries + ``prefill.q_data_type`` because the new-token epilogue quantized it. The + gathered latent still gets the impl's cast to whatever ``kv_b_proj`` + consumes -- free (a no-op ``.to``) for a checkpoint whose ``kv_b_proj`` + takes the fp8 latent directly, and required for a bf16 one, which is + what a stock K3 checkpoint carries. Its output is bf16 either way. The gathered ``k_pe`` is likewise used as-is (fp8 for a plain fp8 cache) and needs no RoPE: it was rotated on the way in. @@ -1130,6 +1112,7 @@ def _compute_prefill_context( fp8_prefill = q.dtype == current_platform.fp8_dtype() workspace = chunked_context.workspace kv_cache = self._attn_read_kv_cache() + kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill) def run_chunk( chunk, out: torch.Tensor | None = None @@ -1137,6 +1120,8 @@ def run_chunk( self._gather_context_latent(chunk, kv_cache, prefill, fp8_prefill) gathered = workspace[: chunk.num_context_tokens] kv_c_normed = gathered[..., : self.kv_lora_rank] + if kv_b_proj_input_dtype is not None: + kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype) kv_nope = self.kv_b_proj(kv_c_normed)[0].view( -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim ) From 865148ab1caec80b7625241bfe2fd101df76c72d Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Thu, 13 Aug 2026 18:38:27 -0700 Subject: [PATCH 45/52] [Kimi-K3][DCP] Publish prefill KV directly in MLA layout Signed-off-by: Summer Yang --- .buildkite/test_areas/distributed.yaml | 8 +- .../dcp_utils/dcp_direct_kv_gather.cu | 171 ++++---- .../test_dcp_direct_a2a_lse_reduce.py | 410 ++++++++++++++++-- tests/v1/attention/test_mla_context_chunks.py | 153 ++++++- .../layers/attention/mla_attention.py | 224 +++++++--- .../layers/attention/sparse_mla_attention.py | 32 +- vllm/models/kimi_k3/nvidia/mla.py | 5 +- vllm/v1/attention/ops/dcp_utils.py | 136 ++++-- 8 files changed, 921 insertions(+), 218 deletions(-) diff --git a/.buildkite/test_areas/distributed.yaml b/.buildkite/test_areas/distributed.yaml index 96cf519fddf3..0a87ba8e316a 100644 --- a/.buildkite/test_areas/distributed.yaml +++ b/.buildkite/test_areas/distributed.yaml @@ -182,7 +182,7 @@ steps: - label: Distributed Tests (8xH100) key: distributed-tests-8xh100 - timeout_in_minutes: 20 + timeout_in_minutes: 30 device: h100 num_devices: 8 working_dir: "/vllm-workspace/tests" @@ -194,6 +194,11 @@ steps: - vllm/v1/engine/llm_engine.py - vllm/v1/executor/uniproc_executor.py - vllm/v1/worker/gpu_worker.py + - vllm/v1/attention/ops/dcp_utils.py + - vllm/model_executor/layers/attention/mla_attention.py + - vllm/model_executor/layers/attention/sparse_mla_attention.py + - csrc/libtorch_stable/attention/dcp_utils/ + - tests/distributed/test_dcp_direct_a2a_lse_reduce.py - tests/distributed/test_mnnvl_alltoall.py commands: @@ -201,6 +206,7 @@ steps: - export NCCL_CUMEM_HOST_ENABLE=0 # test with torchrun tp=2 and dp=4 with ep - torchrun --nproc-per-node=8 ../examples/features/torchrun/torchrun_dp_example_offline.py --tp-size=2 --pp-size=1 --dp-size=4 --enable-ep + - pytest -v -s distributed/test_dcp_direct_a2a_lse_reduce.py - label: Distributed Tests (4xA100) key: distributed-tests-4xa100 diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu index 809f759b6acc..801803c09c9e 100644 --- a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu @@ -20,28 +20,53 @@ using vllm::direct_dcp::wait_for_epoch; constexpr int kThreads = 256; // KV chunks need many blocks in flight to saturate the fabric. constexpr int64_t kMaxMulticastBlocks = 128; -constexpr int64_t kMaxCopyBlocks = 1024; -// Multicast each rank's KV slice and completion epoch to every replica. +// Multicast each rank's valid local rows directly into compact, request-major +// kv_c and k_pe planes. dst_rows maps every padded local input row to its final +// output row, or -1 for padding. Destination rows are disjoint across source +// ranks, so all ranks can publish concurrently without atomics. __global__ void direct_dcp_kv_gather_multimem_kernel( - const uint4* local_kv, uint4* mc_kv, uint32_t* mc_signal, - const uint32_t* received_signal, const int64_t* epoch_ptr, - uint32_t* completion, int64_t world_size, int64_t rank, int64_t slice_bytes, - int64_t slot_stride_bytes) { + const uint4* local_kv, const int32_t* dst_rows, uint4* mc_kv, + uint32_t* mc_signal, const uint32_t* received_signal, + const int64_t* epoch_ptr, uint32_t* completion, int64_t world_size, + int64_t rank, int64_t num_tokens, int64_t items_per_row, + int64_t kv_c_items_per_row, int64_t output_tokens, + int64_t max_gathered_tokens, int64_t buffer_slot, + int64_t slot_stride_items) { uint32_t epoch = static_cast(epoch_ptr[0]); - int64_t buffer_slot = static_cast(epoch & 1u); - int64_t item_count = slice_bytes / sizeof(uint4); - mc_kv += buffer_slot * slot_stride_bytes / sizeof(uint4); - mc_kv += rank * item_count; + mc_kv += buffer_slot * slot_stride_items; + int64_t item_count = num_tokens * items_per_row; int64_t item_stride = static_cast(gridDim.x) * blockDim.x; for (int64_t item = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; item < item_count; item += item_stride) { - multimem_store_16(mc_kv + item, local_kv[item]); + int64_t src_row = item / items_per_row; + int32_t dst_row = dst_rows[src_row]; + if (dst_row < 0) { + continue; + } + if (dst_row >= output_tokens) { + printf( + "direct DCP final-layout kv-gather destination out of bounds " + "source=%lld dst=%d output_tokens=%lld\n", + static_cast(rank), dst_row, + static_cast(output_tokens)); + asm volatile("trap;"); + } + int64_t row_item = item - src_row * items_per_row; + int64_t dst_item; + if (row_item < kv_c_items_per_row) { + dst_item = static_cast(dst_row) * kv_c_items_per_row + row_item; + } else { + int64_t k_pe_items_per_row = items_per_row - kv_c_items_per_row; + dst_item = max_gathered_tokens * kv_c_items_per_row + + static_cast(dst_row) * k_pe_items_per_row + row_item - + kv_c_items_per_row; + } + multimem_store_16(mc_kv + dst_item, local_kv[item]); } - // Publish all multicast writes before incrementing completion. __threadfence_system(); __syncthreads(); if (threadIdx.x != 0) { @@ -56,59 +81,57 @@ __global__ void direct_dcp_kv_gather_multimem_kernel( multimem_store_release_system(mc_signal + buffer_slot * world_size + rank, epoch); - for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { int64_t signal_item = buffer_slot * world_size + source_rank; if (!wait_for_epoch(received_signal + signal_item, epoch)) { - printf("direct DCP kv-gather multimem timeout source=%lld epoch=%u\n", + printf("direct DCP final-layout kv-gather timeout source=%lld epoch=%u\n", static_cast(source_rank), epoch); asm volatile("trap;"); } } } -// Materialize the acquired slot into the caller's workspace. -template -__global__ void materialize_kv_gather_kernel(const copy_t* received_kv, - const int64_t* epoch_ptr, - copy_t* gathered_kv, - int64_t gathered_item_count, - int64_t slot_stride_items) { - uint32_t epoch = static_cast(epoch_ptr[0]); - int64_t buffer_slot = static_cast(epoch & 1u); - received_kv += buffer_slot * slot_stride_items; - int64_t item = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - int64_t item_stride = static_cast(gridDim.x) * blockDim.x; - for (; item < gathered_item_count; item += item_stride) { - gathered_kv[item] = received_kv[item]; - } -} - -void direct_dcp_kv_gather( - const torch::stable::Tensor& local_kv, torch::stable::Tensor& received_kv, - torch::stable::Tensor& received_signal, torch::stable::Tensor& completion, - torch::stable::Tensor& epoch, torch::stable::Tensor& gathered_kv, - int64_t world_size, int64_t rank, int64_t max_gathered_tokens, - int64_t kv_mc_ptr, int64_t signal_mc_ptr) { +void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, + const torch::stable::Tensor& dst_rows, + torch::stable::Tensor& received_kv, + torch::stable::Tensor& received_signal, + torch::stable::Tensor& completion, + torch::stable::Tensor& epoch, int64_t output_tokens, + int64_t plane_split_dim, int64_t buffer_slot, + int64_t world_size, int64_t rank, + int64_t max_gathered_tokens, int64_t kv_mc_ptr, + int64_t signal_mc_ptr) { using torch::headeronly::ScalarType; STD_TORCH_CHECK(local_kv.is_cuda(), "local kv must be a CUDA tensor"); ScalarType dtype = local_kv.scalar_type(); STD_TORCH_CHECK(dtype == ScalarType::Half || dtype == ScalarType::BFloat16 || dtype == ScalarType::Float8_e4m3fn, - "direct DCP kv-gather only supports FP16, BF16, and FP8"); + "direct DCP final-layout kv-gather only supports FP16, " + "BF16, and FP8"); STD_TORCH_CHECK(local_kv.dim() == 2 && local_kv.is_contiguous(), "local kv must be a contiguous [T,D] tensor"); + STD_TORCH_CHECK(dst_rows.is_cuda() && dst_rows.is_contiguous() && + dst_rows.scalar_type() == ScalarType::Int && + dst_rows.dim() == 1 && + dst_rows.numel() == local_kv.size(0), + "final-layout destination rows must be CUDA int32 [T]"); STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1"); STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank"); + STD_TORCH_CHECK(buffer_slot == 0 || buffer_slot == 1, + "final-layout buffer slot must be 0 or 1"); + STD_TORCH_CHECK(output_tokens > 0 && output_tokens <= max_gathered_tokens, + "final-layout output exceeds symmetric buffer capacity"); int64_t num_tokens = local_kv.size(0); int64_t token_dim = local_kv.size(1); int64_t element_size = local_kv.element_size(); STD_TORCH_CHECK(num_tokens > 0 && token_dim > 0, "local kv dimensions must be positive"); + STD_TORCH_CHECK(plane_split_dim > 0 && plane_split_dim < token_dim, + "final-layout plane split must be within the token row"); STD_TORCH_CHECK(num_tokens * world_size <= max_gathered_tokens, - "gathered kv exceeds symmetric kv-gather buffer capacity"); + "padded gathered kv exceeds symmetric buffer capacity"); STD_TORCH_CHECK(received_kv.is_cuda() && received_kv.scalar_type() == dtype && received_kv.is_contiguous() && received_kv.dim() == 3 && @@ -130,65 +153,57 @@ void direct_dcp_kv_gather( epoch.scalar_type() == ScalarType::Long && epoch.numel() == 1, "epoch must be a one-element CUDA int64 tensor"); - STD_TORCH_CHECK(gathered_kv.is_cuda() && gathered_kv.scalar_type() == dtype && - gathered_kv.is_contiguous() && gathered_kv.dim() == 2 && - gathered_kv.size(0) == num_tokens * world_size && - gathered_kv.size(1) == token_dim, - "gathered kv must be contiguous with shape [world_size*T,D]"); + int64_t device_index = local_kv.get_device_index(); - STD_TORCH_CHECK( - received_kv.get_device_index() == device_index && - received_signal.get_device_index() == device_index && - completion.get_device_index() == device_index && - epoch.get_device_index() == device_index && - gathered_kv.get_device_index() == device_index, - "direct DCP kv-gather tensors must be on the same CUDA device"); - - int64_t slice_bytes = num_tokens * token_dim * element_size; - int64_t slot_stride_bytes = max_gathered_tokens * token_dim * element_size; + STD_TORCH_CHECK(dst_rows.get_device_index() == device_index && + received_kv.get_device_index() == device_index && + received_signal.get_device_index() == device_index && + completion.get_device_index() == device_index && + epoch.get_device_index() == device_index, + "direct DCP final-layout tensors must share a CUDA device"); + + int64_t row_bytes = token_dim * element_size; + int64_t kv_c_row_bytes = plane_split_dim * element_size; + int64_t k_pe_row_bytes = row_bytes - kv_c_row_bytes; + int64_t slot_stride_bytes = max_gathered_tokens * row_bytes; bool vectorized = reinterpret_cast(local_kv.data_ptr()) % alignof(uint4) == 0 && reinterpret_cast(received_kv.data_ptr()) % alignof(uint4) == 0 && - reinterpret_cast(gathered_kv.data_ptr()) % alignof(uint4) == - 0 && - slice_bytes % sizeof(uint4) == 0 && + row_bytes % sizeof(uint4) == 0 && kv_c_row_bytes % sizeof(uint4) == 0 && + k_pe_row_bytes % sizeof(uint4) == 0 && slot_stride_bytes % sizeof(uint4) == 0; - STD_TORCH_CHECK( - vectorized, - "direct DCP kv-gather requires 16-byte-aligned pointers and strides"); + STD_TORCH_CHECK(vectorized, + "direct DCP final-layout kv-gather requires 16-byte-aligned " + "planes, rows, and pointers"); STD_TORCH_CHECK(kv_mc_ptr != 0 && signal_mc_ptr != 0, - "direct DCP kv-gather requires multicast pointers"); + "direct DCP final-layout kv-gather requires multicast " + "pointers"); const torch::stable::accelerator::DeviceGuard device_guard(device_index); cudaStream_t stream = get_current_cuda_stream(); increment_epoch_kernel<<<1, 1, 0, stream>>>( epoch.mutable_data_ptr()); - check_cuda_launch("direct DCP kv-gather"); + check_cuda_launch("direct DCP final-layout kv-gather"); - int64_t item_count = slice_bytes / sizeof(uint4); + int64_t items_per_row = row_bytes / sizeof(uint4); + int64_t kv_c_items_per_row = kv_c_row_bytes / sizeof(uint4); + int64_t item_count = num_tokens * items_per_row; int64_t blocks = (item_count + kThreads - 1) / kThreads; blocks = blocks < kMaxMulticastBlocks ? blocks : kMaxMulticastBlocks; direct_dcp_kv_gather_multimem_kernel<<>>( reinterpret_cast(local_kv.data_ptr()), + dst_rows.const_data_ptr(), reinterpret_cast(static_cast(kv_mc_ptr)), reinterpret_cast(static_cast(signal_mc_ptr)), reinterpret_cast( received_signal.const_data_ptr()), epoch.const_data_ptr(), reinterpret_cast(completion.mutable_data_ptr()), - world_size, rank, slice_bytes, slot_stride_bytes); - check_cuda_launch("direct DCP kv-gather"); - - int64_t gathered_item_count = world_size * item_count; - int64_t copy_blocks = (gathered_item_count + kThreads - 1) / kThreads; - copy_blocks = copy_blocks < kMaxCopyBlocks ? copy_blocks : kMaxCopyBlocks; - materialize_kv_gather_kernel<<>>( - reinterpret_cast(received_kv.data_ptr()), - epoch.const_data_ptr(), - reinterpret_cast(gathered_kv.mutable_data_ptr()), - gathered_item_count, slot_stride_bytes / sizeof(uint4)); - check_cuda_launch("direct DCP kv-gather"); + world_size, rank, num_tokens, items_per_row, kv_c_items_per_row, + output_tokens, max_gathered_tokens, buffer_slot, + slot_stride_bytes / sizeof(uint4)); + check_cuda_launch("direct DCP final-layout kv-gather"); } } // namespace @@ -196,9 +211,9 @@ void direct_dcp_kv_gather( STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_kv_gather_ops) { direct_dcp_kv_gather_ops.def( "direct_dcp_kv_gather(" - "Tensor local_kv, Tensor! received_kv, Tensor! received_signal, " - "Tensor! completion, " - "Tensor! epoch, Tensor! gathered_kv, " + "Tensor local_kv, Tensor dst_rows, Tensor! received_kv, " + "Tensor! received_signal, Tensor! completion, Tensor! epoch, " + "int output_tokens, int plane_split_dim, int buffer_slot, " "int world_size, int rank, int max_gathered_tokens, " "int kv_mc_ptr, int signal_mc_ptr) -> ()"); } diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index 9f6c9c56c282..22dc9dbc63a3 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -3,6 +3,7 @@ """Tests for direct symmetric-memory DCP collectives.""" import functools +import time from unittest.mock import MagicMock import multiprocess as mp @@ -15,6 +16,7 @@ from vllm.utils.system_utils import update_environment_variables mp.set_start_method("spawn", force=True) +pytestmark = pytest.mark.skip_global_cleanup def _has_multicast_support() -> bool: @@ -115,8 +117,9 @@ def _distributed_run(fn, world_size: int, extra_env: dict[str, str]) -> None: processes.append(process) process.start() + deadline = time.monotonic() + 120 for process in processes: - process.join(timeout=120) + process.join(timeout=max(0, deadline - time.monotonic())) for process in processes: if process.is_alive(): @@ -222,7 +225,13 @@ def test_kv_gather_env_disabled_returns_none(self, monkeypatch): monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1") dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear() workspace = dcp_utils.get_direct_dcp_kv_gather_workspace( - _FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1 + _FakeGroupCoordinator(), + torch.device("cpu"), + 64, + 576, + 512, + torch.bfloat16, + 1, ) assert workspace is None @@ -240,7 +249,13 @@ def test_kv_gather_flag_is_independent(self, monkeypatch): ) result = dcp_utils.get_direct_dcp_kv_gather_workspace( - _FakeGroupCoordinator(), torch.device("cpu"), 64, 576, torch.bfloat16, 1 + _FakeGroupCoordinator(), + torch.device("cpu"), + 64, + 576, + 512, + torch.bfloat16, + 1, ) assert result is workspace @@ -256,7 +271,7 @@ def test_kv_gather_flag_is_independent(self, monkeypatch): ( "VLLM_USE_DIRECT_DCP_KV_GATHER", "get_direct_dcp_kv_gather_workspace", - (64, 576, torch.bfloat16, 1), + (64, 576, 512, torch.bfloat16, 1), ), ], ) @@ -284,17 +299,38 @@ def test_gather_requires_multicast( def test_kv_gather_rejects_invalid_workspace_geometry(self): with pytest.raises(ValueError, match="ubatch"): dcp_utils.DirectDCPKVGatherWorkspace( - None, torch.device("cpu"), 64, 576, num_ubatches=0 + None, torch.device("cpu"), 64, 576, 512, num_ubatches=0 ) with pytest.raises(ValueError, match="divide evenly"): dcp_utils.DirectDCPKVGatherWorkspace( - _FakeProcessGroup(), torch.device("cpu"), 63, 576 + _FakeProcessGroup(), torch.device("cpu"), 63, 576, 512 ) with pytest.raises(ValueError, match="16-byte"): dcp_utils.DirectDCPKVGatherWorkspace( - _FakeProcessGroup(), torch.device("cpu"), 64, 3 + _FakeProcessGroup(), torch.device("cpu"), 64, 16, 4 ) + @pytest.mark.parametrize( + ("token_dim", "plane_split_dim", "dtype", "supported"), + [ + (576, 512, torch.bfloat16, True), + (576, 512, torch.float16, True), + (576, 512, torch.float8_e4m3fn, True), + (16, 8, torch.bfloat16, True), + (16, 4, torch.bfloat16, False), + (24, 16, torch.float8_e4m3fn, False), + (16, 0, torch.bfloat16, False), + (16, 16, torch.bfloat16, False), + ], + ) + def test_kv_gather_requires_each_plane_aligned( + self, token_dim, plane_split_dim, dtype, supported + ): + assert ( + dcp_utils._kv_gather_layout_supported(token_dim, plane_split_dim, dtype) + is supported + ) + def test_q_gather_rejects_invalid_workspace_geometry(self): with pytest.raises(ValueError, match="ubatch"): dcp_utils.DirectDCPQGatherWorkspace( @@ -365,13 +401,29 @@ def test_mla_dcp_manager_selects_direct_backends(monkeypatch): is_lse_base_on_e=False, use_pcp=False, ) - workspace = torch.empty(96, 8) - assert manager.query_gather == direct_query.gather - manager.init_kv_gather(workspace, 64) - gathered_kv, local_kv = torch.empty(4, 8), torch.empty(2, 8) - manager.kv_gather(gathered_kv, local_kv) - direct_kv.gather.assert_called_once_with(gathered_kv, local_kv) + assert manager.init_kv_gather(64, 16, 8, torch.bfloat16) + dcp_manager.get_direct_dcp_kv_gather_workspace.assert_called_once_with( + group, + torch.device("cpu"), + 64, + 16, + 8, + torch.bfloat16, + 1, + ) + local_kv = torch.empty(2, 8) + dst_rows = torch.tensor([0, 1], dtype=torch.int32) + compact_kv = (torch.empty(2, 1, 6), torch.empty(2, 1, 2)) + direct_kv.gather.return_value = compact_kv + assert manager.use_direct_kv_gather + assert manager.direct_kv_gather(local_kv, dst_rows, 2, 1) is compact_kv + direct_kv.gather.assert_called_once_with( + local_kv, + dst_rows, + 2, + 1, + ) output, lse = torch.empty(1), torch.empty(1) seq_lens = torch.ones(1, dtype=torch.int32) query_start_loc = torch.tensor([0, 1], dtype=torch.int32) @@ -424,8 +476,7 @@ def test_mla_dcp_manager_selects_fallback_backends(monkeypatch): all_gather = MagicMock() monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", all_gather) - workspace = torch.empty(96, 8) - manager.init_kv_gather(workspace, 64) + assert not manager.init_kv_gather(64, 16, 8, torch.bfloat16) output, local = torch.empty(4, 8), torch.empty(2, 8) manager.kv_gather(output, local) all_gather.assert_called_once_with(output, local, group=group.device_group) @@ -528,7 +579,8 @@ def test_mla_chunk_workspace_honors_configured_token_limit(monkeypatch): MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config) -def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): +@pytest.mark.parametrize("use_direct", [False, True]) +def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch, use_direct): import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla monkeypatch.setattr( @@ -548,7 +600,7 @@ def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): ) manager = object.__new__(dcp_utils.MLADCPManager) - manager.init_kv_gather = MagicMock() + manager.init_kv_gather = MagicMock(return_value=use_direct) layer = MagicMock(dcp_manager=manager) config = MagicMock() config.model_config.dtype = torch.bfloat16 @@ -571,9 +623,18 @@ def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): assert builder.dcp_manager is manager manager.init_kv_gather.assert_called_once_with( - builder.chunked_prefill_workspace, builder.chunked_prefill_workspace_size, + 12, + 8, + torch.bfloat16, + ) + local_rows = builder.chunked_prefill_workspace_size // 2 + expected_rows = ( + local_rows + if use_direct + else builder.chunked_prefill_workspace_size + local_rows ) + assert builder.chunked_prefill_workspace.shape == (expected_rows, 12) def test_sparse_mla_workspace_preserves_non_dcp_size(): @@ -770,49 +831,306 @@ def _distributed_direct_kv_gather_worker(env: dict[str, str]) -> None: try: rank = dist.get_rank() world_size = dist.get_world_size() - token_dim = 576 - max_gathered_tokens = 128 * world_size + token_dim, plane_split_dim = 576, 512 + max_gathered_tokens = 64 active_ubatch = [0] dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] - for dtype_idx, dtype_name in enumerate(("bfloat16", "float16")): + full_layout = ( + [4, 4, 6], + [4, 0, 0], + [[8, 8, 7, 6], [4, 2, 2, 2], [6, 6, 5, 4]], + ) + alternate_layout = ( + [4, 4, 6], + [4, 0, 0], + [[8, 8, 8, 6], [3, 2, 2, 2], [6, 6, 5, 4]], + ) + small_layout = ([4], [4], [[8, 8, 7, 6]]) + + def layout_maps(layout) -> tuple[list[list[int]], int]: + padded_lens, local_starts, context_lens = layout + maps: list[list[int]] = [[] for _ in range(world_size)] + output_start = 0 + for padded_len, local_start, request_lens in zip( + padded_lens, local_starts, context_lens, strict=True + ): + valid_lens = [ + min(max(0, length - local_start), padded_len) + for length in request_lens + ] + for source_rank, valid_len in enumerate(valid_lens): + rank_start = output_start + sum(valid_lens[:source_rank]) + maps[source_rank].extend(range(rank_start, rank_start + valid_len)) + maps[source_rank].extend([-1] * (padded_len - valid_len)) + output_start += sum(valid_lens) + valid_rows = sorted(row for rows in maps for row in rows if row >= 0) + assert valid_rows == list(range(output_start)) + return maps, output_start + + def make_local( + source_rank: int, + iteration: int, + num_rows: int, + dtype: torch.dtype, + ) -> torch.Tensor: + values = torch.arange( + num_rows * token_dim, + dtype=torch.int32, + device=device, + ) + values = (values + source_rank * 7 + iteration * 13).remainder(29) - 14 + return values.view(num_rows, token_dim).to(dtype) + + def expected_planes(layout, iteration: int, dtype: torch.dtype): + maps, output_tokens = layout_maps(layout) + output = torch.empty(output_tokens, token_dim, dtype=dtype, device=device) + covered = torch.zeros(output_tokens, dtype=torch.bool, device=device) + for source_rank, rows in enumerate(maps): + source = make_local(source_rank, iteration, len(rows), dtype) + dst_rows = torch.tensor(rows, dtype=torch.int32, device=device) + valid = dst_rows >= 0 + assert not torch.any(covered[dst_rows[valid]]) + output[dst_rows[valid]] = source[valid] + covered[dst_rows[valid]] = True + assert torch.all(covered) + return ( + output[:, :plane_split_dim].unsqueeze(1), + output[:, plane_split_dim:].unsqueeze(1), + ) + + def publish(workspace, dtype, iteration: int, buffer_slot: int, layout): + maps, output_tokens = layout_maps(layout) + actual = workspace.gather( + make_local(rank, iteration, len(maps[rank]), dtype), + torch.tensor(maps[rank], dtype=torch.int32, device=device), + output_tokens, + buffer_slot, + ) + return actual, expected_planes(layout, iteration, dtype) + + def assert_planes(actual, expected) -> None: + for actual_plane, expected_plane in zip(actual, expected, strict=True): + assert torch.equal( + actual_plane.view(torch.uint8), expected_plane.view(torch.uint8) + ) + + def check_case( + workspace: dcp_utils.DirectDCPKVGatherWorkspace, + dtype: torch.dtype, + iteration: int, + ubatch: int, + buffer_slot: int, + layout, + ) -> None: + maps, output_tokens = layout_maps(layout) + local_kv = make_local(rank, iteration, len(maps[rank]), dtype) + dst_rows = torch.tensor(maps[rank], dtype=torch.int32, device=device) + active_ubatch[0] = ubatch + + other_ubatch = 1 - ubatch + torch.accelerator.synchronize() + other_ubatch_before = workspace.received_kv[other_ubatch].clone() + inactive_slot_before = workspace.received_kv[ + ubatch, 1 - buffer_slot + ].clone() + slot = workspace.received_kv[ubatch, buffer_slot].view(-1) + kv_c_capacity = max_gathered_tokens * plane_split_dim + kv_c_storage = slot[:kv_c_capacity].view( + max_gathered_tokens, plane_split_dim + ) + k_pe_storage = slot[kv_c_capacity:].view( + max_gathered_tokens, token_dim - plane_split_dim + ) + kv_c_tail_before = kv_c_storage[output_tokens:].clone() + k_pe_tail_before = k_pe_storage[output_tokens:].clone() + epochs_before = workspace.epoch.clone() + torch.accelerator.synchronize() + + kv_c, k_pe = workspace.gather( + local_kv, + dst_rows, + output_tokens, + buffer_slot, + ) + torch.accelerator.synchronize() + + expected_kv_c, expected_k_pe = expected_planes(layout, iteration, dtype) + assert kv_c.shape == expected_kv_c.shape + assert k_pe.shape == expected_k_pe.shape + assert kv_c.dtype == k_pe.dtype == dtype + assert kv_c.is_contiguous() and k_pe.is_contiguous() + assert torch.equal(kv_c.view(torch.uint8), expected_kv_c.view(torch.uint8)) + assert torch.equal(k_pe.view(torch.uint8), expected_k_pe.view(torch.uint8)) + assert ( + kv_c.data_ptr() == workspace.received_kv[ubatch, buffer_slot].data_ptr() + ) + assert k_pe.data_ptr() == ( + kv_c.data_ptr() + kv_c_capacity * workspace.received_kv.element_size() + ) + assert int(workspace.epoch[ubatch].item()) == ( + int(epochs_before[ubatch].item()) + 1 + ) + assert int(workspace.epoch[other_ubatch].item()) == int( + epochs_before[other_ubatch].item() + ) + assert not torch.count_nonzero(workspace.completion) + assert torch.equal( + workspace.received_kv[other_ubatch].view(torch.uint8), + other_ubatch_before.view(torch.uint8), + ) + assert torch.equal( + workspace.received_kv[ubatch, 1 - buffer_slot].view(torch.uint8), + inactive_slot_before.view(torch.uint8), + ) + assert torch.equal( + kv_c_storage[output_tokens:].view(torch.uint8), + kv_c_tail_before.view(torch.uint8), + ) + assert torch.equal( + k_pe_storage[output_tokens:].view(torch.uint8), + k_pe_tail_before.view(torch.uint8), + ) + # The gather itself synchronizes publication, but a faster rank + # may otherwise start the next test case and multicast into what + # a slower rank is still validating as this case's inactive slot. + # Production keeps this ordering through same-stream attention + # consumption followed by K3's TP-wide rendezvous; the test needs + # an explicit rank barrier around its host-side assertions. + dist.barrier() + + # Exercise the ownership contract once in depth for BF16, then retain + # one byte-exact kernel smoke for each additional supported dtype. + eager_cases = { + "bfloat16": ( + (0, 0, 0, full_layout), + (1, 0, 1, alternate_layout), + (2, 0, 0, small_layout), + ), + "float16": ((3, 0, 0, full_layout),), + "float8_e4m3fn": ((4, 0, 0, full_layout),), + } + for dtype_name, cases in eager_cases.items(): dtype = _dtype_from_name(dtype_name) workspace = dcp_utils.DirectDCPKVGatherWorkspace( dist.group.WORLD, device, max_gathered_tokens, token_dim, + plane_split_dim, dtype, num_ubatches=2, ) + for args in cases: + check_case(workspace, dtype, *args) - # Use disjoint slices of one persistent chunked-context workspace. - storage = torch.zeros( - (world_size + 1) * 128, token_dim, device=device, dtype=dtype - ) - for iteration, num_tokens in enumerate((1, 128, 17)): - generator = torch.Generator(device=device) - generator.manual_seed(7000 + rank * 101 + dtype_idx * 977 + iteration) - local_kv = storage[:num_tokens] - local_kv.copy_( - torch.randn( - num_tokens, - token_dim, - device=device, - dtype=torch.float32, - generator=generator, - ).to(dtype) - ) - gathered = storage[128 : 128 + num_tokens * world_size] - active_ubatch[0] = iteration % 2 - workspace.gather(gathered, local_kv) + if dtype == torch.bfloat16: + active_ubatch[0] = 0 + + # A gather on the other slot is the all-rank rendezvous that + # makes the first slot reusable. Rank 0 deliberately consumes + # slot 0 late; no explicit barrier protects the critical + # slot-0 -> slot-1 -> slot-0 sequence. + first, expected = publish(workspace, dtype, 30, 0, full_layout) + torch.accelerator.synchronize() + if rank == 0: + time.sleep(0.25) + assert_planes(first, expected) + publish(workspace, dtype, 31, 1, alternate_layout) torch.accelerator.synchronize() + reused, expected = publish(workspace, dtype, 32, 0, small_layout) + torch.accelerator.synchronize() + assert_planes(reused, expected) + dist.barrier() # Isolate the next ownership scenario. + + # Model execution may reset to slot 0 at a layer boundary. + # Simulate attention consumption followed by its downstream + # TP collective, then reuse the same slot immediately. + first, expected = publish(workspace, dtype, 33, 0, full_layout) + consumed = sum(plane.float().sum() for plane in first) + expected_consumed = sum(plane.float().sum() for plane in expected) + dist.all_reduce(consumed) + reused, expected = publish(workspace, dtype, 34, 0, small_layout) + torch.accelerator.synchronize() + torch.testing.assert_close(consumed, expected_consumed * world_size) + assert_planes(reused, expected) + dist.barrier() - expected = torch.empty_like(gathered) - dist.all_gather_into_tensor(expected, local_kv.contiguous()) - assert torch.equal( - gathered.view(torch.uint8), expected.view(torch.uint8) + if env.get("TEST_CUDA_GRAPH") != "1" or dtype != torch.bfloat16: + continue + + graph_workspace = dcp_utils.DirectDCPKVGatherWorkspace( + dist.group.WORLD, + device, + max_gathered_tokens, + token_dim, + plane_split_dim, + dtype, + num_ubatches=2, + ) + capture_maps, capture_output_tokens = layout_maps(full_layout) + captured_input = make_local(rank, 20, len(capture_maps[rank]), dtype) + captured_dst_rows = torch.tensor( + capture_maps[rank], dtype=torch.int32, device=device + ) + active_ubatch[0] = 1 + torch.accelerator.synchronize() + dist.barrier() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_kv_c, captured_k_pe = graph_workspace.gather( + captured_input, + captured_dst_rows, + capture_output_tokens, + buffer_slot=1, ) + dist.barrier() + graph.replay() + torch.accelerator.synchronize() + dist.barrier() + expected_kv_c, expected_k_pe = expected_planes(full_layout, 20, dtype) + assert torch.equal( + captured_kv_c.view(torch.uint8), expected_kv_c.view(torch.uint8) + ) + assert torch.equal( + captured_k_pe.view(torch.uint8), expected_k_pe.view(torch.uint8) + ) + dist.barrier() + + # DBO owns a separate pair of buffers per ubatch. An eager call in + # ubatch 0 must not disturb the captured ubatch-1 slot. + check_case( + graph_workspace, + dtype, + iteration=21, + ubatch=0, + buffer_slot=0, + layout=small_layout, + ) + maps, output_tokens = layout_maps(alternate_layout) + assert output_tokens == capture_output_tokens + captured_input.copy_(make_local(rank, 22, len(maps[rank]), dtype)) + captured_dst_rows.copy_( + torch.tensor(maps[rank], dtype=torch.int32, device=device) + ) + epochs_before = graph_workspace.epoch.clone() + torch.accelerator.synchronize() + active_ubatch[0] = 0 # Replay retains the ubatch captured above. + graph.replay() + torch.accelerator.synchronize() + expected_kv_c, expected_k_pe = expected_planes(alternate_layout, 22, dtype) + assert torch.equal( + captured_kv_c.view(torch.uint8), expected_kv_c.view(torch.uint8) + ) + assert torch.equal( + captured_k_pe.view(torch.uint8), expected_k_pe.view(torch.uint8) + ) + assert int(graph_workspace.epoch[1].item()) == ( + int(epochs_before[1].item()) + 1 + ) + assert int(graph_workspace.epoch[0].item()) == int(epochs_before[0].item()) + assert not torch.count_nonzero(graph_workspace.completion) + dist.barrier() finally: dist.destroy_process_group() @@ -833,7 +1151,7 @@ def test_distributed_direct_kv_gather_matches_reference(world_size: int): _distributed_run( _distributed_direct_kv_gather_worker, world_size=world_size, - extra_env={}, + extra_env={"TEST_CUDA_GRAPH": "1"}, ) diff --git a/tests/v1/attention/test_mla_context_chunks.py b/tests/v1/attention/test_mla_context_chunks.py index 006cf301e6c9..048276c79d67 100644 --- a/tests/v1/attention/test_mla_context_chunks.py +++ b/tests/v1/attention/test_mla_context_chunks.py @@ -8,11 +8,15 @@ longer needs an empty-span masking pass). """ +from types import SimpleNamespace + import pytest import torch +import vllm.model_executor.layers.attention.mla_attention as mla_attention from vllm.model_executor.layers.attention.mla_attention import ( _gather_dcp_context_kv, + build_dcp_kv_final_layout_dst_rows, build_mla_chunked_context_metadata, init_mla_context_partial, reorg_kvcache, @@ -29,10 +33,16 @@ def build_chunked_context( dcp_world_size: int = 1, dcp_local_block_size: int = 1, direct_dcp_kv_gather: bool = False, + dcp_manager=None, ): query_start_loc = torch.zeros(len(query_lens) + 1, dtype=torch.int32) query_start_loc[1:] = torch.tensor(query_lens, dtype=torch.int32).cumsum(0) - workspace_rows = workspace_size + workspace_size // dcp_world_size + if dcp_world_size == 1: + workspace_rows = workspace_size + elif dcp_manager is not None and dcp_manager.use_direct_kv_gather: + workspace_rows = workspace_size // dcp_world_size + else: + workspace_rows = workspace_size + workspace_size // dcp_world_size return build_mla_chunked_context_metadata( context_lens_cpu=torch.tensor(context_lens, dtype=torch.int32), prefill_query_start_loc_cpu=query_start_loc, @@ -45,6 +55,7 @@ def build_chunked_context( dcp_local_block_size=dcp_local_block_size, dcp_virtual_block_size=dcp_local_block_size * dcp_world_size, direct_dcp_kv_gather=direct_dcp_kv_gather, + dcp_manager=dcp_manager, ) @@ -307,7 +318,7 @@ def test_dcp_reorg_uses_each_chunks_local_starts(): toks = chunk.num_local_context_tokens rank_buffers = [torch.full((toks, 1, 1), -1) for _ in range(dcp_world_size)] - expected = [] + expected: list[torch.Tensor] = [] src_token_idx = 0 for request, (padded_len, local_lens, local_start) in enumerate( zip( @@ -340,3 +351,141 @@ def test_dcp_reorg_uses_each_chunks_local_starts(): ) torch.testing.assert_close(reorganized, torch.cat(expected)) + + +@pytest.mark.skip_global_cleanup +def test_dcp_final_layout_dst_rows_maps_padding_and_continuations(): + """The publish map is compact, disjoint, and skips padded local rows.""" + padded_local_seq_lens = [4, 3] + local_context_lens_allranks = [[4, 3], [7, 5]] + local_starts = [0, 4] + + rank_0 = build_dcp_kv_final_layout_dst_rows( + padded_local_seq_lens, + local_context_lens_allranks, + local_starts, + [7, 4], + dcp_rank=0, + ) + rank_1 = build_dcp_kv_final_layout_dst_rows( + padded_local_seq_lens, + local_context_lens_allranks, + local_starts, + [7, 4], + dcp_rank=1, + ) + + assert rank_0.tolist() == [0, 1, 2, 3, 7, 8, 9] + assert rank_1.tolist() == [4, 5, 6, -1, 10, -1, -1] + valid_rows = sorted(row for row in [*rank_0.tolist(), *rank_1.tolist()] if row >= 0) + assert valid_rows == list(range(11)) + + zero_valid_maps = [ + build_dcp_kv_final_layout_dst_rows( + padded_local_seq_lens=[2], + local_context_lens_allranks=[[4, 4, 3, 2]], + local_starts=[2], + output_seq_lens=[5], + dcp_rank=rank, + ).tolist() + for rank in range(4) + ] + assert zero_valid_maps == [[0, 1], [2, 3], [4, -1], [-1, -1]] + + +@pytest.mark.skip_global_cleanup +def test_dcp_final_layout_validates_each_request_exact_coverage(): + """Equal batch totals cannot hide a gap in one request and overlap in another.""" + with pytest.raises(ValueError, match="exactly cover request 0"): + build_dcp_kv_final_layout_dst_rows( + padded_local_seq_lens=[2, 2], + local_context_lens_allranks=[[1, 1], [2, 2]], + local_starts=[0, 0], + output_seq_lens=[3, 3], + dcp_rank=0, + ) + + mismatched_manager = SimpleNamespace( + use_direct_kv_gather=True, + group=SimpleNamespace(rank_in_group=0, world_size=4), + ) + with pytest.raises(ValueError, match="world sizes differ"): + build_chunked_context( + [128], + [4], + 1024, + dcp_world_size=2, + dcp_local_block_size=64, + dcp_manager=mismatched_manager, + ) + + +@pytest.mark.skip_global_cleanup +def test_dcp_final_layout_publish_matches_reorg(monkeypatch): + """All source-rank maps jointly produce the existing compact layout.""" + dcp_world_size, interleave, workspace_size = 2, 64, 1024 + monkeypatch.setattr(mla_attention, "np_to_pinned_tensor", torch.from_numpy) + metadata_by_rank = [] + for rank in range(dcp_world_size): + dcp_manager = SimpleNamespace( + use_direct_kv_gather=True, + group=SimpleNamespace(rank_in_group=rank, world_size=dcp_world_size), + ) + metadata = build_chunked_context( + [3000, 200, 200], + [4, 4, 4], + workspace_size, + block_size=128, + dcp_world_size=dcp_world_size, + dcp_local_block_size=interleave, + dcp_manager=dcp_manager, + ) + assert metadata is not None + assert metadata.workspace.shape[0] == workspace_size // dcp_world_size + assert all( + chunk.num_local_context_tokens <= workspace_size // dcp_world_size + for chunk in metadata.chunks + ) + metadata_by_rank.append(metadata) + + for chunk_index, chunk in enumerate(metadata_by_rank[0].chunks): + assert chunk.padded_local_seq_lens is not None + assert chunk.local_context_lens_allranks is not None + assert chunk.local_starts is not None + + toks = chunk.num_local_context_tokens + compact = torch.full((chunk.num_context_tokens,), -1, dtype=torch.int64) + expected: list[list[torch.Tensor]] = [] + for rank in range(dcp_world_size): + local = torch.full((toks,), -2, dtype=torch.int64) + src_token_idx = 0 + for request, (padded_len, local_lens, local_start) in enumerate( + zip( + chunk.padded_local_seq_lens, + chunk.local_context_lens_allranks, + chunk.local_starts, + ) + ): + actual_len = min(max(0, local_lens[rank] - local_start), padded_len) + values = ( + request * 100_000 + + rank * 10_000 + + torch.arange(local_start, local_start + actual_len) + ) + local[src_token_idx : src_token_idx + actual_len] = values + if rank == 0: + expected.append([]) + expected[request].append(values) + src_token_idx += padded_len + + dst_rows = metadata_by_rank[rank].chunks[chunk_index].final_layout_dst_rows + assert dst_rows is not None + valid = dst_rows >= 0 + assert torch.all(compact[dst_rows[valid]] == -1) + compact[dst_rows[valid]] = local[valid] + + assert not torch.any(compact == -1) + torch.testing.assert_close( + compact, + torch.cat([torch.cat(request) for request in expected]), + ) diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 83ea816e59b3..ffa92d4d3bd6 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -2512,6 +2512,7 @@ class ContextChunk: padded_local_token_to_seq: torch.Tensor | None = None num_local_context_tokens: int = 0 local_starts: list[int] | None = None + final_layout_dst_rows: torch.Tensor | None = None @property def num_requests(self) -> int: @@ -2783,6 +2784,69 @@ def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor: return np_to_pinned_tensor(np.asarray(values, dtype=np.int32)) +def build_dcp_kv_final_layout_dst_rows( + padded_local_seq_lens: list[int], + local_context_lens_allranks: list[list[int]], + local_starts: list[int], + output_seq_lens: list[int], + dcp_rank: int, +) -> np.ndarray: + """Map padded local DCP rows to compact request-major output rows. + + Local cache gathering packs every request into an equally padded segment on + every DCP rank. The MLA consumer instead needs requests outermost and only + real rows from rank 0..N within each request. Negative entries identify + local padding that the multicast publisher must skip. + """ + if not ( + len(padded_local_seq_lens) + == len(local_context_lens_allranks) + == len(local_starts) + == len(output_seq_lens) + ): + raise ValueError("DCP final-layout metadata must have equal request counts") + if not local_context_lens_allranks: + return np.empty(0, dtype=np.int32) + world_size = len(local_context_lens_allranks[0]) + if world_size <= 1: + raise ValueError( + f"DCP final-layout metadata requires at least two ranks: {world_size}" + ) + if not 0 <= dcp_rank < world_size: + raise ValueError(f"invalid DCP rank {dcp_rank} for world size {world_size}") + + dst_rows: list[int] = [] + request_output_start = 0 + for request, (padded_len, context_lens, local_start, output_len) in enumerate( + zip( + padded_local_seq_lens, + local_context_lens_allranks, + local_starts, + output_seq_lens, + strict=True, + ) + ): + if len(context_lens) != world_size: + raise ValueError("DCP final-layout metadata has inconsistent world sizes") + valid_lens = [ + min(max(0, context_len - local_start), padded_len) + for context_len in context_lens + ] + covered = sum(valid_lens) + if covered != output_len: + raise ValueError( + "DCP final-layout rows do not exactly cover request " + f"{request}: rank lengths {valid_lens} cover {covered}, " + f"expected {output_len}" + ) + local_valid = valid_lens[dcp_rank] + rank_output_start = request_output_start + sum(valid_lens[:dcp_rank]) + dst_rows.extend(range(rank_output_start, rank_output_start + local_valid)) + dst_rows.extend([-1] * (padded_len - local_valid)) + request_output_start += output_len + return np.asarray(dst_rows, dtype=np.int32) + + def align_mla_chunked_context_workspace_size( vllm_config: VllmConfig, workspace_size: int, @@ -2906,6 +2970,7 @@ def padded_rows(rows: int) -> int: token_to_seq_parts: list[np.ndarray] = [] padded_local_cu_seq_lens_flat: list[int] = [] padded_local_token_to_seq_parts: list[np.ndarray] = [] + final_layout_dst_rows_parts: list[np.ndarray] = [] local_starts_per_chunk: list[list[int]] = [] local_seq_lens_per_chunk: list[list[int]] = [] layouts: list[tuple[slice, slice, slice, slice]] = [] @@ -2947,6 +3012,24 @@ def padded_rows(rows: int) -> int: np.repeat(np.arange(num_requests, dtype=np.int32), local_seq_lens) ) num_local_tokens = sum(local_seq_lens) + if dcp_manager is not None and dcp_manager.use_direct_kv_gather: + assert local_context_lens_allranks is not None + if dcp_manager.group.world_size != dcp_world_size: + raise ValueError( + "DCP group and chunk metadata world sizes differ: " + f"{dcp_manager.group.world_size} != {dcp_world_size}" + ) + request_context_lens = local_context_lens_allranks[ + plan.request_start : plan.request_end + ] + final_layout_dst_rows = build_dcp_kv_final_layout_dst_rows( + local_seq_lens, + request_context_lens, + local_starts, + plan.seq_lens, + dcp_manager.group.rank_in_group, + ) + final_layout_dst_rows_parts.append(final_layout_dst_rows) # The gather takes per-rank local offsets under DCP. starts_flat.extend(local_starts) else: @@ -2980,6 +3063,13 @@ def padded_rows(rows: int) -> int: padded_local_token_to_seq = _flat_int32( np.concatenate(padded_local_token_to_seq_parts) ).to(device, non_blocking=True) + final_layout_dst_rows = ( + _flat_int32(np.concatenate(final_layout_dst_rows_parts)).to( + device, non_blocking=True + ) + if final_layout_dst_rows_parts + else None + ) chunks: list[MLACommonPrefillMetadata.ContextChunk] = [] for index, (plan, layout) in enumerate(zip(plans, layouts)): @@ -3017,6 +3107,8 @@ def padded_rows(rows: int) -> int: local_token_slice ] chunk.local_starts = local_starts_per_chunk[index] + if final_layout_dst_rows is not None: + chunk.final_layout_dst_rows = final_layout_dst_rows[local_token_slice] chunks.append(chunk) return MLACommonPrefillMetadata.ChunkedContextMetadata( @@ -3200,37 +3292,45 @@ def __init__( use_packed_fp8_cache = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" self.dcp_manager: MLADCPManager | None = None if self.dcp_world_size > 1: - # Note(hc): The local kvcache is incomplete when DCP is triggered, - # an additional kvcache allgather across the DCP group is therefore - # required, so the workspace has to be enlarged by 1/DCP relative - # to the original TP allocation. assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0 - self.chunked_prefill_workspace = torch.empty( - ( - self.chunked_prefill_workspace_size - + self.chunked_prefill_workspace_size // self.dcp_world_size, - self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim, - ), - dtype=torch.bfloat16 - if use_packed_fp8_cache - else self.model_config.dtype, - device=device, + workspace_dtype = ( + torch.bfloat16 if use_packed_fp8_cache else self.model_config.dtype + ) + workspace_head_size = ( + self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim ) self.dcp_manager = getattr(attention_layer, "dcp_manager", None) + use_direct_kv_gather = False if self.dcp_manager is not None: if not isinstance(self.dcp_manager, MLADCPManager): raise TypeError( "MLA attention layer dcp_manager must be an " f"MLADCPManager, got {type(self.dcp_manager).__name__}." ) - self.dcp_manager.init_kv_gather( - self.chunked_prefill_workspace, + use_direct_kv_gather = self.dcp_manager.init_kv_gather( self.chunked_prefill_workspace_size, + workspace_head_size, + self.mla_dims.kv_lora_rank, + workspace_dtype, ) elif not self.supports_direct_dcp_kv_gather: raise RuntimeError( f"{type(self).__name__} requires MLADCPManager when DCP is enabled." ) + # The direct path only materializes this rank's padded rows. NCCL + # and the backend-owned fallback additionally need a rank-major + # destination for every rank. + local_rows = self.chunked_prefill_workspace_size // self.dcp_world_size + workspace_rows = ( + local_rows + if use_direct_kv_gather + else self.chunked_prefill_workspace_size + local_rows + ) + self.chunked_prefill_workspace = torch.empty( + (workspace_rows, workspace_head_size), + dtype=workspace_dtype, + device=device, + ) else: self.chunked_prefill_workspace = torch.empty( ( @@ -3861,42 +3961,62 @@ def _context_parallel_compute_prefill_context( batch_size=chunk.num_requests, seq_starts=chunk.starts, ) - # workspace - # |------- N tokens --------|--------- N*dcp_size tokens ----------| - # |<- use for local_gather ->|<--------- use for allgather -------->| - allgather_offset = workspace.shape[0] // (dcp_world_size + 1) - assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0] - assert toks <= allgather_offset - local_gathered_kvcache = workspace[:toks] - cur_allgather_workspace = workspace[ - allgather_offset : allgather_offset * (1 + dcp_world_size) - ] - assert toks * dcp_world_size <= cur_allgather_workspace.shape[0] - cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size] - _gather_dcp_context_kv( - cur_allgather_kvcache, - local_gathered_kvcache, - dcp_manager=chunked_context.dcp_manager, - direct_dcp_kv_gather=chunked_context.direct_dcp_kv_gather, - ) - assert ( - cur_allgather_kvcache.shape[-1] - == self.kv_lora_rank + self.qk_rope_head_dim - ) - allgatered_kv_c_normed, allgatered_k_pe = cur_allgather_kvcache.unsqueeze( - 1 - ).split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) - - kv_c_normed, k_pe = reorg_kvcache( - allgatered_kv_c_normed, - allgatered_k_pe, - padded_local_chunk_seq_lens_lst=chunk.padded_local_seq_lens, - local_context_lens_allranks=chunk.local_context_lens_allranks, - local_starts=chunk.local_starts, - sum_seq_len=chunk.num_context_tokens, - max_seq_len=chunk.max_seq_len, - toks=toks, - ) + dcp_manager = chunked_context.dcp_manager + if dcp_manager is not None and dcp_manager.use_direct_kv_gather: + assert toks <= workspace.shape[0] + local_gathered_kvcache = workspace[:toks] + assert chunk.final_layout_dst_rows is not None + kv_c_normed, k_pe = dcp_manager.direct_kv_gather( + local_gathered_kvcache, + chunk.final_layout_dst_rows, + chunk.num_context_tokens, + # The next gather synchronizes alternating chunks. The TP + # output collective synchronizes the final chunk before + # the next layer/forward resets the slot to zero. + chunk.index & 1, + ) + assert kv_c_normed.is_contiguous() + assert k_pe.is_contiguous() + else: + # workspace + # |------- N tokens --------|------ N*dcp_size tokens -------| + # |<- use for local gather ->|<------ use for allgather ----->| + allgather_offset = workspace.shape[0] // (dcp_world_size + 1) + assert allgather_offset * (dcp_world_size + 1) == workspace.shape[0] + assert toks <= allgather_offset + local_gathered_kvcache = workspace[:toks] + cur_allgather_workspace = workspace[ + allgather_offset : allgather_offset * (1 + dcp_world_size) + ] + assert toks * dcp_world_size <= cur_allgather_workspace.shape[0] + cur_allgather_kvcache = cur_allgather_workspace[ + : toks * dcp_world_size + ] + _gather_dcp_context_kv( + cur_allgather_kvcache, + local_gathered_kvcache, + dcp_manager=dcp_manager, + direct_dcp_kv_gather=chunked_context.direct_dcp_kv_gather, + ) + assert ( + cur_allgather_kvcache.shape[-1] + == self.kv_lora_rank + self.qk_rope_head_dim + ) + allgatered_kv_c_normed, allgatered_k_pe = ( + cur_allgather_kvcache.unsqueeze(1).split( + [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + ) + kv_c_normed, k_pe = reorg_kvcache( + allgatered_kv_c_normed, + allgatered_k_pe, + padded_local_chunk_seq_lens_lst=chunk.padded_local_seq_lens, + local_context_lens_allranks=chunk.local_context_lens_allranks, + local_starts=chunk.local_starts, + sum_seq_len=chunk.num_context_tokens, + max_seq_len=chunk.max_seq_len, + toks=toks, + ) if kv_b_proj_input_dtype is not None: kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype) diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index 132bb5f9203b..cffcca3b0ea5 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -140,16 +140,6 @@ def __init__( self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim ) workspace_rows = self.chunked_prefill_workspace_size - if self.dcp_world_size > 1: - # DCP gathers each rank's local KV shard into the workspace, so it - # needs an extra 1/DCP rows beyond the TP allocation. - assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0 - workspace_rows += self.chunked_prefill_workspace_size // self.dcp_world_size - self.chunked_prefill_workspace = torch.empty( - (workspace_rows, workspace_head_size), - dtype=self.model_config.dtype, - device=device, - ) self.topk_mask_workspace: torch.Tensor | None = None if _is_masked_mha_available( self.model_config.model_arch_config.total_num_attention_heads, @@ -170,13 +160,31 @@ def __init__( layer_prefill_backend = attention_layer.prefill_backend self.dcp_manager: MLADCPManager | None = None if self.dcp_world_size > 1: + assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0 self.dcp_manager = getattr(attention_layer, "dcp_manager", None) assert isinstance(self.dcp_manager, MLADCPManager) if layer_prefill_backend is not None: - self.dcp_manager.init_kv_gather( - self.chunked_prefill_workspace, + use_direct_kv_gather = self.dcp_manager.init_kv_gather( self.chunked_prefill_workspace_size, + workspace_head_size, + self.mla_dims.kv_lora_rank, + self.model_config.dtype, ) + local_rows = self.chunked_prefill_workspace_size // self.dcp_world_size + workspace_rows = ( + local_rows + if use_direct_kv_gather + else self.chunked_prefill_workspace_size + local_rows + ) + else: + workspace_rows += ( + self.chunked_prefill_workspace_size // self.dcp_world_size + ) + self.chunked_prefill_workspace = torch.empty( + (workspace_rows, workspace_head_size), + dtype=self.model_config.dtype, + device=device, + ) self._prefill_backend = ( layer_prefill_backend.clone() if layer_prefill_backend is not None else None ) diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index 0d8243950f61..a49a35cef66b 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -1095,8 +1095,9 @@ def _compute_prefill_context( copied per chunk. Decode context parallelism keeps using - ``impl._context_parallel_compute_prefill_context``; its extra allgather - and reorg are not fused here. + ``impl._context_parallel_compute_prefill_context``. Its NCCL fallback + retains the extra all-gather and reorganization; the direct symmetric + path publishes into the compact MLA layout instead. """ prefill = attn_metadata.prefill assert prefill is not None diff --git a/vllm/v1/attention/ops/dcp_utils.py b/vllm/v1/attention/ops/dcp_utils.py index b28d337a1046..de35a9e0ff65 100644 --- a/vllm/v1/attention/ops/dcp_utils.py +++ b/vllm/v1/attention/ops/dcp_utils.py @@ -472,12 +472,31 @@ def get_direct_dcp_q_gather_workspace( ) -def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool: - return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0 +def _kv_gather_layout_supported( + token_dim: int, + plane_split_dim: int, + dtype: torch.dtype, +) -> bool: + """Whether both packed output planes support 16-byte multicast stores.""" + if not 0 < plane_split_dim < token_dim: + return False + element_size = torch.empty((), dtype=dtype).element_size() + return ( + plane_split_dim * element_size % 16 == 0 + and (token_dim - plane_split_dim) * element_size % 16 == 0 + ) class DirectDCPKVGatherWorkspace(_DirectDCPWorkspace): - """Persistent symmetric buffers for direct DCP KV gather.""" + """Persistent symmetric buffers for direct DCP KV gather. + + Storage is owned by ``(DBO ubatch, buffer slot)``. Different ubatches have + disjoint buffers and may run independently. Within one ubatch, publishing + and consumption must remain stream ordered. Before reusing the same slot, + every rank must have consumed it and reached either a gather on the other + slot or another all-rank rendezvous. Concurrent same-ubatch use from + multiple streams is not supported. + """ def __init__( self, @@ -485,6 +504,7 @@ def __init__( device: torch.device, max_gathered_tokens: int, token_dim: int, + plane_split_dim: int, dtype: torch.dtype = torch.bfloat16, num_ubatches: int = 1, ) -> None: @@ -500,8 +520,10 @@ def __init__( "Direct DCP kv-gather dimensions must be positive, got " f"T={max_gathered_tokens}, D={token_dim}" ) - if not _kv_gather_layout_supported(token_dim, dtype): - raise ValueError("Direct DCP kv-gather requires 16-byte-aligned KV rows.") + if not _kv_gather_layout_supported(token_dim, plane_split_dim, dtype): + raise ValueError( + "Direct DCP kv-gather requires two nonempty 16-byte-aligned KV planes." + ) super().__init__(group, device, num_ubatches) if self.world_size <= 1: raise ValueError("Direct DCP kv-gather requires at least two ranks") @@ -511,6 +533,8 @@ def __init__( f"ranks: {max_gathered_tokens} % {self.world_size} != 0" ) self.max_gathered_tokens = max_gathered_tokens + self.token_dim = token_dim + self.plane_split_dim = plane_split_dim kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim) signal_shape = (num_ubatches, 2, self.world_size) @@ -528,26 +552,48 @@ def __init__( self.completion = self.received_signal.new_zeros((num_ubatches, 2)) torch.accelerator.synchronize() - def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None: + def gather( + self, + local_kv: torch.Tensor, + dst_rows: torch.Tensor, + output_tokens: int, + buffer_slot: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Publish valid rows into compact request-major KV planes.""" ubatch = dbo_current_ubatch_id() if not 0 <= ubatch < self.num_ubatches: raise ValueError( f"DCP kv-gather ubatch {ubatch} exceeds {self.num_ubatches} slots" ) + # The custom op validates the dynamic tensor geometry, dtype, device, + # capacity, and slot before launching. Avoid duplicating those checks on + # this latency-sensitive host path. + token_dim = self.token_dim + plane_split_dim = self.plane_split_dim kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] torch.ops._C.direct_dcp_kv_gather( local_kv, + dst_rows, self.received_kv[ubatch], self.received_signal[ubatch], self.completion[ubatch], self.epoch[ubatch : ubatch + 1], - gathered_kv, + output_tokens, + plane_split_dim, + buffer_slot, self.world_size, self.rank, self.max_gathered_tokens, kv_multicast_ptr, signal_multicast_ptr, ) + slot = self.received_kv[ubatch, buffer_slot].view(-1) + kv_c_capacity = self.max_gathered_tokens * plane_split_dim + kv_c = slot[:kv_c_capacity].view(self.max_gathered_tokens, 1, plane_split_dim) + k_pe = slot[kv_c_capacity:].view( + self.max_gathered_tokens, 1, token_dim - plane_split_dim + ) + return kv_c[:output_tokens], k_pe[:output_tokens] @functools.cache @@ -556,6 +602,7 @@ def get_direct_dcp_kv_gather_workspace( device: torch.device, max_gathered_tokens: int, token_dim: int, + plane_split_dim: int, dtype: torch.dtype, num_ubatches: int, ) -> DirectDCPKVGatherWorkspace | None: @@ -566,13 +613,14 @@ def get_direct_dcp_kv_gather_workspace( _KV_GATHER_SUPPORTED_DTYPES, ): return None - if not _kv_gather_layout_supported(token_dim, dtype): + if not _kv_gather_layout_supported(token_dim, plane_split_dim, dtype): return None return DirectDCPKVGatherWorkspace( group.device_group, device, max_gathered_tokens, token_dim, + plane_split_dim, dtype, num_ubatches, ) @@ -592,7 +640,7 @@ def __call__( class MLADCPManager: """Select and own layer-level collective implementations for MLA DCP.""" - _kv_gather: Callable[[torch.Tensor, torch.Tensor], object] + _kv_gather: Callable[[torch.Tensor, torch.Tensor], object] | None def __init__( self, @@ -614,6 +662,8 @@ def __init__( self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config) self.use_a2a = parallel_config.dcp_comm_backend == "a2a" self.padded_num_heads = padded_num_heads + self._direct_kv_gather_workspace: DirectDCPKVGatherWorkspace | None = None + self._kv_gather = None self.combine = self._init_combine( num_heads, @@ -700,41 +750,77 @@ def _gather_query(self, query: torch.Tensor) -> torch.Tensor: def init_kv_gather( self, - workspace: torch.Tensor, max_gathered_tokens: int, - ) -> None: + token_dim: int, + plane_split_dim: int, + dtype: torch.dtype, + ) -> bool: + """Select the KV collective before allocating its local scratch. + + Returns whether the direct final-layout publisher was selected. That + path only needs one rank's local rows; the fallback additionally needs + a rank-major all-gather destination. + """ world_size = self.group.world_size - assert max_gathered_tokens > 0 - assert max_gathered_tokens % world_size == 0 - assert workspace.ndim == 2 - assert workspace.is_contiguous() - assert workspace.shape[0] == ( - max_gathered_tokens + max_gathered_tokens // world_size - ) - assert workspace.shape[1] > 0 + if max_gathered_tokens <= 0 or max_gathered_tokens % world_size != 0: + raise ValueError( + "DCP KV gather capacity must be positive and divide evenly " + f"across {world_size} ranks, got {max_gathered_tokens}" + ) + if token_dim <= 0: + raise ValueError( + f"DCP KV gather token dimension must be positive: {token_dim}" + ) direct_workspace = get_direct_dcp_kv_gather_workspace( self.group, - workspace.device, + self.device, max_gathered_tokens, - workspace.shape[1], - workspace.dtype, + token_dim, + plane_split_dim, + dtype, self.num_ubatches, ) + self._direct_kv_gather_workspace = direct_workspace if direct_workspace is not None: logger.info_once( - "Using direct symmetric-memory DCP chunked-context KV gather for MLA." + "Using direct symmetric-memory DCP final-layout KV multicast for MLA." ) - self._kv_gather = direct_workspace.gather + self._kv_gather = None else: self._kv_gather = functools.partial( torch.distributed.all_gather_into_tensor, group=self.group.device_group, ) + return direct_workspace is not None + + @property + def use_direct_kv_gather(self) -> bool: + return self._direct_kv_gather_workspace is not None def kv_gather( self, gathered_kv: torch.Tensor, local_kv: torch.Tensor, ) -> object: - return self._kv_gather(gathered_kv, local_kv) + kv_gather = self._kv_gather + if kv_gather is None: + raise RuntimeError("NCCL DCP KV gather is not selected") + return kv_gather(gathered_kv, local_kv) + + def direct_kv_gather( + self, + local_kv: torch.Tensor, + dst_rows: torch.Tensor, + output_tokens: int, + buffer_slot: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + workspace = self._direct_kv_gather_workspace + if not self.use_direct_kv_gather or workspace is None: + raise RuntimeError("direct DCP KV gather is not enabled") + return workspace.gather( + local_kv, + dst_rows, + output_tokens, + buffer_slot, + ) From 96789b501deaaa0c9d8ea45c0de49eb7ea9b8f7e Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Tue, 1 Sep 2026 18:46:24 +0900 Subject: [PATCH 46/52] perf(mla): publish K3 prefill KV over PCIe Extend the exact final-layout DCP publisher to same-node PCIe peer memory, and give backend-owned decode DCP a standalone context-KV gather plan. Kimi consumes the compact planes through the fused context K/V packing loop, avoiding rank-major materialization, Python reorganization, and separate K/V pack launches without changing chunk geometry or attention semantics. Validated with 46 metadata/collective unit tests, 41 Kimi fused-context and CUDA oracle tests, and an SM120 stable-extension build. The multi-rank peer oracle is run during deployment qualification after the target GPUs are drained. AI assistance was used. --- .../dcp_utils/dcp_direct_kv_gather.cu | 151 +++++++++++-- .../test_dcp_direct_a2a_lse_reduce.py | 75 +++--- .../kimi_k3/test_mla_prefill_context.py | 61 +++++ .../layers/attention/mla_attention.py | 33 +-- vllm/models/kimi_k3/nvidia/mla.py | 82 +++++-- vllm/v1/attention/ops/dcp_utils.py | 213 ++++++++++-------- 6 files changed, 447 insertions(+), 168 deletions(-) diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu index 801803c09c9e..240a0e16e524 100644 --- a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu @@ -12,14 +12,17 @@ namespace { using vllm::direct_dcp::check_cuda_launch; +using vllm::direct_dcp::get_peer_ptr; using vllm::direct_dcp::increment_epoch_kernel; using vllm::direct_dcp::multimem_store_16; using vllm::direct_dcp::multimem_store_release_system; +using vllm::direct_dcp::store_release_system; using vllm::direct_dcp::wait_for_epoch; constexpr int kThreads = 256; // KV chunks need many blocks in flight to saturate the fabric. constexpr int64_t kMaxMulticastBlocks = 128; +constexpr int64_t kMaxPeerBlocks = 1024; // Multicast each rank's valid local rows directly into compact, request-major // kv_c and k_pe planes. dst_rows maps every padded local input row to its final @@ -91,8 +94,89 @@ __global__ void direct_dcp_kv_gather_multimem_kernel( } } +// PCIe/UVA fallback for systems without NVLS multicast. Each source rank +// publishes its valid rows directly into every peer's compact, request-major +// planes. This moves the same payload volume as an all-gather while avoiding +// the rank-major materialization and the subsequent reorganization pass. +__global__ void direct_dcp_kv_gather_peer_kernel( + const uint4* local_kv, const int32_t* dst_rows, + const int64_t* peer_kv_ptrs, const int64_t* peer_signal_ptrs, + const uint32_t* received_signal, const int64_t* epoch_ptr, + uint32_t* completion, int64_t world_size, int64_t rank, + int64_t num_tokens, int64_t items_per_row, int64_t kv_c_items_per_row, + int64_t output_tokens, int64_t max_gathered_tokens, + int64_t buffer_slot, int64_t slot_stride_items) { + uint32_t epoch = static_cast(epoch_ptr[0]); + int64_t source_items = num_tokens * items_per_row; + int64_t total_items = world_size * source_items; + int64_t item_stride = static_cast(gridDim.x) * blockDim.x; + for (int64_t linear = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + linear < total_items; linear += item_stride) { + int64_t destination_rank = linear / source_items; + int64_t item = linear - destination_rank * source_items; + int64_t src_row = item / items_per_row; + int32_t dst_row = dst_rows[src_row]; + if (dst_row < 0) { + continue; + } + if (dst_row >= output_tokens) { + printf( + "direct DCP peer final-layout destination out of bounds " + "source=%lld dst=%d output_tokens=%lld\n", + static_cast(rank), dst_row, + static_cast(output_tokens)); + asm volatile("trap;"); + } + int64_t row_item = item - src_row * items_per_row; + int64_t dst_item; + if (row_item < kv_c_items_per_row) { + dst_item = static_cast(dst_row) * kv_c_items_per_row + row_item; + } else { + int64_t k_pe_items_per_row = items_per_row - kv_c_items_per_row; + dst_item = max_gathered_tokens * kv_c_items_per_row + + static_cast(dst_row) * k_pe_items_per_row + + row_item - kv_c_items_per_row; + } + uint4* peer_kv = get_peer_ptr(peer_kv_ptrs, destination_rank); + peer_kv[buffer_slot * slot_stride_items + dst_item] = local_kv[item]; + } + + // Every block publishes its system-scope payload before contributing to the + // completion count. The final block releases one signal to each peer, then + // waits until every source has published into this rank's local replica. + __threadfence_system(); + __syncthreads(); + if (threadIdx.x != 0) { + return; + } + uint32_t completed = atomicAdd(completion + buffer_slot, 1u); + if (completed + 1u != gridDim.x) { + return; + } + atomicExch(completion + buffer_slot, 0u); + + int64_t signal_item = buffer_slot * world_size + rank; + for (int64_t destination_rank = 0; destination_rank < world_size; + ++destination_rank) { + uint32_t* peer_signal = + get_peer_ptr(peer_signal_ptrs, destination_rank); + store_release_system(peer_signal + signal_item, epoch); + } + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { + int64_t source_signal_item = buffer_slot * world_size + source_rank; + if (!wait_for_epoch(received_signal + source_signal_item, epoch)) { + printf("direct DCP peer final-layout timeout source=%lld epoch=%u\n", + static_cast(source_rank), epoch); + asm volatile("trap;"); + } + } +} + void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, const torch::stable::Tensor& dst_rows, + const torch::stable::Tensor& peer_kv_ptrs, + const torch::stable::Tensor& peer_signal_ptrs, torch::stable::Tensor& received_kv, torch::stable::Tensor& received_signal, torch::stable::Tensor& completion, @@ -116,6 +200,17 @@ void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, dst_rows.dim() == 1 && dst_rows.numel() == local_kv.size(0), "final-layout destination rows must be CUDA int32 [T]"); + STD_TORCH_CHECK( + peer_kv_ptrs.is_cuda() && peer_kv_ptrs.is_contiguous() && + peer_kv_ptrs.scalar_type() == ScalarType::Long && + peer_kv_ptrs.dim() == 1 && peer_kv_ptrs.numel() == world_size, + "KV peer pointer table must be CUDA int64 [world_size]"); + STD_TORCH_CHECK( + peer_signal_ptrs.is_cuda() && peer_signal_ptrs.is_contiguous() && + peer_signal_ptrs.scalar_type() == ScalarType::Long && + peer_signal_ptrs.dim() == 1 && + peer_signal_ptrs.numel() == world_size, + "signal peer pointer table must be CUDA int64 [world_size]"); STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1"); STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank"); STD_TORCH_CHECK(buffer_slot == 0 || buffer_slot == 1, @@ -156,6 +251,8 @@ void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, int64_t device_index = local_kv.get_device_index(); STD_TORCH_CHECK(dst_rows.get_device_index() == device_index && + peer_kv_ptrs.get_device_index() == device_index && + peer_signal_ptrs.get_device_index() == device_index && received_kv.get_device_index() == device_index && received_signal.get_device_index() == device_index && completion.get_device_index() == device_index && @@ -176,10 +273,6 @@ void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, STD_TORCH_CHECK(vectorized, "direct DCP final-layout kv-gather requires 16-byte-aligned " "planes, rows, and pointers"); - STD_TORCH_CHECK(kv_mc_ptr != 0 && signal_mc_ptr != 0, - "direct DCP final-layout kv-gather requires multicast " - "pointers"); - const torch::stable::accelerator::DeviceGuard device_guard(device_index); cudaStream_t stream = get_current_cuda_stream(); increment_epoch_kernel<<<1, 1, 0, stream>>>( @@ -189,20 +282,39 @@ void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, int64_t items_per_row = row_bytes / sizeof(uint4); int64_t kv_c_items_per_row = kv_c_row_bytes / sizeof(uint4); int64_t item_count = num_tokens * items_per_row; - int64_t blocks = (item_count + kThreads - 1) / kThreads; - blocks = blocks < kMaxMulticastBlocks ? blocks : kMaxMulticastBlocks; - direct_dcp_kv_gather_multimem_kernel<<>>( - reinterpret_cast(local_kv.data_ptr()), - dst_rows.const_data_ptr(), - reinterpret_cast(static_cast(kv_mc_ptr)), - reinterpret_cast(static_cast(signal_mc_ptr)), - reinterpret_cast( - received_signal.const_data_ptr()), - epoch.const_data_ptr(), - reinterpret_cast(completion.mutable_data_ptr()), - world_size, rank, num_tokens, items_per_row, kv_c_items_per_row, - output_tokens, max_gathered_tokens, buffer_slot, - slot_stride_bytes / sizeof(uint4)); + int64_t blocks; + if (kv_mc_ptr != 0 && signal_mc_ptr != 0) { + blocks = (item_count + kThreads - 1) / kThreads; + blocks = blocks < kMaxMulticastBlocks ? blocks : kMaxMulticastBlocks; + direct_dcp_kv_gather_multimem_kernel<<>>( + reinterpret_cast(local_kv.data_ptr()), + dst_rows.const_data_ptr(), + reinterpret_cast(static_cast(kv_mc_ptr)), + reinterpret_cast(static_cast(signal_mc_ptr)), + reinterpret_cast( + received_signal.const_data_ptr()), + epoch.const_data_ptr(), + reinterpret_cast(completion.mutable_data_ptr()), + world_size, rank, num_tokens, items_per_row, kv_c_items_per_row, + output_tokens, max_gathered_tokens, buffer_slot, + slot_stride_bytes / sizeof(uint4)); + } else { + int64_t peer_item_count = world_size * item_count; + blocks = (peer_item_count + kThreads - 1) / kThreads; + blocks = blocks < kMaxPeerBlocks ? blocks : kMaxPeerBlocks; + direct_dcp_kv_gather_peer_kernel<<>>( + reinterpret_cast(local_kv.data_ptr()), + dst_rows.const_data_ptr(), + peer_kv_ptrs.const_data_ptr(), + peer_signal_ptrs.const_data_ptr(), + reinterpret_cast( + received_signal.const_data_ptr()), + epoch.const_data_ptr(), + reinterpret_cast(completion.mutable_data_ptr()), + world_size, rank, num_tokens, items_per_row, kv_c_items_per_row, + output_tokens, max_gathered_tokens, buffer_slot, + slot_stride_bytes / sizeof(uint4)); + } check_cuda_launch("direct DCP final-layout kv-gather"); } @@ -211,7 +323,8 @@ void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_kv_gather_ops) { direct_dcp_kv_gather_ops.def( "direct_dcp_kv_gather(" - "Tensor local_kv, Tensor dst_rows, Tensor! received_kv, " + "Tensor local_kv, Tensor dst_rows, Tensor peer_kv_ptrs, " + "Tensor peer_signal_ptrs, Tensor! received_kv, " "Tensor! received_signal, Tensor! completion, Tensor! epoch, " "int output_tokens, int plane_split_dim, int buffer_slot, " "int world_size, int rank, int max_gathered_tokens, " diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index 22dc9dbc63a3..a48f960fecdb 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -31,6 +31,20 @@ def _has_multicast_support() -> bool: return False +def _has_peer_access_support(device_count: int) -> bool: + if torch.cuda.device_count() < device_count: + return False + try: + return all( + source == destination + or torch.cuda.can_device_access_peer(source, destination) + for source in range(device_count) + for destination in range(device_count) + ) + except Exception: + return False + + def _dtype_from_name(dtype_name: str) -> torch.dtype: return { "float16": torch.float16, @@ -260,30 +274,9 @@ def test_kv_gather_flag_is_independent(self, monkeypatch): assert result is workspace - @pytest.mark.parametrize( - ("flag_name", "factory_name", "factory_args"), - [ - ( - "VLLM_USE_DIRECT_DCP_Q_GATHER", - "get_direct_dcp_q_gather_workspace", - (16, 2, 32, torch.bfloat16, 1), - ), - ( - "VLLM_USE_DIRECT_DCP_KV_GATHER", - "get_direct_dcp_kv_gather_workspace", - (64, 576, 512, torch.bfloat16, 1), - ), - ], - ) - def test_gather_requires_multicast( - self, - monkeypatch, - flag_name, - factory_name, - factory_args, - ): - factory = getattr(dcp_utils, factory_name) - monkeypatch.setenv(flag_name, "1") + def test_q_gather_requires_multicast(self, monkeypatch): + factory = dcp_utils.get_direct_dcp_q_gather_workspace + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "1") monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: False) factory.cache_clear() @@ -291,11 +284,39 @@ def test_gather_requires_multicast( factory( _FakeGroupCoordinator(), torch.device("cpu"), - *factory_args, + 16, + 2, + 32, + torch.bfloat16, + 1, ) is None ) + def test_kv_gather_accepts_peer_path_without_multicast(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "1") + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: False) + dcp_utils.get_direct_dcp_kv_gather_workspace.cache_clear() + workspace = object() + init_workspace = MagicMock(return_value=workspace) + monkeypatch.setattr( + dcp_utils, + "DirectDCPKVGatherWorkspace", + init_workspace, + ) + + result = dcp_utils.get_direct_dcp_kv_gather_workspace( + _FakeGroupCoordinator(), + torch.device("cpu"), + 64, + 576, + 512, + torch.bfloat16, + 1, + ) + + assert result is workspace + def test_kv_gather_rejects_invalid_workspace_geometry(self): with pytest.raises(ValueError, match="ubatch"): dcp_utils.DirectDCPKVGatherWorkspace( @@ -1141,8 +1162,8 @@ def check_case( pytest.param( 4, marks=pytest.mark.skipif( - torch.accelerator.device_count() < 4 or not _has_multicast_support(), - reason="Need 4 GPUs with symmetric-memory multicast.", + not _has_peer_access_support(4), + reason="Need 4 GPUs with mutual CUDA peer access.", ), ), ], diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py index 6c6e4bf20e87..b8b44e63b511 100644 --- a/tests/models/kimi_k3/test_mla_prefill_context.py +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -9,6 +9,7 @@ """ from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -135,6 +136,7 @@ def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: self._k_scale = k_scale self.kv_lora_rank = _KV_LORA_RANK self.num_local_heads = _NUM_HEADS + self.dcp_world_size = 1 self.qk_nope_head_dim = _QK_NOPE self.v_head_dim = _V_HEAD_DIM @@ -317,3 +319,62 @@ def test_fused_context_rejects_an_unquantized_query() -> None: ) with pytest.raises(AssertionError, match="new-token epilogue"): layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) + + +@torch.inference_mode() +def test_fused_context_consumes_direct_dcp_final_layout(monkeypatch) -> None: + """The K3 loop must consume compact DCP planes without rank-major reorg.""" + device = torch.device("cuda") + workspace = torch.empty((2, _ENTRY), device=device, dtype=torch.bfloat16) + kv_c = torch.randn((3, 1, _KV_LORA_RANK), device=device, dtype=torch.bfloat16) + k_pe = torch.randn((3, 1, _QK_ROPE), device=device, dtype=torch.bfloat16) + dcp_kv_gather = MagicMock(use_direct_kv_gather=True) + dcp_kv_gather.direct_kv_gather.return_value = (kv_c, k_pe) + chunk = SimpleNamespace( + request_slice=slice(0, 1), + padded_local_cu_seq_lens=torch.tensor([0, 2], device=device), + padded_local_token_to_seq=torch.zeros(2, dtype=torch.int32, device=device), + final_layout_dst_rows=torch.tensor([0, 2], dtype=torch.int32, device=device), + num_local_context_tokens=2, + num_context_tokens=3, + num_requests=1, + starts=torch.tensor([0], dtype=torch.int32, device=device), + index=5, + ) + prefill = SimpleNamespace( + block_table=torch.zeros((1, 1), dtype=torch.int32, device=device), + chunked_context=SimpleNamespace( + workspace=workspace, + dcp_manager=dcp_kv_gather, + ), + ) + gather_cache = MagicMock() + monkeypatch.setattr( + "vllm.models.kimi_k3.nvidia.mla.ops.cp_gather_cache", gather_cache + ) + layer = _FusedLayer( + _KVBProj(device, weight_dtype=torch.bfloat16), + torch.empty((1, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16), + "auto", + torch.ones(1, dtype=torch.float32, device=device), + ) + layer.dcp_world_size = 2 + + actual_kv_c, actual_k_pe = layer._gather_context_latent( + chunk, + layer.kv_cache, + prefill, + fp8_prefill=False, + ) + + assert actual_kv_c is kv_c + assert actual_k_pe is k_pe + gather_cache.assert_called_once() + dcp_kv_gather.direct_kv_gather.assert_called_once() + local_rows, dst_rows, output_tokens, slot = ( + dcp_kv_gather.direct_kv_gather.call_args.args + ) + assert local_rows.data_ptr() == workspace.data_ptr() + assert dst_rows is chunk.final_layout_dst_rows + assert output_tokens == 3 + assert slot == 1 diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index ffa92d4d3bd6..701211f25aa2 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -311,7 +311,7 @@ dcp_b12x_all_gather_heads, sanitize_dcp_attn_empty_rows, ) -from vllm.v1.attention.ops.dcp_utils import MLADCPManager +from vllm.v1.attention.ops.dcp_utils import MLADCPKVGather, MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states from vllm.v1.attention.selector import get_attn_backend from vllm.v1.kv_cache_interface import ( @@ -2525,7 +2525,7 @@ class ChunkedContextMetadata: chunks: "list[MLACommonPrefillMetadata.ContextChunk]" context_lens_list: list[int] empty_token_slices: list[slice] - dcp_manager: MLADCPManager | None = None + dcp_manager: MLADCPKVGather | None = None direct_dcp_kv_gather: bool = False block_table: torch.Tensor @@ -2878,7 +2878,7 @@ def build_mla_chunked_context_metadata( dcp_world_size: int, dcp_local_block_size: int, dcp_virtual_block_size: int, - dcp_manager: MLADCPManager | None = None, + dcp_manager: MLADCPKVGather | None = None, direct_dcp_kv_gather: bool = False, ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None": """Build chunked-context metadata for an MLA prefill. @@ -3141,9 +3141,9 @@ class MLACommonMetadataBuilder(AttentionMetadataBuilder[M]): # Whether this builder can flatten a non-causal query block into decode rows. supports_non_causal_multi_token_decode: ClassVar[bool] = False - # A decode backend that owns DCP query and output collectives can omit an - # MLADCPManager. Chunked-context prefill still gathers KV through the - # process DCP group before it runs the configured prefill backend. + # A decode backend that owns DCP query and output collectives can omit a + # full MLADCPManager. Chunked-context prefill still gets a standalone exact + # KV gather plan, including the final-layout PCIe publisher when available. supports_direct_dcp_kv_gather: ClassVar[bool] = False # The threshold for reordering the batch into decode and prefill requests. @@ -3290,7 +3290,7 @@ def __init__( ) use_packed_fp8_cache = vllm_config.cache_config.cache_dtype == "fp8_ds_mla" - self.dcp_manager: MLADCPManager | None = None + self.dcp_manager: MLADCPKVGather | None = None if self.dcp_world_size > 1: assert self.chunked_prefill_workspace_size % self.dcp_world_size == 0 workspace_dtype = ( @@ -3300,12 +3300,19 @@ def __init__( self.mla_dims.kv_lora_rank + self.mla_dims.qk_rope_head_dim ) self.dcp_manager = getattr(attention_layer, "dcp_manager", None) + if self.dcp_manager is None and self.supports_direct_dcp_kv_gather: + self.dcp_manager = MLADCPKVGather( + get_dcp_group(), + device, + parallel_config.num_ubatches, + ) use_direct_kv_gather = False if self.dcp_manager is not None: - if not isinstance(self.dcp_manager, MLADCPManager): + if not isinstance(self.dcp_manager, MLADCPKVGather): raise TypeError( - "MLA attention layer dcp_manager must be an " - f"MLADCPManager, got {type(self.dcp_manager).__name__}." + "MLA attention layer dcp_manager must provide an exact " + "DCP KV gather plan, got " + f"{type(self.dcp_manager).__name__}." ) use_direct_kv_gather = self.dcp_manager.init_kv_gather( self.chunked_prefill_workspace_size, @@ -3627,7 +3634,7 @@ def _gather_dcp_context_kv( gathered_kv: torch.Tensor, local_kv: torch.Tensor, *, - dcp_manager: MLADCPManager | None, + dcp_manager: MLADCPKVGather | None, direct_dcp_kv_gather: bool, ) -> None: """Gather one chunk of DCP-sharded context KV into caller storage.""" @@ -3989,9 +3996,7 @@ def _context_parallel_compute_prefill_context( allgather_offset : allgather_offset * (1 + dcp_world_size) ] assert toks * dcp_world_size <= cur_allgather_workspace.shape[0] - cur_allgather_kvcache = cur_allgather_workspace[ - : toks * dcp_world_size - ] + cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size] _gather_dcp_context_kv( cur_allgather_kvcache, local_gathered_kvcache, diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index a49a35cef66b..f557746400c9 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -1111,23 +1111,21 @@ def _compute_prefill_context( ) fp8_prefill = q.dtype == current_platform.fp8_dtype() - workspace = chunked_context.workspace kv_cache = self._attn_read_kv_cache() kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill) def run_chunk( chunk, out: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor]: - self._gather_context_latent(chunk, kv_cache, prefill, fp8_prefill) - gathered = workspace[: chunk.num_context_tokens] - kv_c_normed = gathered[..., : self.kv_lora_rank] + kv_c_normed, k_pe = self._gather_context_latent( + chunk, kv_cache, prefill, fp8_prefill + ) if kv_b_proj_input_dtype is not None: kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype) kv_nope = self.kv_b_proj(kv_c_normed)[0].view( -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim ) k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) - k_pe = gathered[..., self.kv_lora_rank :] if fp8_prefill: k, v = fused_mla_kv_concat_quant_fp8(k_nope, k_pe, v) else: @@ -1201,16 +1199,62 @@ def _gather_context_latent( kv_cache: torch.Tensor, prefill, fp8_prefill: bool, - ) -> None: + ) -> tuple[torch.Tensor, torch.Tensor]: """Gather one chunk's paged context latent into the workspace. - Dispatched exactly as in ``impl._compute_prefill_context``: an fp8 query - reads the plain fp8 cache in its stored layout, anything else lands in the - workspace as the model dtype. + Under DCP, gather the local paged shard and let the exact direct + publisher place every rank's valid rows straight into compact + request-major KV planes. Non-DCP keeps the merged-row workspace layout. """ - workspace = prefill.chunked_context.workspace - toks = chunk.num_context_tokens + chunked_context = prefill.chunked_context + assert chunked_context is not None + workspace = chunked_context.workspace block_table = prefill.block_table[chunk.request_slice] + if self.dcp_world_size > 1: + dcp_kv_gather = chunked_context.dcp_manager + assert dcp_kv_gather is not None and dcp_kv_gather.use_direct_kv_gather + assert chunk.padded_local_cu_seq_lens is not None + assert chunk.padded_local_token_to_seq is not None + assert chunk.final_layout_dst_rows is not None + toks = chunk.num_local_context_tokens + if self.kv_cache_dtype == "fp8_ds_mla": + ops.cp_gather_and_upconvert_fp8_kv_cache( + src_cache=kv_cache, + dst=workspace[:toks], + block_table=block_table, + workspace_starts=chunk.padded_local_cu_seq_lens, + batch_size=chunk.num_requests, + seq_starts=chunk.starts, + ) + elif is_quantized_kv_cache(self.kv_cache_dtype): + ops.gather_and_maybe_dequant_cache( + src_cache=kv_cache, + dst=workspace, + block_table=block_table, + cu_seq_lens=chunk.padded_local_cu_seq_lens, + token_to_seq=chunk.padded_local_token_to_seq, + num_tokens=toks, + kv_cache_dtype=self.kv_cache_dtype, + scale=self._k_scale, + seq_starts=chunk.starts, + ) + else: + ops.cp_gather_cache( + src_cache=kv_cache, + dst=workspace, + block_table=block_table, + cu_seq_lens=chunk.padded_local_cu_seq_lens, + batch_size=chunk.num_requests, + seq_starts=chunk.starts, + ) + return dcp_kv_gather.direct_kv_gather( + workspace[:toks], + chunk.final_layout_dst_rows, + chunk.num_context_tokens, + chunk.index & 1, + ) + + toks = chunk.num_context_tokens if self.kv_cache_dtype == "fp8_ds_mla": ops.cp_gather_and_upconvert_fp8_kv_cache( src_cache=kv_cache, @@ -1241,6 +1285,11 @@ def _gather_context_latent( batch_size=chunk.num_requests, seq_starts=chunk.starts, ) + gathered = workspace[:toks] + return ( + gathered[..., : self.kv_lora_rank], + gathered[..., self.kv_lora_rank :], + ) def _forward_prefill_fused( self, @@ -1256,9 +1305,9 @@ def _forward_prefill_fused( """Prefill using the fused key-concat + cache-insert kernel. Replaces ``_concat_k_nope_k_pe`` and the prefill cache write with one - fused kernel launch, dispatched by cache dtype. Chunked context runs - through this layer's ``_compute_prefill_context``, except under DCP where - it is delegated to the impl. + fused kernel launch, dispatched by cache dtype. Chunked context uses + this layer's fused packing loop for non-DCP and direct final-layout DCP; + only the NCCL rank-major fallback remains delegated to the impl. Supported configs (K3 fp8 policy): - bf16 cache -> bf16 prefill query @@ -1376,7 +1425,10 @@ def _forward_prefill_fused( # need to be live at the same time. out.copy_(suffix_output[..., : self.v_head_dim]) del output_prefill, suffix_output - if self.dcp_world_size > 1: + dcp_kv_gather = prefill.chunked_context.dcp_manager + if self.dcp_world_size > 1 and not ( + dcp_kv_gather is not None and dcp_kv_gather.use_direct_kv_gather + ): context_output, context_lse = ( self.impl._context_parallel_compute_prefill_context( # type: ignore[attr-defined] q, diff --git a/vllm/v1/attention/ops/dcp_utils.py b/vllm/v1/attention/ops/dcp_utils.py index de35a9e0ff65..23c94a770f68 100644 --- a/vllm/v1/attention/ops/dcp_utils.py +++ b/vllm/v1/attention/ops/dcp_utils.py @@ -538,17 +538,18 @@ def __init__( kv_shape = (num_ubatches, 2, max_gathered_tokens, token_dim) signal_shape = (num_ubatches, 2, self.world_size) - self.received_kv, _ = self._allocate(kv_shape, dtype) - self.received_signal, _ = self._allocate(signal_shape, torch.int32) + self.received_kv, self.peer_kv_ptrs = self._allocate(kv_shape, dtype) + self.received_signal, self.peer_signal_ptrs = self._allocate( + signal_shape, torch.int32 + ) kv_multicast_ptrs = self._multicast_ptrs(self.received_kv) signal_multicast_ptrs = self._multicast_ptrs(self.received_signal) self.multicast_ptrs = list( zip(kv_multicast_ptrs, signal_multicast_ptrs, strict=True) ) - if not all(kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs): - raise RuntimeError( - "Direct DCP kv-gather requires NVLS symmetric-memory multicast." - ) + self.uses_multicast = all( + kv_ptr and signal_ptr for kv_ptr, signal_ptr in self.multicast_ptrs + ) self.completion = self.received_signal.new_zeros((num_ubatches, 2)) torch.accelerator.synchronize() @@ -574,6 +575,8 @@ def gather( torch.ops._C.direct_dcp_kv_gather( local_kv, dst_rows, + self.peer_kv_ptrs[ubatch], + self.peer_signal_ptrs[ubatch], self.received_kv[ubatch], self.received_signal[ubatch], self.completion[ubatch], @@ -606,7 +609,7 @@ def get_direct_dcp_kv_gather_workspace( dtype: torch.dtype, num_ubatches: int, ) -> DirectDCPKVGatherWorkspace | None: - if not _direct_dcp_multicast_enabled( + if not _direct_dcp_enabled( group, dtype, envs.VLLM_USE_DIRECT_DCP_KV_GATHER, @@ -626,6 +629,110 @@ def get_direct_dcp_kv_gather_workspace( ) +class MLADCPKVGather: + """Own the exact DCP context-KV collective independently of decode DCP. + + Backends such as B12X own their decode query/output exchange, but chunked + prefill still needs to exchange compressed KV. Keeping that collective in + a small standalone object lets those backends use the final-layout + symmetric publisher without duplicating the decode collectives. + """ + + def __init__( + self, + group: GroupCoordinator, + device: torch.device, + num_ubatches: int, + ) -> None: + self.group = group + self.device = torch.device(device) + self.num_ubatches = max(num_ubatches, 1) + self._direct_kv_gather_workspace: DirectDCPKVGatherWorkspace | None = None + self._kv_gather: Callable[[torch.Tensor, torch.Tensor], object] | None = None + + def init_kv_gather( + self, + max_gathered_tokens: int, + token_dim: int, + plane_split_dim: int, + dtype: torch.dtype, + ) -> bool: + """Select the KV collective before allocating its local scratch. + + Returns whether the direct final-layout publisher was selected. That + path only needs one rank's local rows; the fallback additionally needs + a rank-major all-gather destination. + """ + world_size = self.group.world_size + if max_gathered_tokens <= 0 or max_gathered_tokens % world_size != 0: + raise ValueError( + "DCP KV gather capacity must be positive and divide evenly " + f"across {world_size} ranks, got {max_gathered_tokens}" + ) + if token_dim <= 0: + raise ValueError( + f"DCP KV gather token dimension must be positive: {token_dim}" + ) + + direct_workspace = get_direct_dcp_kv_gather_workspace( + self.group, + self.device, + max_gathered_tokens, + token_dim, + plane_split_dim, + dtype, + self.num_ubatches, + ) + self._direct_kv_gather_workspace = direct_workspace + if direct_workspace is not None: + transport = ( + "NVLS multicast" if direct_workspace.uses_multicast else "PCIe peer" + ) + logger.info_once( + "Using direct symmetric-memory DCP final-layout KV %s publisher " + "for MLA.", + transport, + ) + self._kv_gather = None + else: + self._kv_gather = functools.partial( + torch.distributed.all_gather_into_tensor, + group=self.group.device_group, + ) + return direct_workspace is not None + + @property + def use_direct_kv_gather(self) -> bool: + return self._direct_kv_gather_workspace is not None + + def kv_gather( + self, + gathered_kv: torch.Tensor, + local_kv: torch.Tensor, + ) -> object: + kv_gather = self._kv_gather + if kv_gather is None: + raise RuntimeError("NCCL DCP KV gather is not selected") + return kv_gather(gathered_kv, local_kv) + + def direct_kv_gather( + self, + local_kv: torch.Tensor, + dst_rows: torch.Tensor, + output_tokens: int, + buffer_slot: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + workspace = self._direct_kv_gather_workspace + if workspace is None: + raise RuntimeError("direct DCP KV gather is not enabled") + return workspace.gather( + local_kv, + dst_rows, + output_tokens, + buffer_slot, + ) + + class DCPCombine(Protocol): def __call__( self, @@ -637,11 +744,9 @@ def __call__( ) -> torch.Tensor: ... -class MLADCPManager: +class MLADCPManager(MLADCPKVGather): """Select and own layer-level collective implementations for MLA DCP.""" - _kv_gather: Callable[[torch.Tensor, torch.Tensor], object] | None - def __init__( self, vllm_config: VllmConfig, @@ -656,15 +761,14 @@ def __init__( use_pcp: bool, ) -> None: parallel_config = vllm_config.parallel_config - self.group = get_dcp_group() - self.device = torch.device(device) - self.num_ubatches = max(parallel_config.num_ubatches, 1) + super().__init__( + get_dcp_group(), + device, + parallel_config.num_ubatches, + ) self.max_num_tokens = get_dcp_workspace_max_num_tokens(vllm_config) self.use_a2a = parallel_config.dcp_comm_backend == "a2a" self.padded_num_heads = padded_num_heads - self._direct_kv_gather_workspace: DirectDCPKVGatherWorkspace | None = None - self._kv_gather = None - self.combine = self._init_combine( num_heads, output_head_dim, @@ -747,80 +851,3 @@ def _gather_query(self, query: torch.Tensor) -> torch.Tensor: if self.padded_num_heads is not None: query = reserve_query_head_storage(query, self.padded_num_heads) return query - - def init_kv_gather( - self, - max_gathered_tokens: int, - token_dim: int, - plane_split_dim: int, - dtype: torch.dtype, - ) -> bool: - """Select the KV collective before allocating its local scratch. - - Returns whether the direct final-layout publisher was selected. That - path only needs one rank's local rows; the fallback additionally needs - a rank-major all-gather destination. - """ - world_size = self.group.world_size - if max_gathered_tokens <= 0 or max_gathered_tokens % world_size != 0: - raise ValueError( - "DCP KV gather capacity must be positive and divide evenly " - f"across {world_size} ranks, got {max_gathered_tokens}" - ) - if token_dim <= 0: - raise ValueError( - f"DCP KV gather token dimension must be positive: {token_dim}" - ) - - direct_workspace = get_direct_dcp_kv_gather_workspace( - self.group, - self.device, - max_gathered_tokens, - token_dim, - plane_split_dim, - dtype, - self.num_ubatches, - ) - self._direct_kv_gather_workspace = direct_workspace - if direct_workspace is not None: - logger.info_once( - "Using direct symmetric-memory DCP final-layout KV multicast for MLA." - ) - self._kv_gather = None - else: - self._kv_gather = functools.partial( - torch.distributed.all_gather_into_tensor, - group=self.group.device_group, - ) - return direct_workspace is not None - - @property - def use_direct_kv_gather(self) -> bool: - return self._direct_kv_gather_workspace is not None - - def kv_gather( - self, - gathered_kv: torch.Tensor, - local_kv: torch.Tensor, - ) -> object: - kv_gather = self._kv_gather - if kv_gather is None: - raise RuntimeError("NCCL DCP KV gather is not selected") - return kv_gather(gathered_kv, local_kv) - - def direct_kv_gather( - self, - local_kv: torch.Tensor, - dst_rows: torch.Tensor, - output_tokens: int, - buffer_slot: int, - ) -> tuple[torch.Tensor, torch.Tensor]: - workspace = self._direct_kv_gather_workspace - if not self.use_direct_kv_gather or workspace is None: - raise RuntimeError("direct DCP KV gather is not enabled") - return workspace.gather( - local_kv, - dst_rows, - output_tokens, - buffer_slot, - ) From 1ce3580a1bbba938ffa8828b2cb1e015f264e847 Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Tue, 1 Sep 2026 19:31:14 +0900 Subject: [PATCH 47/52] perf(mla): decouple context tiles and enable SM120 FA4 Keep the externally visible scheduler/recompute cadence unchanged while allowing a larger transient MLA context tile. Add an opt-in Kimi 192/128 dense-prefill route to the existing SM120 FA4 kernel with one deterministic split, native asymmetric V output, and natural-log LSE for exact chunk merging. Global attention and decode backend selection remain unchanged. Rejected during implementation: FlashInfer modular CuTe DSL uses unsupported tcgen05 MMA on SM120; TRT-LLM ragged reports unsupported architecture; FMHAv2 DeepSeek separate-Q/K/V returned incorrect output. The vLLM SM120 FA4 path matches FA2 within 6.1e-5 output and 3e-6 LSE in focused oracles. Tests: 77 unit/registry tests, 61 Kimi context/kernel tests, and 3 focused SM120 FA4 GPU oracles passed. AI assistance was used. --- .../test_dcp_direct_a2a_lse_reduce.py | 14 ++ .../attention/test_kimi_k3_mla_sm120_fa4.py | 134 ++++++++++++++++++ .../test_mla_prefill_quant_output.py | 62 +++++++- vllm/envs.py | 13 ++ .../layers/attention/mla_attention.py | 16 ++- .../backends/mla/prefill/flash_attn.py | 47 +++++- 6 files changed, 279 insertions(+), 7 deletions(-) create mode 100644 tests/kernels/attention/test_kimi_k3_mla_sm120_fa4.py diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py index a48f960fecdb..268c35c89492 100644 --- a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -589,6 +589,15 @@ def test_mla_chunk_workspace_honors_configured_token_limit(monkeypatch): == 4096 ) + # The internal tile is deliberately independent of the externally visible + # scheduler/recompute budget and takes precedence over the legacy knob. + monkeypatch.setenv("VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE", "16384") + assert ( + MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config) + == 16384 + ) + monkeypatch.setenv("VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE", "0") + monkeypatch.setenv("VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE", "0") assert ( MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config) @@ -599,6 +608,11 @@ def test_mla_chunk_workspace_honors_configured_token_limit(monkeypatch): with pytest.raises(ValueError, match="must be non-negative"): MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config) + monkeypatch.setenv("VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE", "4096") + monkeypatch.setenv("VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE", "-1") + with pytest.raises(ValueError, match="INTERNAL_CONTEXT.*non-negative"): + MLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config) + @pytest.mark.parametrize("use_direct", [False, True]) def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch, use_direct): diff --git a/tests/kernels/attention/test_kimi_k3_mla_sm120_fa4.py b/tests/kernels/attention/test_kimi_k3_mla_sm120_fa4.py new file mode 100644 index 000000000000..a6e29cd61bdf --- /dev/null +++ b/tests/kernels/attention/test_kimi_k3_mla_sm120_fa4.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Exactness oracles for Kimi-K3's opt-in SM120 FA4 prefill path.""" + +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.attention.backends.mla.prefill.flash_attn import ( + FlashAttnPrefillBackend, +) +from vllm.vllm_flash_attn import flash_attn_varlen_func + +pytestmark = pytest.mark.skipif( + not current_platform.is_device_capability_family(120), + reason="requires consumer Blackwell SM120/SM121", +) + +_HEADS = 12 +_QK_DIM = 192 +_V_DIM = 128 + + +def _run( + version: int, + q: torch.Tensor, + k: torch.Tensor, + v: torch.Tensor, + indptr_q: torch.Tensor, + indptr_k: torch.Tensor, + max_q: int, + max_k: int, + *, + causal: bool, + out: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + if version == 2: + v = torch.nn.functional.pad(v, (0, _QK_DIM - _V_DIM)) + output, lse = flash_attn_varlen_func( + q, + k, + v, + cu_seqlens_q=indptr_q, + cu_seqlens_k=indptr_k, + max_seqlen_q=max_q, + max_seqlen_k=max_k, + softmax_scale=_QK_DIM**-0.5, + causal=causal, + return_softmax_lse=True, + out=out, + num_splits=1 if version == 4 else 0, + fa_version=version, + ) + return output[..., :_V_DIM], lse + + +@torch.inference_mode() +def test_sm120_fa4_kimi_context_matches_fa2_and_honors_out() -> None: + torch.manual_seed(23) + q_len, kv_len = 128, 2048 + q = torch.randn(q_len, _HEADS, _QK_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(kv_len, _HEADS, _QK_DIM, device="cuda", dtype=torch.bfloat16) + v = torch.randn(kv_len, _HEADS, _V_DIM, device="cuda", dtype=torch.bfloat16) + qo = torch.tensor([0, q_len], device="cuda", dtype=torch.int32) + kv = torch.tensor([0, kv_len], device="cuda", dtype=torch.int32) + + expected, expected_lse = _run(2, q, k, v, qo, kv, q_len, kv_len, causal=False) + out = torch.empty_like(expected) + actual, actual_lse = _run(4, q, k, v, qo, kv, q_len, kv_len, causal=False, out=out) + + assert actual.data_ptr() == out.data_ptr() + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(actual_lse, expected_lse, rtol=0, atol=3e-6) + + +@torch.inference_mode() +def test_sm120_fa4_kimi_varlen_causal_matches_fa2() -> None: + torch.manual_seed(29) + lengths = (1024, 512) + total = sum(lengths) + q = torch.randn(total, _HEADS, _QK_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn_like(q) + v = torch.randn(total, _HEADS, _V_DIM, device="cuda", dtype=torch.bfloat16) + indptr = torch.tensor([0, lengths[0], total], device="cuda", dtype=torch.int32) + + expected, expected_lse = _run( + 2, q, k, v, indptr, indptr, max(lengths), max(lengths), causal=True + ) + actual, actual_lse = _run( + 4, q, k, v, indptr, indptr, max(lengths), max(lengths), causal=True + ) + + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(actual_lse, expected_lse, rtol=0, atol=3e-6) + + +@torch.inference_mode() +def test_sm120_fa4_kimi_backend_context_wiring(monkeypatch) -> None: + monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "1") + q_len, kv_len = 64, 1024 + backend = FlashAttnPrefillBackend( + num_heads=_HEADS, + scale=_QK_DIM**-0.5, + kv_lora_rank=512, + qk_nope_head_dim=128, + qk_rope_head_dim=64, + v_head_dim=_V_DIM, + vllm_config=MagicMock(), + ) + assert backend.vllm_flash_attn_version == 4 + assert backend.supports_out() + + q = torch.randn(q_len, _HEADS, _QK_DIM, device="cuda", dtype=torch.bfloat16) + k = torch.randn(kv_len, _HEADS, _QK_DIM, device="cuda", dtype=torch.bfloat16) + v = torch.randn(kv_len, _HEADS, _V_DIM, device="cuda", dtype=torch.bfloat16) + qo = torch.tensor([0, q_len], device="cuda", dtype=torch.int32) + kv = torch.tensor([0, kv_len], device="cuda", dtype=torch.int32) + chunk = SimpleNamespace( + query_start_loc=qo, + cu_seq_lens=kv, + max_query_len=q_len, + max_seq_len=kv_len, + ) + out = torch.empty(q_len, _HEADS, _V_DIM, device="cuda", dtype=torch.bfloat16) + + actual, actual_lse = backend.run_prefill_context_chunk(chunk, q, k, v, out=out) + expected, expected_lse = _run(2, q, k, v, qo, kv, q_len, kv_len, causal=False) + + assert actual.data_ptr() == out.data_ptr() + torch.testing.assert_close(actual, expected, rtol=1e-3, atol=1e-3) + torch.testing.assert_close(actual_lse, expected_lse, rtol=0, atol=3e-6) diff --git a/tests/v1/attention/test_mla_prefill_quant_output.py b/tests/v1/attention/test_mla_prefill_quant_output.py index d7659485aa9f..77877d1a94db 100644 --- a/tests/v1/attention/test_mla_prefill_quant_output.py +++ b/tests/v1/attention/test_mla_prefill_quant_output.py @@ -10,7 +10,7 @@ + standalone static-FP8-quant path it replaces (GPU-only, SM100/SM110). """ -from unittest.mock import patch +from unittest.mock import MagicMock, patch import pytest import torch @@ -23,7 +23,9 @@ from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.backends.mla.prefill.base import MLAPrefillBackend from vllm.v1.attention.backends.mla.prefill.flash_attn import ( + FA4_MLA_PREFILL_KERNEL, FlashAttnPrefillBackend, + _use_sm120_fa4_prefill, ) _FA_MODULE = "vllm.v1.attention.backends.mla.prefill.flash_attn" @@ -115,6 +117,64 @@ def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): assert "output_scale" in base_params +def test_sm120_fa4_prefill_is_opt_in_and_shape_scoped(monkeypatch): + with patch(f"{_FA_MODULE}.current_platform") as plat: + plat.is_device_capability_family.return_value = True + monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "0") + assert not _use_sm120_fa4_prefill(192, 128) + + monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "1") + assert _use_sm120_fa4_prefill(192, 128) + assert not _use_sm120_fa4_prefill(256, 128) + assert not _use_sm120_fa4_prefill(192, 192) + + plat.is_device_capability_family.return_value = False + assert not _use_sm120_fa4_prefill(192, 128) + + +def test_sm120_fa4_prefill_forces_one_split_without_padding(): + backend = object.__new__(FlashAttnPrefillBackend) + backend.requires_v_padding = False + backend._is_vllm_fa = True + backend._sm120_fa4_prefill = True + output = torch.empty(2, 3, 128) + runtime = MagicMock(return_value=output) + backend.flash_attn_varlen_func = runtime + q = torch.empty(2, 3, 192) + k = torch.empty(4, 3, 192) + v = torch.empty(4, 3, 128) + + actual = backend._flash_attn_varlen_diff_headdims(q=q, k=k, v=v) + + assert actual is output + runtime.assert_called_once() + kwargs = runtime.call_args.kwargs + assert kwargs["v"] is v + assert kwargs["num_splits"] == 1 + + +def test_sm120_fa4_prefill_warmup_covers_lse(monkeypatch): + from vllm.model_executor.layers.attention import mla_attention + + monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "1") + config = MagicMock() + config.model_config.dtype = torch.bfloat16 + config.model_config.get_num_attention_heads.return_value = 12 + dims = MagicMock(qk_nope_head_dim=128, qk_rope_head_dim=64, v_head_dim=128) + with ( + patch(f"{_FA_MODULE}.current_platform") as plat, + patch.object(mla_attention, "get_mla_dims", return_value=dims), + ): + plat.is_device_capability.return_value = False + plat.is_device_capability_family.side_effect = lambda family: family == 120 + keys = FA4_MLA_PREFILL_KERNEL.get_warmup_keys(config) + + assert keys + assert {key.fa_version for key in keys} == {4} + assert {key.num_splits for key in keys} == {1} + assert {key.return_softmax_lse for key in keys} == {False, True} + + def test_mla_impl_forward_mha_accepts_output_scale(): """The abstract MLA impl forward_mha must carry output_scale so every override (and the unconditional forward_impl call) stays compatible.""" diff --git a/vllm/envs.py b/vllm/envs.py index 0ca61813913c..677ec8ec4aac 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -65,6 +65,8 @@ VLLM_B12X_ABSORB_BMM: bool = False VLLM_DSPARK_FP8_DRAFT_HEAD: bool = False VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE: int = 0 + VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE: int = 0 + VLLM_MLA_SM120_FA4_PREFILL: bool = False VLLM_K3_KV_GROUP_SIZE: int = 0 VLLM_DSPARK_DRAFT_KV_WINDOW: int = 0 VLLM_DSPARK_COMPACT_ROPE: bool = False @@ -1160,6 +1162,17 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE": lambda: int( os.getenv("VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE", "0") ), + # Override only the transient MLA context tile. Unlike the scheduler token + # budget, this does not change admission, preemption, prefix-commit, Mamba + # snapshot, or LMCache recompute granularity. + "VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE": lambda: int( + os.getenv("VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE", "0") + ), + # Opt in to the SM120 FA4 kernel for dense MLA prefill only. Decode and + # non-MLA attention retain their configured FlashAttention version. + "VLLM_MLA_SM120_FA4_PREFILL": lambda: bool( + int(os.getenv("VLLM_MLA_SM120_FA4_PREFILL", "0")) + ), # Bound the number of physical Kimi-K3 layers sharing each hybrid-cache # block table. Zero preserves the general grouping heuristic. "VLLM_K3_KV_GROUP_SIZE": lambda: int(os.getenv("VLLM_K3_KV_GROUP_SIZE", "0")), diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 701211f25aa2..820a836f1db4 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -3160,12 +3160,26 @@ def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int: model_config = vllm_config.model_config configured_workspace_size = envs.VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE + internal_workspace_size = envs.VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE if configured_workspace_size < 0: raise ValueError( "VLLM_MLA_CHUNKED_PREFILL_WORKSPACE_SIZE must be non-negative, " f"got {configured_workspace_size}." ) - if configured_workspace_size: + if internal_workspace_size < 0: + raise ValueError( + "VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE must be non-negative, " + f"got {internal_workspace_size}." + ) + if internal_workspace_size: + chunked_prefill_workspace_size = internal_workspace_size + logger.info_once( + "MLA internal context workspace is %d tokens; scheduler and " + "recompute budgets remain %d tokens.", + internal_workspace_size, + scheduler_config.max_num_batched_tokens, + ) + elif configured_workspace_size: chunked_prefill_workspace_size = configured_workspace_size else: chunked_prefill_workspace_size = min( diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 8c65edd64cdc..90fa67c307f0 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -10,6 +10,7 @@ import torch import vllm.envs as envs +from vllm.logger import init_logger from vllm.model_executor.layers.quantization.utils.quant_utils import ( kFp8StaticTensorSym, ) @@ -38,6 +39,8 @@ else: flash_attn_varlen_func = None # type: ignore[assignment] +logger = init_logger(__name__) + FA4_STANDARD_DTYPES = (torch.bfloat16, torch.float16) @@ -57,6 +60,15 @@ FA4_MLA_PREFILL_LSE_OPTIONS = (False, True) +def _use_sm120_fa4_prefill(qk_head_dim: int, v_head_dim: int) -> bool: + """Select the qualified SM120 FA4 MLA shape without changing global FA.""" + return ( + envs.VLLM_MLA_SM120_FA4_PREFILL + and current_platform.is_device_capability_family(120) + and (qk_head_dim, v_head_dim) == (192, 128) + ) + + @dataclass(frozen=True) class _FA4MLAPrefillShapeProbe: max_seqlen_q: int @@ -157,7 +169,15 @@ def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]: mla_dims = get_mla_dims(vllm_config.model_config) qk_head_dim = mla_dims.qk_nope_head_dim + mla_dims.qk_rope_head_dim - fa_version = get_flash_attn_version(head_size=qk_head_dim) + force_sm120_fa4 = _use_sm120_fa4_prefill(qk_head_dim, mla_dims.v_head_dim) + fa_version = ( + 4 + if force_sm120_fa4 + else get_flash_attn_version( + head_size=qk_head_dim, + head_size_v=mla_dims.v_head_dim, + ) + ) if fa_version != 4: return [] @@ -171,7 +191,7 @@ def get_warmup_keys(self, vllm_config: "VllmConfig") -> list[CompileKey]: if not (is_sm90 or is_sm100_family or is_sm120): return [] - num_splits = 1 if envs.VLLM_BATCH_INVARIANT else 0 + num_splits = 1 if envs.VLLM_BATCH_INVARIANT or force_sm120_fa4 else 0 if is_sm120 and num_splits != 1: return [] @@ -336,10 +356,23 @@ def __init__( "Ensure FlashAttnPrefillBackend.is_available() is checked first." ) qk_head_dim = qk_nope_head_dim + qk_rope_head_dim + self._sm120_fa4_prefill = _use_sm120_fa4_prefill(qk_head_dim, v_head_dim) self.flash_attn_varlen_func = flash_attn_varlen_func - self.vllm_flash_attn_version = get_flash_attn_version( - head_size=qk_head_dim, head_size_v=v_head_dim + self.vllm_flash_attn_version = ( + 4 + if self._sm120_fa4_prefill + else get_flash_attn_version( + head_size=qk_head_dim, + head_size_v=v_head_dim, + ) ) + if self._sm120_fa4_prefill: + logger.info_once( + "Using SM120 FA4 for dense MLA prefill (QK=%d, V=%d, " + "num_splits=1); global FlashAttention selection is unchanged.", + qk_head_dim, + v_head_dim, + ) if self.vllm_flash_attn_version is not None: self.flash_attn_varlen_func = functools.partial( flash_attn_varlen_func, fa_version=self.vllm_flash_attn_version @@ -399,7 +432,11 @@ def _flash_attn_varlen_diff_headdims( # called "return_attn_probs" instead of return_softmax_lse kwargs["return_attn_probs"] = return_softmax_lse assert out is None and output_scale is None - if envs.VLLM_BATCH_INVARIANT: + if self._sm120_fa4_prefill: + # The consumer-Blackwell kernel currently supports one deterministic + # split. This also avoids batch-shape-dependent split heuristics. + kwargs["num_splits"] = 1 + elif envs.VLLM_BATCH_INVARIANT: kwargs["num_splits"] = 1 attn_out = FA4_MLA_PREFILL_KERNEL( From 8a89a1d2d187c36c502f4e7f840a8eb4e47692c0 Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Tue, 1 Sep 2026 20:03:07 +0900 Subject: [PATCH 48/52] perf(kimi): retain large context projection workspace Reserve one BF16 kv_b_proj output buffer per ubatch before KV allocation and share it across Kimi MLA layers. Large internal context tiles write GEMM output directly into retained storage, while scheduler-sized tails keep the existing low-latency projection path. This removes 144 MiB hot-path allocations without changing attention math or external recompute cadence. Tests: ruff check/format; 11 Kimi context GPU tests; 12 Kimi Eagle3/model assembly tests. AI assistance was used. Co-authored-by: OpenAI Codex --- .../kimi_k3/test_mla_prefill_context.py | 61 +++++++++- vllm/models/kimi_k3/nvidia/mla.py | 107 +++++++++++++++++- vllm/models/kimi_k3/nvidia/model.py | 59 +++++++++- 3 files changed, 222 insertions(+), 5 deletions(-) diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py index b8b44e63b511..25e4dce6c01d 100644 --- a/tests/models/kimi_k3/test_mla_prefill_context.py +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -19,7 +19,10 @@ MLACommonPrefillMetadata, build_mla_chunked_context_metadata, ) -from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.mla import ( + KimiK3PrefillProjectionWorkspace, + MultiHeadLatentAttention, +) from vllm.platforms import current_platform pytestmark = pytest.mark.skipif( @@ -100,6 +103,8 @@ class _KVBProj(torch.nn.Module): def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: super().__init__() + from vllm.model_executor.layers.linear import UnquantizedLinearMethod + weight = ( torch.randn( _NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM), @@ -110,8 +115,13 @@ def __init__(self, device: torch.device, weight_dtype: torch.dtype) -> None: * 0.05 ) self.register_buffer("weight", weight.to(weight_dtype)) + self.quant_method = UnquantizedLinearMethod() + self.bias = None + self.gather_output = False + self.forward_calls = 0 def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, None]: + self.forward_calls += 1 if self.weight.dtype == torch.bfloat16 and x.dtype != torch.bfloat16: raise RuntimeError( "a bfloat16 kv_b_proj cannot consume the gathered latent as " @@ -129,7 +139,14 @@ class _FusedLayer: _gather_context_latent = MultiHeadLatentAttention._gather_context_latent _attn_read_kv_cache = MultiHeadLatentAttention._attn_read_kv_cache - def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: + def __init__( + self, + kv_b_proj, + kv_cache, + kv_cache_dtype, + k_scale, + prefill_projection_workspace=None, + ) -> None: self.kv_b_proj = kv_b_proj self.kv_cache = kv_cache self.kv_cache_dtype = kv_cache_dtype @@ -139,6 +156,7 @@ def __init__(self, kv_b_proj, kv_cache, kv_cache_dtype, k_scale) -> None: self.dcp_world_size = 1 self.qk_nope_head_dim = _QK_NOPE self.v_head_dim = _V_HEAD_DIM + self.prefill_projection_workspace = prefill_projection_workspace class _ReferenceImpl: @@ -252,7 +270,24 @@ def test_fused_context_matches_generic_impl( * 0.2 ).to(q_data_type) - layer = _FusedLayer(kv_b_proj, kv_cache, kv_cache_dtype, k_scale) + projection_workspace = None + if not kv_b_proj_quantized: + projection_workspace = KimiK3PrefillProjectionWorkspace( + num_ubatches=1, min_tokens=0 + ) + projection_workspace.reserve( + max_tokens=_WORKSPACE_TOKENS, + output_size=_NUM_HEADS * (_QK_NOPE + _V_HEAD_DIM), + dtype=torch.bfloat16, + device=device, + ) + layer = _FusedLayer( + kv_b_proj, + kv_cache, + kv_cache_dtype, + k_scale, + prefill_projection_workspace=projection_workspace, + ) fused_out, fused_lse = layer._compute_prefill_context( q, SimpleNamespace(prefill=prefill_fused) ) @@ -278,6 +313,10 @@ def test_fused_context_matches_generic_impl( ) torch.testing.assert_close(fused_out, ref_out, atol=0, rtol=0) torch.testing.assert_close(fused_lse, ref_lse, atol=0, rtol=0) + expected_projection_calls = len(prefill_ref.chunked_context.chunks) + if kv_b_proj_quantized: + expected_projection_calls *= 2 + assert kv_b_proj.forward_calls == expected_projection_calls # Every chunk but the continuation should have been written in place, i.e. # straight into the returned accumulator, with no intermediate copy. @@ -321,6 +360,22 @@ def test_fused_context_rejects_an_unquantized_query() -> None: layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) +def test_prefill_projection_workspace_reuses_reserved_storage() -> None: + device = torch.device("cuda", torch.cuda.current_device()) + workspace = KimiK3PrefillProjectionWorkspace(num_ubatches=1, min_tokens=4) + workspace.reserve(8, 16, torch.bfloat16, device) + + assert workspace.get(3, 16, torch.bfloat16, device) is None + short = workspace.get(4, 16, torch.bfloat16, device) + full = workspace.get(8, 16, torch.bfloat16, device) + assert short is not None and full is not None + assert short.data_ptr() == full.data_ptr() + assert workspace.nbytes == 8 * 16 * torch.bfloat16.itemsize + + with pytest.raises(ValueError, match="needs 9 rows"): + workspace.get(9, 16, torch.bfloat16, device) + + @torch.inference_mode() def test_fused_context_consumes_direct_dcp_final_layout(monkeypatch) -> None: """The K3 loop must consume compact DCP planes without rank-major reorg.""" diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index f557746400c9..d3f94d0b8f5c 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -114,6 +114,7 @@ SlidingWindowMLASpec, get_kv_quant_mode, ) +from vllm.v1.worker.ubatching import dbo_current_ubatch_id if TYPE_CHECKING: from vllm.model_executor.layers.attention.mla_attention import MLACommonMetadata @@ -127,6 +128,77 @@ _MLA_CALLER_OUTPUT_MIN_TOKENS = 1024 +class KimiK3PrefillProjectionWorkspace: + """Retained output storage for large dense context projections.""" + + def __init__(self, num_ubatches: int, min_tokens: int) -> None: + if num_ubatches < 1: + raise ValueError("num_ubatches must be positive") + if min_tokens < 0: + raise ValueError("min_tokens must be non-negative") + self.num_ubatches = num_ubatches + self.min_tokens = min_tokens + self._buffer: torch.Tensor | None = None + + def reserve( + self, + max_tokens: int, + output_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> None: + if max_tokens < self.min_tokens: + raise ValueError( + f"max_tokens ({max_tokens}) must be at least min_tokens " + f"({self.min_tokens})" + ) + self._buffer = torch.empty( + (self.num_ubatches, max_tokens, output_size), + dtype=dtype, + device=device, + ) + + @property + def nbytes(self) -> int: + buffer = self._buffer + return 0 if buffer is None else buffer.numel() * buffer.element_size() + + def get( + self, + num_tokens: int, + output_size: int, + dtype: torch.dtype, + device: torch.device, + ) -> torch.Tensor | None: + if num_tokens < self.min_tokens: + return None + buffer = self._buffer + if buffer is None: + raise RuntimeError("Kimi-K3 prefill projection workspace is not reserved") + if num_tokens > buffer.shape[1]: + raise ValueError( + f"context projection needs {num_tokens} rows, but the retained " + f"workspace has {buffer.shape[1]}" + ) + if output_size != buffer.shape[2]: + raise ValueError( + f"context projection needs {output_size} columns, but the retained " + f"workspace has {buffer.shape[2]}" + ) + if dtype != buffer.dtype or device != buffer.device: + raise ValueError( + "context projection input and retained workspace must have the " + "same dtype and device" + ) + ubatch_id = dbo_current_ubatch_id() + if ubatch_id >= self.num_ubatches: + raise RuntimeError( + f"ubatch {ubatch_id} has no Kimi-K3 prefill projection workspace; " + f"configured slots: {self.num_ubatches}" + ) + return buffer[ubatch_id, :num_tokens] + + def _parse_k3_qrep_layers(spec: str) -> frozenset[int] | None: if spec.strip().lower() == "all": return None @@ -295,6 +367,7 @@ def __init__( quant_config: QuantizationConfig | None = None, prefix: str = "", aux_stream: torch.cuda.Stream | None = None, + prefill_projection_workspace: KimiK3PrefillProjectionWorkspace | None = None, use_rope: bool = False, non_causal_multi_token_decode: bool = False, ) -> None: @@ -322,6 +395,7 @@ def __init__( self.scale = self.qk_head_dim**-0.5 self.rms_norm_eps = config.rms_norm_eps self.layer_name = prefix + self.prefill_projection_workspace = prefill_projection_workspace self.rotary_emb: RotaryEmbedding | None = None if use_rope: @@ -1114,6 +1188,37 @@ def _compute_prefill_context( kv_cache = self._attn_read_kv_cache() kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill) + def project_context(kv_c_normed: torch.Tensor) -> torch.Tensor: + workspace = self.prefill_projection_workspace + weight = getattr(self.kv_b_proj, "weight", None) + if workspace is None or not isinstance(weight, torch.Tensor): + return self.kv_b_proj(kv_c_normed)[0] + rows = kv_c_normed.numel() // self.kv_lora_rank + projection = workspace.get( + rows, + self.num_local_heads * (self.qk_nope_head_dim + self.v_head_dim), + kv_c_normed.dtype, + kv_c_normed.device, + ) + if projection is None: + return self.kv_b_proj(kv_c_normed)[0] + if not isinstance(self.kv_b_proj.quant_method, UnquantizedLinearMethod): + raise RuntimeError( + "Kimi-K3 retained context projection requires an " + "unquantized kv_b_proj" + ) + if self.kv_b_proj.bias is not None or self.kv_b_proj.gather_output: + raise RuntimeError( + "Kimi-K3 retained context projection requires a local, " + "bias-free kv_b_proj" + ) + torch.mm( + kv_c_normed.reshape(rows, self.kv_lora_rank), + weight.t(), + out=projection, + ) + return projection + def run_chunk( chunk, out: torch.Tensor | None = None ) -> tuple[torch.Tensor, torch.Tensor]: @@ -1122,7 +1227,7 @@ def run_chunk( ) if kv_b_proj_input_dtype is not None: kv_c_normed = kv_c_normed.to(kv_b_proj_input_dtype) - kv_nope = self.kv_b_proj(kv_c_normed)[0].view( + kv_nope = project_context(kv_c_normed).view( -1, self.num_local_heads, self.qk_nope_head_dim + self.v_head_dim ) k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1) diff --git a/vllm/models/kimi_k3/nvidia/model.py b/vllm/models/kimi_k3/nvidia/model.py index 30ed84ffc256..a3ab4eaf2edb 100644 --- a/vllm/models/kimi_k3/nvidia/model.py +++ b/vllm/models/kimi_k3/nvidia/model.py @@ -25,6 +25,9 @@ from vllm.forward_context import get_forward_context, is_forward_context_available from vllm.logger import init_logger from vllm.model_executor.layers.activation import SiluAndMul, SituAndMul +from vllm.model_executor.layers.attention.mla_attention import ( + align_mla_chunked_context_workspace_size, +) from vllm.model_executor.layers.fused_moe import ( FusedMoEFactory, fused_moe_make_expert_params_mapping, @@ -111,7 +114,10 @@ from vllm.models.kimi_k3.nvidia.low_latency_gemm import ( enable_kimi_k3_low_latency_gemm, ) -from vllm.models.kimi_k3.nvidia.mla import MultiHeadLatentAttention +from vllm.models.kimi_k3.nvidia.mla import ( + KimiK3PrefillProjectionWorkspace, + MultiHeadLatentAttention, +) from vllm.models.kimi_k3.nvidia.ops import attn_res from vllm.models.kimi_k3.nvidia.tp_projection import ( gather_kimi_sharded_projection, @@ -1467,6 +1473,7 @@ def __init__( vllm_config: VllmConfig, prefix: str = "", aux_stream: torch.cuda.Stream | None = None, + prefill_projection_workspace: KimiK3PrefillProjectionWorkspace | None = None, ) -> None: super().__init__() self.hidden_size = config.hidden_size @@ -1537,6 +1544,7 @@ def __init__( quant_config=quant_config, prefix=f"{prefix}.self_attn", aux_stream=aux_stream, + prefill_projection_workspace=prefill_projection_workspace, ) self._self_attn_writes_output = False @@ -1791,6 +1799,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): config = vllm_config.model_config.hf_text_config self.config = config + self._vllm_config = vllm_config self.attn_res_block_size: int | None = config.attn_res_block_size self.use_attn_res = self.attn_res_block_size is not None self.reuse_attn_res_output = ( @@ -1820,6 +1829,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): # attention front-end (DeepseekV4 convention: created at the model # level and threaded into each attention layer). aux_stream = torch.cuda.Stream() + self._mla_prefill_projection_workspace = KimiK3PrefillProjectionWorkspace( + num_ubatches=2 if parallel_config.enable_dbo else 1, + min_tokens=int(vllm_config.scheduler_config.max_num_batched_tokens) + 1, + ) def get_layer(prefix: str): return KimiDecoderLayer( @@ -1827,6 +1840,7 @@ def get_layer(prefix: str): vllm_config, prefix, aux_stream=aux_stream, + prefill_projection_workspace=self._mla_prefill_projection_workspace, ) self.start_layer, self.end_layer, self.layers = make_layers( @@ -2059,6 +2073,48 @@ def reserve_attn_res_workspace(self) -> None: self._max_num_batched_tokens, ) + def reserve_mla_prefill_projection_workspace(self) -> None: + """Reserve one large context projection output shared by MLA layers.""" + internal_tokens = envs.VLLM_MLA_INTERNAL_CONTEXT_WORKSPACE_SIZE + if internal_tokens <= self._max_num_batched_tokens: + return + workspace_tokens = align_mla_chunked_context_workspace_size( + self._vllm_config, internal_tokens + ) + mla_layers = [ + layer.self_attn + for layer in self.layers + if isinstance(getattr(layer, "self_attn", None), MultiHeadLatentAttention) + ] + if not mla_layers: + return + first = mla_layers[0] + if envs.VLLM_BATCH_INVARIANT or not all( + isinstance(layer.kv_b_proj.quant_method, UnquantizedLinearMethod) + and layer.kv_b_proj.bias is None + and not layer.kv_b_proj.gather_output + for layer in mla_layers + ): + logger.warning_once( + "Kimi-K3 retained context projection is unavailable for the " + "configured kv_b_proj method." + ) + return + weight = first.kv_b_proj.weight + _release_cuda_cache_before_retained_allocation(weight.device) + self._mla_prefill_projection_workspace.reserve( + max_tokens=workspace_tokens, + output_size=weight.shape[0], + dtype=weight.dtype, + device=weight.device, + ) + logger.info_once( + "Kimi-K3 retained %.2f MiB/rank for the %d-token MLA context " + "projection workspace shared across layers.", + self._mla_prefill_projection_workspace.nbytes / (1024**2), + workspace_tokens, + ) + def forward( self, input_ids: torch.Tensor | None, @@ -2533,6 +2589,7 @@ def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: return loaded def process_weights_after_loading(self) -> None: + self.model.reserve_mla_prefill_projection_workspace() self.model.reserve_attn_res_workspace() From 6da7e417834437fd8238a1ca09d0f9950c2f460f Mon Sep 17 00:00:00 2001 From: myshytf <9619163+myshytf@users.noreply.github.com> Date: Wed, 2 Sep 2026 07:42:08 +0900 Subject: [PATCH 49/52] perf(dcp): rotate destinations in the PCIe peer KV gather The peer publisher walked its work destination-major: every rank pushed its whole context-KV slice to destination 0, then 1, and so on. All ranks run the same schedule at the same time, so the eight sources funnelled into one destination's inbound PCIe link while the other seven idled. On the Kimi-K3 TP8/DCP8 target (PCIe Gen4 x16) a 24,576-token context piece took 3.8 ms (~3.7 GB/s per rank against ~25 GB/s links), 273 ms per 1,536-token prefill chunk at 90k context, growing linearly with context. Destinations are now assigned per thread block and rotated by the source rank, so every inbound link receives from a different source at any time; the host rounds the block count up to a multiple of world_size so each destination owns a block set. Payload, layout, completion counting and epoch signalling are unchanged, so the gathered bytes are identical. Validation: tests/distributed/test_dcp_direct_a2a_lse_reduce.py::test_distributed_direct_kv_gather_matches_reference (4-GPU byte-exact oracle incl. CUDA-graph replay) run against this kernel as a standalone extension on the target's GPUs; serving numbers in the deployment ledger. --- .../dcp_utils/dcp_direct_kv_gather.cu | 28 +++++++++++++------ 1 file changed, 19 insertions(+), 9 deletions(-) diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu index 240a0e16e524..b2fa4725ca22 100644 --- a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu @@ -98,6 +98,13 @@ __global__ void direct_dcp_kv_gather_multimem_kernel( // publishes its valid rows directly into every peer's compact, request-major // planes. This moves the same payload volume as an all-gather while avoiding // the rank-major materialization and the subsequent reorganization pass. +// +// Destinations are assigned per thread block and rotated by the source rank +// (block b serves destination (b + rank) mod world_size). All ranks run the +// same schedule at the same time, so with a destination-major walk every +// source would push into one destination's inbound link while the other +// links idle; the rotation keeps every inbound link busy with a different +// source. The host launches a block count that is a multiple of world_size. __global__ void direct_dcp_kv_gather_peer_kernel( const uint4* local_kv, const int32_t* dst_rows, const int64_t* peer_kv_ptrs, const int64_t* peer_signal_ptrs, @@ -108,13 +115,15 @@ __global__ void direct_dcp_kv_gather_peer_kernel( int64_t buffer_slot, int64_t slot_stride_items) { uint32_t epoch = static_cast(epoch_ptr[0]); int64_t source_items = num_tokens * items_per_row; - int64_t total_items = world_size * source_items; - int64_t item_stride = static_cast(gridDim.x) * blockDim.x; - for (int64_t linear = - static_cast(blockIdx.x) * blockDim.x + threadIdx.x; - linear < total_items; linear += item_stride) { - int64_t destination_rank = linear / source_items; - int64_t item = linear - destination_rank * source_items; + int64_t destination_rank = + (static_cast(blockIdx.x) + rank) % world_size; + int64_t blocks_per_destination = static_cast(gridDim.x) / world_size; + int64_t block_in_destination = static_cast(blockIdx.x) / world_size; + int64_t item_stride = blocks_per_destination * blockDim.x; + uint4* peer_kv = get_peer_ptr(peer_kv_ptrs, destination_rank) + + buffer_slot * slot_stride_items; + for (int64_t item = block_in_destination * blockDim.x + threadIdx.x; + item < source_items; item += item_stride) { int64_t src_row = item / items_per_row; int32_t dst_row = dst_rows[src_row]; if (dst_row < 0) { @@ -138,8 +147,7 @@ __global__ void direct_dcp_kv_gather_peer_kernel( static_cast(dst_row) * k_pe_items_per_row + row_item - kv_c_items_per_row; } - uint4* peer_kv = get_peer_ptr(peer_kv_ptrs, destination_rank); - peer_kv[buffer_slot * slot_stride_items + dst_item] = local_kv[item]; + peer_kv[dst_item] = local_kv[item]; } // Every block publishes its system-scope payload before contributing to the @@ -302,6 +310,8 @@ void direct_dcp_kv_gather(const torch::stable::Tensor& local_kv, int64_t peer_item_count = world_size * item_count; blocks = (peer_item_count + kThreads - 1) / kThreads; blocks = blocks < kMaxPeerBlocks ? blocks : kMaxPeerBlocks; + // One block set per destination: round up to a multiple of world_size. + blocks = ((blocks + world_size - 1) / world_size) * world_size; direct_dcp_kv_gather_peer_kernel<<>>( reinterpret_cast(local_kv.data_ptr()), dst_rows.const_data_ptr(), From 118a8523b0945cb626439058583ec691ef6500a0 Mon Sep 17 00:00:00 2001 From: g0san Date: Thu, 3 Sep 2026 04:58:29 +0900 Subject: [PATCH 50/52] fix(kquant): return owned storage from the W4A8 prefill launch The W4A8 prefill launch writes its rows into the hybrid runtime's shared prefill output buffer and applied the coupled outer transform in place, so `_apply_once` returned a view of storage that every hybrid layer reuses. The W4A16 paths never had this property: their fp32 result is converted to the model dtype into a fresh tensor. The outer transform now writes into a fresh tensor of the same shape (it is the last pass over the rows, so this replaces the in-place write rather than adding a copy), and the routed result no longer aliases the shared buffer. Numerics are unchanged; the eager prefill path gains one allocation per launch from the caching allocator. Validation: tests/quantization/test_kquant_hybrid.py in the SM120 production image. Co-Authored-By: Claude Code Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D --- vllm/model_executor/layers/quantization/kquant_hybrid.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/layers/quantization/kquant_hybrid.py b/vllm/model_executor/layers/quantization/kquant_hybrid.py index 5580dbfac052..2975f4ab33d3 100644 --- a/vllm/model_executor/layers/quantization/kquant_hybrid.py +++ b/vllm/model_executor/layers/quantization/kquant_hybrid.py @@ -1991,9 +1991,15 @@ def _apply_once( # normalize both contracts for downstream layers. out_trellis = fused_moe.run(binding=binding)[:m] if use_w4a8_prefill: + # The kernel wrote model-dtype rows into the runtime's shared + # prefill output buffer; the outer transform is the last pass + # over them, so it writes into a fresh tensor and the routed + # result never aliases storage that the next layer's launch + # overwrites. The W4A16 paths get the same ownership from the + # fp32 -> model-dtype conversion below. out_trellis = run_w4a8_coupled_outer_transform( out_trellis, - out_trellis, + torch.empty_like(out_trellis), prepared_value.down_svh, output_transform=True, ) From 444e101ced7bd0705e22f8b733a89b944525ef98 Mon Sep 17 00:00:00 2001 From: g0san Date: Thu, 3 Sep 2026 05:09:45 +0900 Subject: [PATCH 51/52] fix(mla): bucket K3 verify plans, cover the local shard, own fp8 context output Three dense-MLA metadata and output-storage fixes for the fused DCP verification path: - Verify plans (fp8 KV, four-query tiles) are created per power-of-two batch capacity (`_dense_mla_plan_row_caps`) and `build` selects the smallest covering capacity, like the decode plans; the batch range is bounded by the flattened row capacity (four rows per request). One plan per batch value grew linearly with max_num_seqs and exceeded the 1,024-row plan limit from batch 257. - The plan's page table must cover the largest local KV shard: `build` copies the worker's block table into the plan-width flattened table and drops columns past that width (KV-block rounding can make the worker table wider while no local sequence references those columns); a plan narrower than the shard would drop referenced pages, so the builder now rejects it (a sliding-window spec shrinking the plan) instead of clamping. - `_reuse_consumed_query_for_context_output` allocates fresh storage when the consumed query holds fewer bytes than the compact bf16 context output (an fp8 Kimi-K3 query row is 192 bytes, the output row 256), instead of raising on every fp8 prefill with chunked context. Validation: tests/v1/attention/test_b12x_mla.py (38 passed, new covering- bucket test) and tests/models/kimi_k3/test_mla_padding.py (14 passed; the fp8 case now uses the production 192-wide query) in the SM120 image. Co-Authored-By: Claude Code Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D --- tests/models/kimi_k3/test_mla_padding.py | 22 +++++--- tests/v1/attention/test_b12x_mla.py | 61 ++++++++++++++++++++++ vllm/models/kimi_k3/nvidia/mla.py | 25 ++++++--- vllm/v1/attention/backends/mla/b12x_mla.py | 43 ++++++++++++--- 4 files changed, 130 insertions(+), 21 deletions(-) diff --git a/tests/models/kimi_k3/test_mla_padding.py b/tests/models/kimi_k3/test_mla_padding.py index cf137217e41a..89469e7f2e10 100644 --- a/tests/models/kimi_k3/test_mla_padding.py +++ b/tests/models/kimi_k3/test_mla_padding.py @@ -234,30 +234,38 @@ def write_active_prefill(*args): torch.testing.assert_close(output[2:], torch.zeros_like(output[2:])) -@pytest.mark.parametrize("query_dtype", [torch.bfloat16, torch.float8_e4m3fn]) -def test_kimi_mla_context_output_reuses_consumed_query_bytes(query_dtype): +@pytest.mark.parametrize( + ("query_dtype", "aliases_query"), + [(torch.bfloat16, True), (torch.float8_e4m3fn, False)], +) +def test_kimi_mla_context_output_reuses_consumed_query_bytes( + query_dtype, aliases_query +): + """The Kimi-K3 query row (192 elements) backs the bf16 context output + row (128 elements) only in bf16; an fp8 query holds 192 bytes per row + against the 256 the output needs, so it gets its own storage.""" from vllm.models.kimi_k3.nvidia import mla - query = torch.empty((4, 2, 256), dtype=query_dtype) + query = torch.empty((4, 2, 192), dtype=query_dtype) output = torch.randn((4, 2, 128), dtype=torch.bfloat16) compact = mla._reuse_consumed_query_for_context_output(query, output) compact.copy_(output) - assert compact.data_ptr() == query.data_ptr() + assert (compact.data_ptr() == query.data_ptr()) is aliases_query assert compact.shape == output.shape assert compact.dtype == output.dtype assert compact.is_contiguous() torch.testing.assert_close(compact, output) -def test_kimi_mla_context_output_rejects_insufficient_query_storage(): +def test_kimi_mla_context_output_requires_a_contiguous_query(): from vllm.models.kimi_k3.nvidia import mla - query = torch.empty((4, 2, 64), dtype=torch.bfloat16) + query = torch.empty((4, 2, 384), dtype=torch.bfloat16)[..., ::2] output = torch.empty((4, 2, 128), dtype=torch.bfloat16) - with pytest.raises(ValueError, match="too small"): + with pytest.raises(ValueError, match="contiguous"): mla._reuse_consumed_query_for_context_output(query, output) diff --git a/tests/v1/attention/test_b12x_mla.py b/tests/v1/attention/test_b12x_mla.py index 8f41cac1fbf9..78f36473882c 100644 --- a/tests/v1/attention/test_b12x_mla.py +++ b/tests/v1/attention/test_b12x_mla.py @@ -518,6 +518,67 @@ def test_b12x_mla_builder_preserves_tiled_q4_dcp_verification( assert getattr(result, "dense_mla_flat_block_table", None) is None +def test_b12x_mla_builder_selects_the_covering_verify_plan(monkeypatch) -> None: + """Verify plans are bucketed by power-of-two batch capacity; a batch of + three uses the four-request plan, and a batch beyond every capacity + keeps the flattened decode path.""" + builder = object.__new__(B12xMLAMetadataBuilder) + builder._dense_mla_plan = _FakePlan() + builder._dense_mla_plans = {32: _FakePlan()} + plans = { + batch: SimpleNamespace( + caps=SimpleNamespace(max_page_table_width=4), batch=batch + ) + for batch in (1, 2, 4) + } + builder._dense_mla_verify_plans = plans + builder._dense_mla_scratch = torch.empty(256, dtype=torch.uint8) + builder._dense_mla_padded_q = None + builder._dense_mla_padded_output = None + builder._max_dense_mla_rows = 32 + builder._dense_mla_flat_block_table = torch.zeros(32, 4, dtype=torch.int32) + builder._dense_mla_flat_seq_lens = torch.empty(32, dtype=torch.int32) + builder._dense_mla_flat_query_start_loc = torch.arange(33, dtype=torch.int32) + builder._dense_mla_causal_offsets = torch.arange(-3, 1, dtype=torch.int32) + builder._dense_mla_flat_global_seq_lens = torch.empty(32, dtype=torch.int32) + builder._dense_mla_flat_dcp_remainder = torch.empty(32, dtype=torch.int32) + builder.dcp_world_size = 1 + builder._dcp_rank = 0 + builder.cp_kv_cache_interleave_size = 1 + + def metadata_for(num_decodes: int): + return SimpleNamespace( + causal=True, + num_decodes=num_decodes, + num_decode_tokens=4 * num_decodes, + decode=SimpleNamespace( + block_table=torch.arange( + 4 * num_decodes, dtype=torch.int32 + ).view(num_decodes, 4), + seq_lens=torch.full((num_decodes,), 4, dtype=torch.int32), + dcp_tot_seq_lens=None, + ), + ) + + monkeypatch.setattr( + b12x_mla.MLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata_for(3), + ) + result = builder.build(0, SimpleNamespace()) + assert result.dense_mla_plan is plans[4] + assert result.dense_mla_verify_block_table.shape == (3, 4) + + monkeypatch.setattr( + b12x_mla.MLACommonMetadataBuilder, + "build", + lambda *args, **kwargs: metadata_for(5), + ) + result = builder.build(0, SimpleNamespace()) + assert result.dense_mla_plan is builder._dense_mla_plans[32] + assert result.dense_mla_flat_block_table.shape == (20, 4) + + def test_b12x_mla_builder_bounds_single_token_draft_table(monkeypatch) -> None: builder = object.__new__(B12xMLAMetadataBuilder) builder._dense_mla_plan = _FakePlan() diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index d3f94d0b8f5c..a3b7083386e3 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -295,17 +295,30 @@ def _reuse_consumed_query_for_context_output( query: torch.Tensor, output: torch.Tensor, ) -> torch.Tensor: - """Return contiguous semantic-output storage backed by a consumed query.""" + """Return contiguous storage for the compact context output. + + The consumed query backs the output when it holds enough bytes (a bf16 + query row is wider than its bf16 output row); otherwise the output gets + fresh storage, as for an fp8 query whose 192-byte head rows cannot hold + the 256-byte bf16 output rows. + + Args: + query: The contiguous prefill query the attention kernels have consumed. + output: The output tensor whose shape and dtype the storage must match. + + Returns: + A contiguous tensor shaped like ``output``, aliasing ``query`` when it + fits and newly allocated otherwise. + + Raises: + ValueError: If ``query`` is not contiguous. + """ if not query.is_contiguous(): raise ValueError("Kimi-K3 MLA prefill query storage must be contiguous") required_bytes = output.numel() * output.element_size() query_bytes = query.view(torch.uint8).flatten() if query_bytes.numel() < required_bytes: - raise ValueError( - "Kimi-K3 MLA prefill query storage is too small for compact context " - f"output: available={query_bytes.numel()} bytes, " - f"required={required_bytes} bytes" - ) + return torch.empty(output.shape, dtype=output.dtype, device=output.device) return query_bytes[:required_bytes].view(output.dtype).view_as(output) diff --git a/vllm/v1/attention/backends/mla/b12x_mla.py b/vllm/v1/attention/backends/mla/b12x_mla.py index 98c017fd5afc..08447840afc8 100644 --- a/vllm/v1/attention/backends/mla/b12x_mla.py +++ b/vllm/v1/attention/backends/mla/b12x_mla.py @@ -314,12 +314,22 @@ def __init__( self._max_dense_mla_rows = max_dense_mla_rows self._effective_heads = self.num_heads * self.dcp_world_size self._kernel_heads = _kernel_query_heads(self.num_heads, self.dcp_world_size) - max_cache_tokens = _max_dcp_local_cache_tokens( + local_shard_tokens = _max_dcp_local_cache_tokens( vllm_config, dcp_size=self.dcp_world_size ) + max_cache_tokens = local_shard_tokens sliding_window = getattr(kv_cache_spec, "sliding_window", None) if sliding_window is not None: max_cache_tokens = min(max_cache_tokens, int(sliding_window)) + if max_cache_tokens < local_shard_tokens: + # The kernel attends to every local token of a request; the plan's + # page table (and the flattened copy `build` makes of the worker's + # block table) must therefore cover the largest local shard. + raise ValueError( + "B12X_MLA plans must cover the largest local KV shard: " + f"planned={max_cache_tokens} tokens, shard={local_shard_tokens} " + f"(sliding_window={sliding_window})." + ) self._dense_mla_plans = { rows: _create_dense_mla_plan( vllm_config, @@ -332,8 +342,19 @@ def __init__( ) for rows in _dense_mla_plan_row_caps(max_dense_mla_rows) } + # Four-query verify tiles (fp8 KV): one plan per power-of-two batch + # capacity, selected by the smallest covering capacity at build time + # like the decode plans. A batch needs four flattened rows per + # request, so the batch range is bounded by the row capacity. self._dense_mla_verify_plans: dict[int, Any] = {} - if _planned_kv_dtype(vllm_config) == torch.float8_e4m3fn: + max_verify_batch = min( + int(vllm_config.scheduler_config.max_num_seqs), + max_dense_mla_rows // 4, + ) + if ( + _planned_kv_dtype(vllm_config) == torch.float8_e4m3fn + and max_verify_batch >= 1 + ): self._dense_mla_verify_plans = { batch: _create_dense_mla_plan( vllm_config, @@ -347,10 +368,7 @@ def __init__( dcp_size=self.dcp_world_size, max_cache_tokens=max_cache_tokens, ) - for batch in range( - 1, - int(vllm_config.scheduler_config.max_num_seqs) + 1, - ) + for batch in _dense_mla_plan_row_caps(max_verify_batch) } self._dense_mla_plan = self._dense_mla_plans[max_dense_mla_rows] workspace_specs = [ @@ -565,8 +583,15 @@ def build( ) verify_plans = getattr(self, "_dense_mla_verify_plans", {}) tiled_verify = ( - metadata.causal and query_len == 4 and metadata.num_decodes in verify_plans + metadata.causal + and query_len == 4 + and bool(verify_plans) + and metadata.num_decodes <= max(verify_plans) ) + # The worker's block table can be wider than the plan's page table + # when the KV block size rounds the per-request row up; every local + # sequence fits the plan (it covers the largest local shard), so the + # columns past the plan width are never referenced and are dropped. if tiled_verify: verify_table = self._dense_mla_flat_block_table[: metadata.num_decodes] source_width = min( @@ -574,7 +599,9 @@ def build( int(verify_table.shape[1]), ) verify_table[:, :source_width].copy_(source_table[:, :source_width]) - metadata.dense_mla_plan = verify_plans[metadata.num_decodes] + metadata.dense_mla_plan = _select_dense_mla_plan( + verify_plans, metadata.num_decodes + ) metadata.dense_mla_verify_block_table = verify_table metadata.dense_mla_query_cache_seq_lens = flat_lens return metadata From 14c7b33606eb1be31a2c6daea09c1e2e55c9f33f Mon Sep 17 00:00:00 2001 From: g0san Date: Thu, 3 Sep 2026 05:11:58 +0900 Subject: [PATCH 52/52] fix(kimi-k3): fall back from an unreserved projection workspace, guard SM120 FA4 - The retained context-projection workspace is consulted only when the model could have reserved it (local, bias-free, unquantized kv_b_proj outside batch-invariant mode); every other configuration projects through kv_b_proj. A chunk at or above the workspace's row threshold used to reach `workspace.get` on an unreserved workspace and raise. - `VLLM_MLA_SM120_FA4_PREFILL=1` raises with the import reason when the FA4 CuTe interface is unavailable instead of running with `fa_version=4` and failing at the first prefill; the selector documents its contract. - `copyChunk8UnitFp8` skips the bf16 conversion on pre-Ampere targets like `copyChunk8`, so builds that include sm_75 compile (the bf16 converter is defined only for `__CUDA_ARCH__ >= 800`). - Test helper annotation: `_build_prefill_metadata` returns the metadata and the block count. Validation (SM120 image): tests/models/kimi_k3/test_mla_prefill_context.py 12 passed (new unreserved-workspace fallback test), tests/v1/attention/test_mla_prefill_quant_output.py passed (new missing-FA4 test); the fused key-concat kernel compiles with nvcc for sm_120 and sm_75. Co-Authored-By: Claude Code Claude-Session: https://claude.ai/code/session_01HPWxmKzfikaemyykd3p89D --- ..._kimi_k3_mla_key_concat_kv_cache_kernel.cu | 9 ++++ .../kimi_k3/test_mla_prefill_context.py | 41 ++++++++++++++++++- .../test_mla_prefill_quant_output.py | 18 +++++++- vllm/models/kimi_k3/nvidia/mla.py | 29 +++++++------ .../backends/mla/prefill/flash_attn.py | 29 ++++++++++++- 5 files changed, 109 insertions(+), 17 deletions(-) diff --git a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu index db4805dde0fd..592657c0c3e9 100644 --- a/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu +++ b/csrc/libtorch_stable/fused_kimi_k3_mla_key_concat_kv_cache_kernel.cu @@ -198,6 +198,12 @@ template __device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst, const scalar_t* src) { #ifndef USE_ROCM +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) + // _typeConvert is unavailable on pre-Ampere, as in copyChunk8. + if constexpr (std::is_same_v) { + return; + } else { +#endif uint4 const input = *reinterpret_cast(src); using Converter = vllm::_typeConvert; auto const* input2 = @@ -216,6 +222,9 @@ __device__ __forceinline__ void copyChunk8UnitFp8(uint8_t* dst, } } *reinterpret_cast(dst) = output; +#if (!defined(__CUDA_ARCH__) || __CUDA_ARCH__ < 800) + } +#endif #else copyChunk8(dst, src, 1.0f); #endif diff --git a/tests/models/kimi_k3/test_mla_prefill_context.py b/tests/models/kimi_k3/test_mla_prefill_context.py index 25e4dce6c01d..c8a3d3426ce8 100644 --- a/tests/models/kimi_k3/test_mla_prefill_context.py +++ b/tests/models/kimi_k3/test_mla_prefill_context.py @@ -181,7 +181,7 @@ def _build_prefill_metadata( workspace_dtype: torch.dtype, q_data_type: torch.dtype, backend: _RecordingPrefillBackend, -) -> MLACommonPrefillMetadata: +) -> tuple[MLACommonPrefillMetadata, int]: query_start_loc_cpu = torch.zeros(len(_QUERY_LENS) + 1, dtype=torch.int32) query_start_loc_cpu[1:] = torch.tensor(_QUERY_LENS, dtype=torch.int32).cumsum(0) workspace = torch.empty( @@ -360,6 +360,45 @@ def test_fused_context_rejects_an_unquantized_query() -> None: layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) +@torch.inference_mode() +def test_fused_context_projects_through_kv_b_proj_when_unreserved() -> None: + """A workspace that was never reserved (the model skips the reservation + for a gathering, biased or quantized kv_b_proj) must not be consulted: + chunks at or above its row threshold project through the layer.""" + torch.manual_seed(0) + device = torch.device("cuda") + kv_b_proj = _KVBProj(device, weight_dtype=torch.bfloat16) + kv_b_proj.gather_output = True + k_scale = torch.ones(1, dtype=torch.float32, device=device) + backend = _RecordingPrefillBackend() + prefill, num_blocks = _build_prefill_metadata( + device, torch.bfloat16, torch.bfloat16, backend + ) + kv_cache = torch.randn( + (num_blocks, _BLOCK_SIZE, _ENTRY), device=device, dtype=torch.bfloat16 + ) + q = torch.randn( + (sum(_QUERY_LENS), _NUM_HEADS, _QK_NOPE + _QK_ROPE), + device=device, + dtype=torch.bfloat16, + ) + layer = _FusedLayer( + kv_b_proj, + kv_cache, + "auto", + k_scale, + prefill_projection_workspace=KimiK3PrefillProjectionWorkspace( + num_ubatches=1, min_tokens=1 + ), + ) + + out, lse = layer._compute_prefill_context(q, SimpleNamespace(prefill=prefill)) + + assert kv_b_proj.forward_calls == len(backend.calls) > 0 + assert out.shape[0] == lse.shape[-1] == sum(_QUERY_LENS) + assert torch.isfinite(out).all() + + def test_prefill_projection_workspace_reuses_reserved_storage() -> None: device = torch.device("cuda", torch.cuda.current_device()) workspace = KimiK3PrefillProjectionWorkspace(num_ubatches=1, min_tokens=4) diff --git a/tests/v1/attention/test_mla_prefill_quant_output.py b/tests/v1/attention/test_mla_prefill_quant_output.py index 77877d1a94db..0ec5800eb67c 100644 --- a/tests/v1/attention/test_mla_prefill_quant_output.py +++ b/tests/v1/attention/test_mla_prefill_quant_output.py @@ -118,13 +118,15 @@ def test_flash_attn_prefill_backend_signature_accepts_fused_kwargs(): def test_sm120_fa4_prefill_is_opt_in_and_shape_scoped(monkeypatch): + fa_interface = "vllm.vllm_flash_attn.flash_attn_interface" with patch(f"{_FA_MODULE}.current_platform") as plat: plat.is_device_capability_family.return_value = True monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "0") assert not _use_sm120_fa4_prefill(192, 128) monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "1") - assert _use_sm120_fa4_prefill(192, 128) + with patch(f"{fa_interface}.FA4_AVAILABLE", True): + assert _use_sm120_fa4_prefill(192, 128) assert not _use_sm120_fa4_prefill(256, 128) assert not _use_sm120_fa4_prefill(192, 192) @@ -132,6 +134,20 @@ def test_sm120_fa4_prefill_is_opt_in_and_shape_scoped(monkeypatch): assert not _use_sm120_fa4_prefill(192, 128) +def test_sm120_fa4_prefill_refuses_a_missing_fa4_interface(monkeypatch): + """The opt-in fails instead of silently running another FA version.""" + fa_interface = "vllm.vllm_flash_attn.flash_attn_interface" + monkeypatch.setenv("VLLM_MLA_SM120_FA4_PREFILL", "1") + with ( + patch(f"{_FA_MODULE}.current_platform") as plat, + patch(f"{fa_interface}.FA4_AVAILABLE", False), + patch(f"{fa_interface}.FA4_UNAVAILABLE_REASON", "cute missing"), + ): + plat.is_device_capability_family.return_value = True + with pytest.raises(RuntimeError, match="cute missing"): + _use_sm120_fa4_prefill(192, 128) + + def test_sm120_fa4_prefill_forces_one_split_without_padding(): backend = object.__new__(FlashAttnPrefillBackend) backend.requires_v_padding = False diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index a3b7083386e3..54c40fec5fe3 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -1201,11 +1201,24 @@ def _compute_prefill_context( kv_cache = self._attn_read_kv_cache() kv_b_proj_input_dtype = _get_kv_b_proj_input_dtype(self.kv_b_proj, fp8_prefill) + # The retained workspace is reserved only for a local, bias-free, + # unquantized kv_b_proj outside batch-invariant mode; every other + # configuration projects through the layer's own path. + retained_projection = ( + self.prefill_projection_workspace is not None + and isinstance(getattr(self.kv_b_proj, "weight", None), torch.Tensor) + and not envs.VLLM_BATCH_INVARIANT + and isinstance(self.kv_b_proj.quant_method, UnquantizedLinearMethod) + and self.kv_b_proj.bias is None + and not self.kv_b_proj.gather_output + ) + def project_context(kv_c_normed: torch.Tensor) -> torch.Tensor: - workspace = self.prefill_projection_workspace - weight = getattr(self.kv_b_proj, "weight", None) - if workspace is None or not isinstance(weight, torch.Tensor): + if not retained_projection: return self.kv_b_proj(kv_c_normed)[0] + workspace = self.prefill_projection_workspace + assert workspace is not None + weight = self.kv_b_proj.weight rows = kv_c_normed.numel() // self.kv_lora_rank projection = workspace.get( rows, @@ -1215,16 +1228,6 @@ def project_context(kv_c_normed: torch.Tensor) -> torch.Tensor: ) if projection is None: return self.kv_b_proj(kv_c_normed)[0] - if not isinstance(self.kv_b_proj.quant_method, UnquantizedLinearMethod): - raise RuntimeError( - "Kimi-K3 retained context projection requires an " - "unquantized kv_b_proj" - ) - if self.kv_b_proj.bias is not None or self.kv_b_proj.gather_output: - raise RuntimeError( - "Kimi-K3 retained context projection requires a local, " - "bias-free kv_b_proj" - ) torch.mm( kv_c_normed.reshape(rows, self.kv_lora_rank), weight.t(), diff --git a/vllm/v1/attention/backends/mla/prefill/flash_attn.py b/vllm/v1/attention/backends/mla/prefill/flash_attn.py index 90fa67c307f0..d227d201e42b 100644 --- a/vllm/v1/attention/backends/mla/prefill/flash_attn.py +++ b/vllm/v1/attention/backends/mla/prefill/flash_attn.py @@ -61,12 +61,37 @@ def _use_sm120_fa4_prefill(qk_head_dim: int, v_head_dim: int) -> bool: - """Select the qualified SM120 FA4 MLA shape without changing global FA.""" - return ( + """Select the qualified SM120 FA4 MLA shape without changing global FA. + + Args: + qk_head_dim: Combined nope + rope query/key head width. + v_head_dim: Value head width. + + Returns: + True when ``VLLM_MLA_SM120_FA4_PREFILL`` is set on an SM120-family + device and the shape is the qualified (192, 128) Kimi-K3 layout. + + Raises: + RuntimeError: If the selection is requested but the FA4 CuTe interface + cannot be imported; the opt-in never falls back silently. + """ + selected = bool( envs.VLLM_MLA_SM120_FA4_PREFILL and current_platform.is_device_capability_family(120) and (qk_head_dim, v_head_dim) == (192, 128) ) + if selected: + from vllm.vllm_flash_attn.flash_attn_interface import ( + FA4_AVAILABLE, + FA4_UNAVAILABLE_REASON, + ) + + if not FA4_AVAILABLE: + raise RuntimeError( + "VLLM_MLA_SM120_FA4_PREFILL=1 requires the FA4 CuTe interface: " + f"{FA4_UNAVAILABLE_REASON}" + ) + return selected @dataclass(frozen=True)