From 79c865b838e34f7a98a936771284773819d79c8f Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 10 Aug 2026 10:52:10 -0700 Subject: [PATCH 01/30] [Perf] Narrow DeepSeek V4 eager CUDA graph region (#51430) Signed-off-by: Woosuk Kwon Co-authored-by: OpenAI Codex --- vllm/models/deepseek_v4/amd/model.py | 2 +- vllm/models/deepseek_v4/attention.py | 183 ++++++++++++------------ vllm/models/deepseek_v4/nvidia/model.py | 2 +- 3 files changed, 90 insertions(+), 97 deletions(-) diff --git a/vllm/models/deepseek_v4/amd/model.py b/vllm/models/deepseek_v4/amd/model.py index d626febdc0dc..d7fd88c79930 100644 --- a/vllm/models/deepseek_v4/amd/model.py +++ b/vllm/models/deepseek_v4/amd/model.py @@ -551,7 +551,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.rms_norm_eps = config.rms_norm_eps # Three aux streams: one per non-default input GEMM in - # DeepseekV4Attention.attn_gemm_parallel_execute + # DeepseekV4Attention._run_parallel_input_projections # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. # Disable them on ROCm because of hang issues. diff --git a/vllm/models/deepseek_v4/attention.py b/vllm/models/deepseek_v4/attention.py index c56a36bcd2ff..286ea11bf3bf 100644 --- a/vllm/models/deepseek_v4/attention.py +++ b/vllm/models/deepseek_v4/attention.py @@ -362,11 +362,10 @@ def forward( device=hidden_states.device, ) - # Metadata-independent input GEMMs + RMSNorm stay in the captured - # graph; the metadata-dependent rest (q up-proj + kv-insert, indexer, - # compressor, MLA attention) runs in the eager break. + # Keep the attention input preparation in the captured graph. Only the + # sparse indexer and MLA attention run in the eager break below. qr_kv, kv_score, indexer_kv_score, indexer_weights = ( - self.attn_gemm_parallel_execute(hidden_states) + self._run_parallel_input_projections(hidden_states) ) qr, kv = qr_kv.split([self.q_lora_rank, self.head_dim], dim=-1) qr, kv = fused_q_kv_rmsnorm( @@ -377,16 +376,62 @@ def forward( self.eps, ) - # attention_impl is wrapped with @eager_break_during_capture: this is - # where the breakable cudagraph capture breaks (the attention op runs - # eagerly between captured graph segments). - self.attention_impl( + attn_metadata = get_forward_context().attn_metadata + indexer = self.indexer + compressor = self.compressor + aux_streams = self.aux_stream_list + + def project_query_and_cache_kv() -> torch.Tensor: + q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) + return self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) + + index_q: torch.Tensor | None = None + index_q_scale: torch.Tensor | None = None + index_weights_out: torch.Tensor | None = None + + # Keep Q projection and KV insertion on the default stream. The indexer + # and MLA compressor use aux streams 0 and 1; aux 2 is internal to the + # indexer. ROCm runs the same work sequentially without aux streams. + if indexer is not None: + assert compressor is not None + q, (indexer_inputs, _) = execute_in_parallel( + project_query_and_cache_kv, + [ + lambda: indexer( + hidden_states, + qr, + indexer_kv_score, + indexer_weights, + positions, + self.indexer_rotary_emb, + ), + lambda: compressor(kv_score, positions, self.rotary_emb), + ], + self.ln_events[0], + [self.ln_events[1], self.ln_events[2]], + [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, + enable=aux_streams is not None, + ) + index_q, index_q_scale, index_weights_out = indexer_inputs + elif compressor is not None: + aux_stream = aux_streams[0] if aux_streams is not None else None + q, _ = maybe_execute_in_parallel( + project_query_and_cache_kv, + lambda: compressor(kv_score, positions, self.rotary_emb), + self.ln_events[0], + self.ln_events[1], + aux_stream, + ) + else: + q = project_query_and_cache_kv() + + self._sparse_indexer_and_attn( hidden_states, - qr, + index_q, + index_q_scale, + index_weights_out, + q, kv, - kv_score, - indexer_kv_score, - indexer_weights, positions, o_padded, ) @@ -395,12 +440,14 @@ def forward( # Inverse-RoPE + wo_a + wo_b output projection (platform-specific). return self._o_proj(o, positions) - def _fused_wqa_wkv_gemm(self, hidden_states: torch.Tensor) -> torch.Tensor: - # MergedColumnParallelLinear returns (output, bias); bias is None. - qr_kv, _ = self.fused_wqa_wkv(hidden_states) - return qr_kv - - def attn_gemm_parallel_execute(self, hidden_states) -> tuple[Any, ...]: + def _run_parallel_input_projections( + self, hidden_states: torch.Tensor + ) -> tuple[ + torch.Tensor, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + ]: aux_streams = self.aux_stream_list if aux_streams is not None: assert len(aux_streams) >= 3 @@ -443,11 +490,8 @@ def indexer_compressor_kv_score() -> torch.Tensor: aux_fns[1] = indexer_weights_proj aux_fns[2] = indexer_compressor_kv_score - def fused_wqa_wkv() -> torch.Tensor: - return self._fused_wqa_wkv_gemm(hidden_states) - qr_kv, (kv_score, indexer_weights, indexer_kv_score) = execute_in_parallel( - fused_wqa_wkv, + lambda: self.fused_wqa_wkv(hidden_states)[0], aux_fns, self.ln_events[0], self.ln_events[1:4], @@ -459,81 +503,26 @@ def fused_wqa_wkv() -> torch.Tensor: return qr_kv, kv_score, indexer_kv_score, indexer_weights @eager_break_during_capture - def attention_impl( + def _sparse_indexer_and_attn( self, hidden_states: torch.Tensor, - qr: torch.Tensor, + index_q: torch.Tensor | None, + index_q_scale: torch.Tensor | None, + index_weights: torch.Tensor | None, + q: torch.Tensor, kv: torch.Tensor, - kv_score: torch.Tensor, - indexer_kv_score: torch.Tensor, - indexer_weights: torch.Tensor, positions: torch.Tensor, - out: torch.Tensor, # [num_tokens, padded_heads, head_dim], written in place + out: torch.Tensor, ) -> None: - forward_context = get_forward_context() - attn_metadata = forward_context.attn_metadata - - # wq_b + kv_insert (+ MLA compressor when an indexer is present) ride - # on the default stream so q stays on its consumer stream (forward_mqa - # downstream reads q on default). Indexer/compressor go on aux for - # overlap with default's GEMM + cache write. - if self.indexer is not None: - aux_streams = self.aux_stream_list - indexer = self.indexer - # Local ref so the closure keeps a non-None type for mypy. - assert self.compressor is not None - compressor = self.compressor - - def wq_b_kv_insert() -> torch.Tensor: - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - return q - - # 3-way overlap (matches TRT-LLM PR #14142 Level 1): default runs - # wq_b+kv_insert; slot [0] runs the full indexer; slot [1] runs the - # MLA compressor. Slot [2] is reserved for the indexer's inner - # overlap. ROCm (aux_streams is None) falls back to sequential. - q, _ = execute_in_parallel( - wq_b_kv_insert, - [ - lambda: indexer( - hidden_states, - qr, - indexer_kv_score, - indexer_weights, - positions, - self.indexer_rotary_emb, - ), - lambda: compressor(kv_score, positions, self.rotary_emb), - ], - self.ln_events[0], - [self.ln_events[1], self.ln_events[2]], - [aux_streams[0], aux_streams[1]] if aux_streams is not None else None, - enable=aux_streams is not None, + if self.indexer is not None and index_q is not None: + assert index_weights is not None + q_quant = (index_q, index_q_scale) if index_q_scale is not None else index_q + self.indexer.indexer_op( + hidden_states, + q_quant, + None, + index_weights, ) - elif self.compressor is not None: - # wq_b + kv_insert on default, compressor on aux. - aux_stream = ( - self.aux_stream_list[0] if self.aux_stream_list is not None else None - ) - compressor = self.compressor - - def wq_b_kv_insert() -> torch.Tensor: - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) - return q - - q, _ = maybe_execute_in_parallel( - wq_b_kv_insert, - lambda: compressor(kv_score, positions, self.rotary_emb), - self.ln_events[0], - self.ln_events[1], - aux_stream, - ) - else: - # SWA-only layer: no compressor, no overlap. - q = self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim) - q = self._fused_qnorm_rope_kv_insert(q, kv, positions, attn_metadata) # MLA attention writes into the pre-allocated `out` buffer # ([num_tokens, padded_heads, head_dim]). @@ -836,7 +825,7 @@ def forward( indexer_weights: torch.Tensor, positions: torch.Tensor, rotary_emb: nn.Module, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor | None, torch.Tensor | None, torch.Tensor | None]: compressor = self.compressor attn_metadata = get_forward_context().attn_metadata @@ -860,7 +849,7 @@ def forward( PADDED_TOP_K=triton.next_power_of_2(self.topk_tokens), num_warps=8, ) - return self.topk_indices_buffer + return None, None, None def wq_b_and_q_quant(): # ReplicatedLinear returns (output, bias); bias is None. @@ -882,11 +871,15 @@ def wq_b_and_q_quant(): # compressor returns None and writes K to the indexer KV cache; the # join orders that write before indexer_op (skip_k_cache_insert=True). - (q_quant, weights), k = maybe_execute_in_parallel( + (q_quant, weights), _ = maybe_execute_in_parallel( wq_b_and_q_quant, lambda: compressor(compressed_kv_score, positions, rotary_emb), self.ln_events[0], self.ln_events[1], self.aux_stream, ) - return self.indexer_op(hidden_states, q_quant, k, weights) + if isinstance(q_quant, tuple): + q, q_scale = q_quant + else: + q, q_scale = q_quant, None + return q, q_scale, weights diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index 8925890a3cd4..c0d29d0b42a7 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1018,7 +1018,7 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): self.rms_norm_eps = config.rms_norm_eps # Three aux streams: one per non-default input GEMM in - # DeepseekV4Attention.attn_gemm_parallel_execute + # DeepseekV4Attention._run_parallel_input_projections # (compressor kv_score, indexer.weights_proj, indexer.compressor # kv_score). fused_wqa_wkv stays on the default stream. aux_stream_list = [torch.cuda.Stream() for _ in range(3)] From 405bc86768d1267f05dd2d4139048eb30a84b6cb Mon Sep 17 00:00:00 2001 From: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:55:35 +0800 Subject: [PATCH 02/30] [Perf] Launch the top-k/top-p Triton sampler kernel with 8 warps (#51507) Signed-off-by: BabyDrangoner <148877251+BabyDrangoner@users.noreply.github.com> Signed-off-by: Nick Hill Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com> Co-authored-by: Nick Hill Co-authored-by: Claude Opus 5 (1M context) --- vllm/v1/sample/ops/topk_topp_triton.py | 7 +++++++ 1 file changed, 7 insertions(+) diff --git a/vllm/v1/sample/ops/topk_topp_triton.py b/vllm/v1/sample/ops/topk_topp_triton.py index c284ff61876b..7d593e69cd7f 100755 --- a/vllm/v1/sample/ops/topk_topp_triton.py +++ b/vllm/v1/sample/ops/topk_topp_triton.py @@ -931,12 +931,18 @@ def apply_top_k_top_p_triton( # Smaller tiles compile and run faster on CPU; GPU benefits from larger tiles. # On XPU, large BLOCK_SIZE causes precision loss in the single-pass pivot # approximation; use smaller tiles for accurate top-p results. + launch_kwargs = {} if logits.device.type == "cpu": block_size, block_size_trunc = 256, 128 elif logits.device.type == "xpu": block_size, block_size_trunc = 4096, 2048 else: block_size, block_size_trunc = 8192, 4096 + # Each program serially sweeps the vocab row in BLOCK_SIZE tiles, so + # per-tile latency bounds kernel latency, and Triton's default of 4 + # warps leaves an 8192-wide tile at 16 elements per lane. 8 warps is + # faster on every arch measured (SM90, SM100, SM120, gfx950); 16 is not. + launch_kwargs["num_warps"] = 8 _topk_topp_kernel[(NUM_PROGRAMS,)]( logits, @@ -953,6 +959,7 @@ def apply_top_k_top_p_triton( BLOCK_SIZE_TRUNC=block_size_trunc, TOPK_ENABLED=topk_enabled, TOPP_ENABLED=topp_enabled, + **launch_kwargs, ) return logits From 05f0a80016a4a12c5f4e241e92e233acd1c2c0b0 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Mon, 10 Aug 2026 14:29:20 -0400 Subject: [PATCH 03/30] [Bugfix][Structured Output] Mask request stop tokens in xgrammar until grammar terminates (#49227) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- .../test_backend_xgrammar_stop_tokens.py | 82 +++++++++++++++++++ vllm/v1/structured_output/__init__.py | 9 +- vllm/v1/structured_output/backend_guidance.py | 5 +- .../backend_lm_format_enforcer.py | 5 +- vllm/v1/structured_output/backend_outlines.py | 5 +- vllm/v1/structured_output/backend_types.py | 8 +- vllm/v1/structured_output/backend_xgrammar.py | 6 +- 7 files changed, 114 insertions(+), 6 deletions(-) create mode 100644 tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py diff --git a/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py b/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py new file mode 100644 index 000000000000..e910d29aa38c --- /dev/null +++ b/tests/v1/structured_output/test_backend_xgrammar_stop_tokens.py @@ -0,0 +1,82 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Regression tests for issue #42403. + +A request's stop-token set (the generation_config eos list plus any +user-supplied ``stop_token_ids``) is invisible to xgrammar, which only knows +the tokenizer's single eos. Such tokens can therefore escape the grammar +bitmask while the FSM is still mid-object and truncate structured output. +``compile_grammar`` now forwards ``all_stop_token_ids`` to the matcher as +``override_stop_tokens`` so xgrammar masks them until the grammar completes. +""" + +import pytest +from transformers import AutoTokenizer + +from vllm.config import StructuredOutputsConfig, VllmConfig +from vllm.v1.structured_output.backend_types import StructuredOutputOptions +from vllm.v1.structured_output.backend_xgrammar import XgrammarBackend + +TOKENIZER = "openai-community/gpt2" +VOCAB_SIZE = 50257 + +# gpt2 token ids used to drive a `{"type": "string"}` grammar deterministically. +EOS = 50256 # <|endoftext|> -- the tokenizer's only default stop token +QUOTE = 1 # standalone `"`; opens then closes the JSON string +LETTER = 55 # `X`: valid string content, not a special/stop token by default + + +def _token_allowed(row, token_id: int) -> bool: + word = int(row[token_id // 32].item()) & 0xFFFFFFFF + return bool(word & (1 << (token_id % 32))) + + +@pytest.fixture(scope="module") +def backend() -> XgrammarBackend: + vllm_config = VllmConfig( + structured_outputs_config=StructuredOutputsConfig(backend="xgrammar") + ) + return XgrammarBackend( + vllm_config, + tokenizer=AutoTokenizer.from_pretrained(TOKENIZER), + vocab_size=VOCAB_SIZE, + ) + + +def test_request_stop_tokens_gated_to_grammar_terminal(backend: XgrammarBackend): + schema = '{"type": "string"}' + default = backend.compile_grammar(StructuredOutputOptions.JSON, schema) + override = backend.compile_grammar( + StructuredOutputOptions.JSON, schema, stop_token_ids={EOS, LETTER} + ) + + # Open the string: both grammars are now in a non-terminal state. + for grammar in (default, override): + assert grammar.accept_tokens("req", [QUOTE]) + + bm_default = backend.allocate_token_bitmask(1) + bm_override = backend.allocate_token_bitmask(1) + default.fill_bitmask(bm_default, 0) + override.fill_bitmask(bm_override, 0) + + # Mid-string, the plain token is valid content, so the default grammar + # leaves it samplable -- this is the leak. Registering it as a stop token + # masks it until the grammar can terminate. + assert _token_allowed(bm_default[0], LETTER) + assert not _token_allowed(bm_override[0], LETTER) + + # Close the string -> accepting state (grammar complete, not yet terminated). + for grammar in (default, override): + assert grammar.accept_tokens("req", [QUOTE]) + assert not grammar.is_terminated() + + default.fill_bitmask(bm_default, 0) + override.fill_bitmask(bm_override, 0) + + # The extra stop token may now terminate under the override, never under + # the default grammar -- and the tokenizer's own eos still terminates both, + # so default termination is preserved. + assert not _token_allowed(bm_default[0], LETTER) + assert _token_allowed(bm_override[0], LETTER) + assert _token_allowed(bm_default[0], EOS) + assert _token_allowed(bm_override[0], EOS) diff --git a/vllm/v1/structured_output/__init__.py b/vllm/v1/structured_output/__init__.py index 08bb237c04ff..f2ba4ac368e4 100644 --- a/vllm/v1/structured_output/__init__.py +++ b/vllm/v1/structured_output/__init__.py @@ -184,7 +184,14 @@ def _create_grammar(self, request: "Request") -> StructuredOutputGrammar: try: request_type, grammar_spec = struct_request.structured_output_key assert self.backend is not None - return self.backend.compile_grammar(request_type, grammar_spec) + stop_token_ids = ( + request.sampling_params.all_stop_token_ids + if request.sampling_params is not None + else None + ) + return self.backend.compile_grammar( + request_type, grammar_spec, stop_token_ids=stop_token_ids + ) except Exception: logger.exception( "Failed to compile grammar for request %s", request.request_id diff --git a/vllm/v1/structured_output/backend_guidance.py b/vllm/v1/structured_output/backend_guidance.py index 19b2c76e9d3d..30ecbfa065a6 100644 --- a/vllm/v1/structured_output/backend_guidance.py +++ b/vllm/v1/structured_output/backend_guidance.py @@ -106,7 +106,10 @@ def __post_init__(self): ) def compile_grammar( - self, request_type: StructuredOutputOptions, grammar_spec: str + self, + request_type: StructuredOutputOptions, + grammar_spec: str, + stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: self.serialized_grammar = serialize_guidance_grammar( request_type, diff --git a/vllm/v1/structured_output/backend_lm_format_enforcer.py b/vllm/v1/structured_output/backend_lm_format_enforcer.py index 69eded30fc08..898aaa136e40 100644 --- a/vllm/v1/structured_output/backend_lm_format_enforcer.py +++ b/vllm/v1/structured_output/backend_lm_format_enforcer.py @@ -99,7 +99,10 @@ def __post_init__(self): ) def compile_grammar( - self, request_type: StructuredOutputOptions, grammar_spec: str + self, + request_type: StructuredOutputOptions, + grammar_spec: str, + stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: character_level_parser: lmformatenforcer.CharacterLevelParser if request_type == StructuredOutputOptions.JSON: diff --git a/vllm/v1/structured_output/backend_outlines.py b/vllm/v1/structured_output/backend_outlines.py index 91627ff154c2..e66ef6361a70 100644 --- a/vllm/v1/structured_output/backend_outlines.py +++ b/vllm/v1/structured_output/backend_outlines.py @@ -71,7 +71,10 @@ def _compile_index( return index def compile_grammar( - self, request_type: StructuredOutputOptions, grammar_spec: str + self, + request_type: StructuredOutputOptions, + grammar_spec: str, + stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: if request_type == StructuredOutputOptions.JSON: regex = json_schema.build_regex_from_schema(grammar_spec) diff --git a/vllm/v1/structured_output/backend_types.py b/vllm/v1/structured_output/backend_types.py index 5c09b7b0634f..d9826c18f534 100644 --- a/vllm/v1/structured_output/backend_types.py +++ b/vllm/v1/structured_output/backend_types.py @@ -105,7 +105,10 @@ class StructuredOutputBackend(ABC): @abstractmethod def compile_grammar( - self, request_type: StructuredOutputOptions, grammar_spec: str + self, + request_type: StructuredOutputOptions, + grammar_spec: str, + stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: """ Compiles a grammar specification into a structured output grammar. @@ -114,6 +117,9 @@ def compile_grammar( request_type (StructuredOutputOptions): The type of structured output request. grammar_spec (str): The grammar specification to compile. + stop_token_ids (set[int] | None): The request's EOS and user + stop-token ids (``SamplingParams.all_stop_token_ids``), masked + until the grammar terminates. Returns: StructuredOutputGrammar: The compiled structured output grammar. diff --git a/vllm/v1/structured_output/backend_xgrammar.py b/vllm/v1/structured_output/backend_xgrammar.py index 4f199a1a2735..58b726bf72a4 100644 --- a/vllm/v1/structured_output/backend_xgrammar.py +++ b/vllm/v1/structured_output/backend_xgrammar.py @@ -76,7 +76,10 @@ def __post_init__(self): ) def compile_grammar( - self, request_type: StructuredOutputOptions, grammar_spec: str + self, + request_type: StructuredOutputOptions, + grammar_spec: str, + stop_token_ids: set[int] | None = None, ) -> StructuredOutputGrammar: if request_type == StructuredOutputOptions.JSON: ctx = self.compiler.compile_json_schema( @@ -119,6 +122,7 @@ def compile_grammar( return XgrammarGrammar( matcher=xgr.GrammarMatcher( ctx, + override_stop_tokens=list(stop_token_ids) if stop_token_ids else None, max_rollback_tokens=self.num_speculative_tokens, ), vocab_size=self.vocab_size, From 1482d2e015cc71e6ce0d3d92249cb78f4738f2e7 Mon Sep 17 00:00:00 2001 From: Chaitanya Sri Krishna Lolla Date: Tue, 11 Aug 2026 00:09:26 +0530 Subject: [PATCH 04/30] [ROCm][DistInf] Enable vLLM DI CI with buildkite/slurm (#47030) Signed-off-by: lcskrishna Co-authored-by: Sheral Kumar Co-authored-by: avininjamay8 Co-authored-by: tej <37236721+itej89@users.noreply.github.com> Co-authored-by: Cursor --- .buildkite/amd-disagg/cluster.sh | 166 +++++ .buildkite/amd-disagg/models.yaml | 97 +++ .buildkite/amd-disagg/pipeline-disagg.yaml | 202 ++++++ .../amd-disagg/run-slurm-disagg-test.sh | 338 ++++++++++ .buildkite/amd-disagg/run_xPyD_disagg.slurm | 536 +++++++++++++++ .buildkite/amd-disagg/vllm_disagg.sh | 635 ++++++++++++++++++ 6 files changed, 1974 insertions(+) create mode 100644 .buildkite/amd-disagg/cluster.sh create mode 100644 .buildkite/amd-disagg/models.yaml create mode 100644 .buildkite/amd-disagg/pipeline-disagg.yaml create mode 100644 .buildkite/amd-disagg/run-slurm-disagg-test.sh create mode 100644 .buildkite/amd-disagg/run_xPyD_disagg.slurm create mode 100644 .buildkite/amd-disagg/vllm_disagg.sh diff --git a/.buildkite/amd-disagg/cluster.sh b/.buildkite/amd-disagg/cluster.sh new file mode 100644 index 000000000000..4b377dca08c3 --- /dev/null +++ b/.buildkite/amd-disagg/cluster.sh @@ -0,0 +1,166 @@ +#!/bin/bash +# ============================================================================= +# cluster.sh — single, generic config for the vLLM disaggregated (P/D) launcher. +# ----------------------------------------------------------------------------- +# This script is *sourced* by vllm_disagg.sh. +# +# Every value is `${VAR:-default}`, so the environment always wins: +# environment variable > built-in default below +# +# So you can override anything inline: +# PREFILL_IP=10.0.0.1 DECODE_IP=10.0.0.2 ./vllm_disagg.sh prefill +# +# Site-specific values (model dir, IPs, NIC list, partition) are the defaults +# in the "site defaults" sections below — edit those for a new cluster. +# Model-SPECIFIC perf flags live in models.yaml, NOT here. +# ============================================================================= + +_CLUSTER_SH_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# ----------------------------------------------------------------- model / mode +# MODEL_NAME indexes into models.yaml; MODEL_DIR is the parent dir holding it. +# MODEL_PATH (resolved by the launcher) = ${MODEL_DIR}/${MODEL_NAME}. [site] +export MODEL_NAME="${MODEL_NAME:-DeepSeek-V3}" +export MODEL_DIR="${MODEL_DIR:-/data/models2}" + +# Shared NFS root (5 TB, visible on every node): model weights + per-run logs. [site] +export SHARED_MOUNT="${SHARED_MOUNT:-/data}" + +# Parallelism mode (launcher derives PARALLEL_MODE tp|ep from this): +# WIDE_EP_MODE=0 tp : each node is an independent TP server (TP8, 1P1D) +# WIDE_EP_MODE=1 ep : data-parallel + expert-parallel across xP/yD nodes (wideep) +export WIDE_EP_MODE="${WIDE_EP_MODE:-0}" + +# ----------------------------------------------------------------- topology +# xP prefill nodes + yD decode nodes. IPADDRS is the ordered, comma-separated +# node list (prefill IPs first, then decode IPs). NODE_RANK is this node's global +# 0-based rank; under SLURM it defaults to $SLURM_PROCID. Leave IPADDRS empty to +# use the PREFILL_IP/DECODE_IP fallback defaults below (1P1D only). +export xP="${xP:-1}" +export yD="${yD:-1}" +export IPADDRS="${IPADDRS:-}" +export NODE_RANK="${NODE_RANK:-${SLURM_PROCID:-}}" +export PREFILL_IP="${PREFILL_IP:-10.0.0.1}" +export DECODE_IP="${DECODE_IP:-10.0.0.2}" + +# Per-node GPU count and TP degree +export GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +export TP_SIZE="${TP_SIZE:-${GPUS_PER_NODE}}" + +# ----------------------------------------------------------------- ports +# TP-mode (WIDE_EP_MODE=0) server ports. +export PREFILL_PORT="${PREFILL_PORT:-8100}" +export DECODE_PORT="${DECODE_PORT:-8200}" + +# EP-mode (WIDE_EP_MODE=1) ports: API serve port, DP RPC port, KV transfer port, and +# per-node MoRIIO local ping port (must differ from PROXY_PING_PORT). +export SERVE_PORT="${SERVE_PORT:-20005}" +export RPC_PORT="${RPC_PORT:-13345}" +export KV_PORT="${KV_PORT:-9711}" +export LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}" + +# MoRIIO proxy: HTTP port clients/benchmark hit, plus the connector control ports. +# PROXY_PING_PORT MUST be 36367 — the proxy hardcodes its zmq service-discovery +# socket on that port; prefill/decode register to PROXY_IP:PROXY_PING_PORT. +export PROXY_IP="${PROXY_IP:-${PREFILL_IP}}" +export PROXY_PORT="${PROXY_PORT:-10001}" +export PROXY_PING_PORT="${PROXY_PING_PORT:-36367}" +export HANDSHAKE_PORT="${HANDSHAKE_PORT:-6301}" +export NOTIFY_PORT="${NOTIFY_PORT:-61005}" + +export PROXY_SCRIPT="${PROXY_SCRIPT:-/app/vllm/examples/disaggregated/disaggregated_serving/moriio_toy_proxy_server.py}" + +# MoRIIO KV transfer direction (injected into --kv-transfer-config by the launcher): +# 0 -> omit read_mode (default; MoRIIO write mode: prefill pushes to decode) +# 1 -> "read_mode": true (decode pulls KV from prefill; matches upstream disagg) +export MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" + +# ----------------------------------------------------------------- router / gateway +# Selection for client (bench/accuracy) traffic: +# proxy -> the in-container MoRIIO proxy started by the launcher +# vllm-router -> an external `vllm/vllm-router` container started by the SLURM job +# on the rank-0 node (default) +# Both use the SAME MoRIIO discovery mechanism (prefill/decode register to +# PROXY_IP:PROXY_PING_PORT=36367); only the client HTTP front door differs. +export ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}" +export ROUTER_PORT="${ROUTER_PORT:-30000}" +export ROUTER_POLICY="${ROUTER_POLICY:-round_robin}" +export VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}" +# Single client-facing port bench/accuracy target: the router port when routing, +# else the proxy port. Env override always wins. +if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then + export GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}" +else + export GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}" +fi + +# Where per-run logs / benchmark results are written. A $SLURM_JOB_ID subdir is +# appended so each CI run is self-scoped (falls back to 'local' off-SLURM). [site] +_LOG_BASE="${LOG_BASE:-/data/${USER:-$(id -un)}/disagg_logs}" +export LOG_PATH="${LOG_PATH:-${_LOG_BASE}/${SLURM_JOB_ID:-local}}" + +# ----------------------------------------------------------------- vLLM runtime +# Engine/platform/transport-level env (NOT model-specific). Model-architecture +# AITER kernel toggles live in models.yaml under each model's `env:` block. +export VLLM_USE_V1="${VLLM_USE_V1:-1}" +export HSA_NO_SCRATCH_RECLAIM="${HSA_NO_SCRATCH_RECLAIM:-1}" + +#export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}" +#export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}" +# +#export HOME=/tmp +export HF_HOME="${HF_HOME:-/tmp/hf_home}" +export XDG_CACHE_HOME="${XDG_CACHE_HOME:-/tmp/.cache}" +export VLLM_ENGINE_READY_TIMEOUT_S="${VLLM_ENGINE_READY_TIMEOUT_S:-3600}" + +# ----------------------------------------------------------------- scale +# DP/EP group formation timeout across nodes (seconds). +export DISTRIBUTED_TIMEOUT_SECONDS="${DISTRIBUTED_TIMEOUT_SECONDS:-7200}" + +# ----------------------------------------------------------------- RDMA / NCCL +# AMD Pensando AINIC RoCE fabric: 8 NICs exposed as ionic_0..7 (netdevs eth2..9), +# each rail on its own /24. GID index 1 + traffic class 104 +_IB_DEVICES="${IB_DEVICES:-ionic_0,ionic_1,ionic_2,ionic_3,ionic_4,ionic_5,ionic_6,ionic_7}" +_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-1}" +export IB_DEVICES="${IB_DEVICES:-${_IB_DEVICES}}" +export NCCL_IB_HCA="${NCCL_IB_HCA:-${_IB_DEVICES}}" +export NCCL_IB_GID_INDEX="${NCCL_IB_GID_INDEX:-${_IB_GID_INDEX}}" +export NCCL_IB_DISABLE="${NCCL_IB_DISABLE:-0}" +# In-box RCCL net transport (no external plugin); pin bootstrap to the VPC iface. +export NCCL_NET_PLUGIN="${NCCL_NET_PLUGIN:-none}" +export NCCL_SOCKET_IFNAME="${NCCL_SOCKET_IFNAME:-eth0}" +export GLOO_SOCKET_IFNAME="${GLOO_SOCKET_IFNAME:-${NCCL_SOCKET_IFNAME}}" +export NCCL_CROSS_NIC="${NCCL_CROSS_NIC:-0}" +export NCCL_PXN_DISABLE="${NCCL_PXN_DISABLE:-0}" +export NCCL_NET_DISABLE_INTRA="${NCCL_NET_DISABLE_INTRA:-1}" +export NCCL_IB_TC="${NCCL_IB_TC:-104}" +export NCCL_IB_FIFO_TC="${NCCL_IB_FIFO_TC:-192}" +export NCCL_IB_QPS_PER_CONNECTION="${NCCL_IB_QPS_PER_CONNECTION:-1}" +export NCCL_IB_TIMEOUT="${NCCL_IB_TIMEOUT:-22}" +export NCCL_IB_RETRY_CNT="${NCCL_IB_RETRY_CNT:-12}" + +# MoRI uses the same NIC set as NCCL. +export MORI_RDMA_DEVICES="${MORI_RDMA_DEVICES:-${_IB_DEVICES}}" +export MORI_IB_GID_INDEX="${MORI_IB_GID_INDEX:-${_IB_GID_INDEX}}" +export MORI_SHMEM_HEAP_SIZE="${MORI_SHMEM_HEAP_SIZE:-16G}" + +# Pin to gfx950 to avoid jit compilation failures with other archs on this cluster. +export MORI_GPU_ARCHS="gfx950" + +# ----------------------------------------------------------------- benchmark +export BENCHMARK_COMBINATIONS="${BENCHMARK_COMBINATIONS:-1024/128 2048/128}" +export BENCHMARK_CON="${BENCHMARK_CON:-32 64}" +export NUM_PROMPTS_FACTOR="${NUM_PROMPTS_FACTOR:-2}" +export BENCHMARK_MIN_PROMPTS="${BENCHMARK_MIN_PROMPTS:-32}" + +# ----------------------------------------------------------------- accuracy +# `accuracy` role runs lm_eval (local-completions backend) against the proxy. +export ACCURACY_TASKS="${ACCURACY_TASKS:-gsm8k}" +export ACCURACY_NUM_CONCURRENT="${ACCURACY_NUM_CONCURRENT:-64}" +export ACCURACY_MAX_RETRIES="${ACCURACY_MAX_RETRIES:-3}" +export ACCURACY_METRIC="${ACCURACY_METRIC:-exact_match}" +export ACCURACY_THRESHOLD="${ACCURACY_THRESHOLD:-0.90}" + +# ----------------------------------------------------------------- SLURM (submit) +# Used by run-slurm-disagg-test.sh on the login node (harmless to export here). [site] +export SLURM_PARTITION="${SLURM_PARTITION:-default}" diff --git a/.buildkite/amd-disagg/models.yaml b/.buildkite/amd-disagg/models.yaml new file mode 100644 index 000000000000..e919475db85e --- /dev/null +++ b/.buildkite/amd-disagg/models.yaml @@ -0,0 +1,97 @@ +# ============================================================================= +# Model catalog for the vLLM disaggregated (P/D) inference CI. +# ----------------------------------------------------------------------------- +# Loaded by vllm_disagg.sh for the selected MODEL_NAME. WIDE_EP_MODE selects the +# flag set (0 -> tp keys, 1 -> ep keys): +# tp -> base_flags + prefill.tp / decode.tp +# ep -> base_flags + prefill.ep / decode.ep +# +# Keep ONLY model-specific perf flags here. The launcher owns everything +# topological / mode-structural: +# - tp: --host/--port, --tensor-parallel-size, --kv-transfer-config +# - ep: --tp 1, --data-parallel-size[-local], --data-parallel-address/-rpc-port, +# --data-parallel-start-rank/--headless (children), --enable-expert-parallel, +# --all2all-backend mori, --no-enable-prefix-caching, +# --api-server-count + --kv-transfer-config (masters only) +# Do NOT put any of those in this file. +# +# Format: a top-level `models:` list. The launcher selects the entry whose +# `model:` matches MODEL_NAME. Each entry: +# model : name (must match MODEL_NAME; also the dir under MODEL_DIR) +# env : model/arch-specific environment variables. The launcher +# exports these BEFORE `vllm serve`, but only if they are +# not already set, so precedence is: +# caller/inline env > models.yaml env: +# (Transport/fabric/engine env stays in cluster.sh, +# e.g. MORI_RDMA_DEVICES, NCCL_IB_HCA, VLLM_USE_V1.) +# base_flags : always applied (both roles, both modes) +# prefill/decode. : role + mode specific perf flags +# experimental_flags : optional extra flags (omit if empty) +# ============================================================================= + +models: + - model: DeepSeek-V3 + env: + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + VLLM_ROCM_USE_AITER_MOE: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "0" + base_flags: "--trust-remote-code --kv-cache-dtype fp8" + prefill: + tp: "--gpu-memory-utilization 0.85" + ep: "--gpu-memory-utilization 0.85 --enforce-eager" + decode: + tp: "--gpu-memory-utilization 0.85" + #ep: '--gpu-memory-utilization 0.75 --enforce-eager' + ep: '--gpu-memory-utilization 0.75 --compilation-config {"cudagraph_mode":"PIECEWISE","custom_ops":["+quant_fp8"]}' + + - model: MiniMax-M3-MXFP8 + env: + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_MOE: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + VLLM_USE_BREAKABLE_CUDAGRAPH: "0" + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT6" + base_flags: "--trust-remote-code --attention-backend TRITON_ATTN --block-size 128 --language-model-only --kv-cache-dtype fp8" + prefill: + tp: "--gpu-memory-utilization 0.85 --enforce-eager" + decode: + tp: "--gpu-memory-utilization 0.85" + + - model: DeepSeek-R1-MXFP4 + env: + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + VLLM_ROCM_USE_AITER_MOE: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + base_flags: "--trust-remote-code --kv-cache-dtype fp8" + prefill: + tp: "--gpu-memory-utilization 0.85" + decode: + tp: "--gpu-memory-utilization 0.85" + + - model: Kimi-K2.5-MXFP4 + env: + VLLM_ROCM_USE_AITER: "1" + VLLM_ROCM_USE_AITER_MOE: "1" + VLLM_ROCM_USE_AITER_RMSNORM: "1" + base_flags: "--trust-remote-code" + prefill: + tp: "--gpu-memory-utilization 0.85" + decode: + tp: "--gpu-memory-utilization 0.85" + + - model: Kimi-K2.6-MXFP4 + env: + VLLM_ROCM_USE_AITER: "1" + AMDGCN_USE_BUFFER_OPS: "1" + VLLM_ROCM_USE_AITER_MLA: "1" + VLLM_ROCM_QUICK_REDUCE_QUANTIZATION: "INT4" + VLLM_ROCM_USE_SKINNY_GEMM: "0" + VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS: "1" + base_flags: "--trust-remote-code --kv-cache-dtype fp8 --mm-encoder-tp-mode data --block-size 1 --attention-backend ROCM_AITER_MLA" + prefill: + tp: "--gpu-memory-utilization 0.9" + decode: + tp: "--gpu-memory-utilization 0.9" diff --git a/.buildkite/amd-disagg/pipeline-disagg.yaml b/.buildkite/amd-disagg/pipeline-disagg.yaml new file mode 100644 index 000000000000..cf8fe0d23023 --- /dev/null +++ b/.buildkite/amd-disagg/pipeline-disagg.yaml @@ -0,0 +1,202 @@ +# Native Buildkite pipeline — disaggregated PD SLURM tests (MoRIIO). +# +# ROUTER_TYPE must be set explicitly on every step: run-slurm-disagg-test.sh +# defaults to `vllm-router` +# + +steps: + +#-------------------------------------- mi350 · disagg 1P1D (TP8, 2 nodes) · SLURM --------------------------------------# + +- label: "DeepSeek-V3-PD-1P1D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-1P1D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-R1-MXFP4-PD-1P1D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-R1-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-R1-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=DeepSeek-R1-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.5-MXFP4-PD-1P1D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=Kimi-K2.5-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.5-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=Kimi-K2.5-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.6-MXFP4-PD-1P1D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=Kimi-K2.6-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.6-MXFP4-PD-1P1D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=Kimi-K2.6-MXFP4 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "MiniMax-M3-MXFP8-PD-1P1D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=MiniMax-M3-MXFP8 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "MiniMax-M3-MXFP8-PD-1P1D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/2node" + commands: + - MODEL_NAME=MiniMax-M3-MXFP8 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +#-------------------------------------- mi350 · disagg 2P2D (TP8, 4 nodes) · SLURM --------------------------------------# + +- label: "DeepSeek-V3-PD-2P2D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-V3-PD-2P2D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-V3 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-R1-MXFP4-PD-2P2D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-R1-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "DeepSeek-R1-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=DeepSeek-R1-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.5-MXFP4-PD-2P2D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=Kimi-K2.5-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.5-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=Kimi-K2.5-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.6-MXFP4-PD-2P2D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=Kimi-K2.6-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "Kimi-K2.6-MXFP4-PD-2P2D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=Kimi-K2.6-MXFP4 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "MiniMax-M3-MXFP8-PD-2P2D-TP8-MoRIIO-proxy" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=MiniMax-M3-MXFP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +- label: "MiniMax-M3-MXFP8-PD-2P2D-TP8-MoRIIO-vllm-router" + agents: + queue: amd_mi350_ainic + timeout_in_minutes: 120 + concurrency: 2 + concurrency_group: "amd-disagg/mi350/4node" + commands: + - MODEL_NAME=MiniMax-M3-MXFP8 xP=2 yD=2 NODES=4 GPUS_PER_NODE=8 WIDE_EP_MODE=0 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=vllm-router WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh + +#-------------------------------------- mi350 · disagg wide-EP (disabled) --------------------------------------# + +# - label: "DeepSeek-V3-PD-1P1D-EP8/DP8-WideEP-MoRIIO-proxy" +# agents: +# queue: amd_mi350_ainic +# timeout_in_minutes: 120 +# concurrency: 2 +# concurrency_group: "amd-disagg/mi350/2node" +# commands: +# - IMAGE=vllm/vllm-openai-rocm:nightly MODEL_NAME=DeepSeek-V3 NODES=2 GPUS_PER_NODE=8 WIDE_EP_MODE=1 MORIIO_READ_MODE=0 RUN_AFTER_HEALTH=accuracy ROUTER_TYPE=proxy WAIT=1 SLURM_TIME_LIMIT=02:00:00 bash .buildkite/amd-disagg/run-slurm-disagg-test.sh diff --git a/.buildkite/amd-disagg/run-slurm-disagg-test.sh b/.buildkite/amd-disagg/run-slurm-disagg-test.sh new file mode 100644 index 000000000000..e4beeb577ff2 --- /dev/null +++ b/.buildkite/amd-disagg/run-slurm-disagg-test.sh @@ -0,0 +1,338 @@ +#!/bin/bash +# ============================================================================= +# run-slurm-disagg-test.sh — foreground (login-node) submitter for the disagg +# P/D gate. The single command the Buildkite step runs. +# ----------------------------------------------------------------------------- +# Thin glue: submits run_xPyD_disagg.slurm. By default (WAIT=0, Spur-safe) it +# fire-and-forgets and returns 0 on a good submit; with WAIT=1 it polls the +# scheduler and exits with the job's exit code so a CI step pass/fails. +# stdout carries ONLY the SLURM job info +# (image/nodes line, job id, log path); after sbatch nothing else is written to +# stdout — the full job log persists at the --output path and pass/fail +# diagnostics go to stderr. The actual work lives in run_xPyD_disagg.slurm: it +# selects nodes, +# fans out one container per node via a single srun, and hands off to +# vllm_disagg.sh (rank-based prefill/decode self-select). +# +# Default target: 1P1D TP8 (NODES=2), nightly image, accuracy gate. +# +# Spur usage (fire-and-forget; the default here): +# bash run-slurm-disagg-test.sh # 1P1D TP8 +# WIDE_EP_MODE=1 bash run-slurm-disagg-test.sh # 1P1D wide-EP +# NODES=4 WIDE_EP_MODE=1 xP=2 yD=2 bash run-slurm-disagg-test.sh # 2P2D EP +# Prints the job id + NFS log paths and returns immediately (does NOT poll +# squeue/sacct, which hang on Spur). Track the run via the printed `tail -f`. +# +# Classic-Slurm / CI usage (block until done, exit with the gate's pass/fail): +# WAIT=1 bash run-slurm-disagg-test.sh +# ============================================================================= +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +JOB_SCRIPT="${JOB_SCRIPT:-${SCRIPT_DIR}/run_xPyD_disagg.slurm}" + +# ---- knobs (override from the Buildkite step env) -------------------------- +# Defaults tuned for the Spur AMD MI350X cluster +IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}" +NODES="${NODES:-2}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" +PARTITION="${SLURM_PARTITION:-}" +TIME_LIMIT="${SLURM_TIME_LIMIT:-02:00:00}" +WIDE_EP_MODE="${WIDE_EP_MODE:-0}" # 0 -> 1P1D TP8 (default); 1 -> wide-EP +xP="${xP:-1}" +yD="${yD:-1}" +RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" +HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}" # P/D bring-up budget; +900s grace must fit the 2h wall +SHARED_MOUNT="${SHARED_MOUNT:-/data}" +LOG_ROOT="${LOG_ROOT:-${SHARED_MOUNT}/${USER:-$(whoami)}/disagg_logs}" +DRY_RUN="${DRY_RUN:-0}" +MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" + +WAIT="${WAIT:-0}" + +# ROUTER Type - defaults to vllm-router +ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}" +ROUTER_PORT="${ROUTER_PORT:-30000}" +VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}" +# Dry-run only validates wiring; cap its walltime low so it never holds the queue. +[[ "${DRY_RUN}" == "1" ]] && TIME_LIMIT="${SLURM_TIME_LIMIT:-00:10:00}" + +mkdir -p "${LOG_ROOT}" + +# Spur - sbatch scheduler. +DISAGG_SCRIPTS_STAGE="${DISAGG_SCRIPTS_STAGE:-/data/scratch/buildkite-agent}" +STAGED_DIR="${DISAGG_SCRIPTS_STAGE}/${BUILDKITE_COMMIT:-local}" +mkdir -p "${STAGED_DIR}" +cp -rL --no-preserve=ownership,timestamps "${SCRIPT_DIR}/." "${STAGED_DIR}/" +chmod -R u+rwX "${STAGED_DIR}" 2>/dev/null || true +export DISAGG_SCRIPTS_DIR="${STAGED_DIR}" +echo "[slurm-submit] staged scripts for compute nodes: ${DISAGG_SCRIPTS_DIR}" >&2 +export IMAGE MODEL_NAME WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S +export SHARED_MOUNT LOG_ROOT DRY_RUN MORIIO_READ_MODE +export ROUTER_TYPE ROUTER_PORT VLLM_ROUTER_IMAGE + +# Model selection. +[[ -n "${MODEL_NAME:-}" ]] && export MODEL_NAME +[[ -n "${MODEL_DIR:-}" ]] && export MODEL_DIR + +[[ -n "${PARTITION}" ]] && echo "[slurm-submit] NOTE: PARTITION='${PARTITION}' ignored (Spur sbatch has no --partition)" >&2 + +SUBMIT_SCRIPT="${JOB_SCRIPT}" +FILE_NODES="$(grep -oE '^#SBATCH[[:space:]]+--nodes=[0-9]+' "${JOB_SCRIPT}" | grep -oE '[0-9]+' | head -n1)" +if [[ -n "${NODES}" && -n "${FILE_NODES}" && "${NODES}" != "${FILE_NODES}" ]]; then + SUBMIT_SCRIPT="/tmp/$(basename "${JOB_SCRIPT%.slurm}")-n${NODES}-$$.slurm" + sed -E "s/^#SBATCH([[:space:]]+)--nodes=[0-9]+/#SBATCH\\1--nodes=${NODES}/" "${JOB_SCRIPT}" > "${SUBMIT_SCRIPT}" + echo "[slurm-submit] node count ${FILE_NODES}->${NODES}: submitting patched copy ${SUBMIT_SCRIPT}" >&2 +fi + +# Job name comes from the script's #SBATCH --job-name (SLURM sets SLURM_JOB_NAME +JOB_NAME="$(grep -oE '^#SBATCH[[:space:]]+--job-name=[^[:space:]]+' "${JOB_SCRIPT}" | sed -E 's/.*--job-name=//' | head -n1)" +JOB_NAME="${JOB_NAME:-vllm-disagg-pd}" + +echo "[slurm-submit] image=${IMAGE} nodes=${NODES} gpus/node=${GPUS_PER_NODE} mode=$([[ ${WIDE_EP_MODE} == 0 ]] && echo tp || echo ep) router=${ROUTER_TYPE} walltime=${TIME_LIMIT}" +SUBMIT_OUT="$(sbatch --time="${TIME_LIMIT}" "${SUBMIT_SCRIPT}")" +echo "${SUBMIT_OUT}" +# "Submitted batch job 114" -> 114 (last integer on the line). +JOB_ID="$(printf '%s\n' "${SUBMIT_OUT}" | grep -oE '[0-9]+' | tail -n1 || true)" +if [[ -z "${JOB_ID}" ]]; then + echo "[slurm-submit] ERROR: could not parse a job id from sbatch output above" >&2 + exit 1 +fi +echo "[slurm-submit] submitted job ${JOB_ID}" + +# The real per-job log is written to NFS from inside the job body by +# run_xPyD_disagg.slurm (`exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log"`). +# SLURM_JOB_NAME resolves to our --job-name, so we can compute the path here. The +# per-role server/proxy/bench logs land under ${LOG_ROOT}/${JOB_ID}/. +LOG_FILE="${LOG_ROOT}/${JOB_NAME}-${JOB_ID}.log" +LOG_DIR="${LOG_ROOT}/${JOB_ID}" +echo "[slurm-submit] job log: ${LOG_FILE}" +echo "[slurm-submit] role logs: ${LOG_DIR}/" + +# --- Spur-safe by default: fire-and-forget -------------------------------------- +if [[ "${WAIT}" != "1" ]]; then + echo "[slurm-submit] submitted (WAIT=0, not polling). Track with:" >&2 + echo " tail -f ${LOG_FILE}" >&2 + echo " grep -aE '(PASS|FAIL): |exact_match' ${LOG_FILE}" >&2 + exit 0 +fi + +# --- WAIT=1: phase-aware poll of the NFS job log (scontrol-based) ------------- +# Detection is phased so failures surface fast instead of waiting out the full +# walltime: +# 1) submitted -> running : scontrol catches infra/scheduler failures +# (NODE_FAIL/BOOT_FAIL/CANCELLED/TIMEOUT/...) within +# one poll, or a stuck-PENDING/never-started job. +# 2) running -> healthy : advance once every endpoint reports healthy (+ the +# vllm-router, when used); fail on bring-up errors or +# the health budget. +# 3) healthy -> verdict : PASS/FAIL from the accuracy gate, capped. +# scontrol is the scheduler authority on this cluster +echo "[slurm-submit] WAIT=1: phase-aware poll of ${LOG_FILE} (timeout ${TIME_LIMIT})" >&2 + +_h=0 _m=0 _s=0 +IFS=: read -r _h _m _s <<< "${TIME_LIMIT}" +WAIT_DEADLINE=$(( $(date +%s) + 10#${_h}*3600 + 10#${_m}*60 + 10#${_s:-0} )) +unset -v _h _m _s + +# Per-phase budgets (all overridable from the Buildkite step env). +POLL_INTERVAL="${POLL_INTERVAL:-20}" +SUBMIT_GRACE_S="${SUBMIT_GRACE_S:-900}" # reach RUNNING within 15m +PENDING_MAX_S="${PENDING_MAX_S:-1800}" # tolerate 30m queued +HEALTH_PHASE_TIMEOUT_S="${HEALTH_PHASE_TIMEOUT_S:-$(( HEALTH_TIMEOUT_S + 900 ))}" +WORKLOAD_TIMEOUT_S="${WORKLOAD_TIMEOUT_S:-1800}" # accuracy/bench cap +COMPLETED_GRACE_S="${COMPLETED_GRACE_S:-60}" # verdict lag after job exit + +SENTINEL="${LOG_DIR}/.disagg_done" + +# scontrol field extractor (authoritative here; timeout-guarded so a momentary +# scheduler stall can't wedge the poll). Returns the value or "". +job_field() { # $1=jobid $2=field -> value | "" + timeout 15 scontrol show job "$1" 2>/dev/null \ + | grep -oE "$2=[^ ]+" | head -n1 | cut -d= -f2- || true +} +have() { grep -aqE "$1" "${LOG_FILE}" 2>/dev/null; } + +# The job's own verdict sources, in authority order: the .disagg_done sentinel +# (rank-0 rc, written by a single writer) and then the gate line in the log. +# Sets STATE/RC/REASON and returns 0 when a verdict exists, 1 when it does not. +read_verdict() { + if [[ -f "${SENTINEL}" ]]; then + RC="$(tr -dc '0-9' < "${SENTINEL}" 2>/dev/null || true)"; RC="${RC:-1}" + if [[ "${RC}" == "0" ]]; then STATE="COMPLETED"; else STATE="FAILED"; fi + REASON="sentinel"; return 0 + fi + if have '(PASS|FAIL): '; then + if have 'FAIL: '; then STATE="FAILED"; RC=1; else STATE="COMPLETED"; RC=0; fi + REASON="gate"; return 0 + fi + return 1 +} + +# Never let the job outlive this poller. Cancelling a Buildkite build kills the +# agent's bootstrap, and without this the sbatch job keeps its whole allocation +# until the walltime expires. Only armed on the WAIT=1 path — under WAIT=0, +# leaving the job running is the point. +CANCEL_GRACE_S="${CANCEL_GRACE_S:-120}" +_CLEANED=0 +# shellcheck disable=SC2329 # invoked from cleanup_job, which the traps below call +job_active() { + case "$(job_field "${JOB_ID}" JobState)" in + RUNNING|PENDING|COMPLETING|CONFIGURING|SUSPENDED|REQUEUED) return 0 ;; + *) return 1 ;; + esac +} +# shellcheck disable=SC2329 # invoked from the traps below +cleanup_job() { + [[ "${_CLEANED}" == "1" || -z "${JOB_ID:-}" ]] && return 0 + _CLEANED=1 + # A job that reported its own verdict is mid-teardown; give it a bounded + # moment to finish. Timeouts and infra failures get cancelled immediately — + # there is nothing to wait for and the nodes should come back now. + if [[ "${REASON:-}" == "sentinel" || "${REASON:-}" == "gate" || "${STATE:-}" == "COMPLETED" ]]; then + local deadline=$(( $(date +%s) + CANCEL_GRACE_S )) + while (( $(date +%s) < deadline )) && job_active; do sleep 5; done + fi + if job_active; then + echo "[slurm-submit] cleanup: cancelling job ${JOB_ID}" >&2 + scancel "${JOB_ID}" 2>/dev/null \ + || echo "[slurm-submit] WARN: scancel ${JOB_ID} failed; job may still hold nodes" >&2 + fi + return 0 +} +trap cleanup_job EXIT +trap 'cleanup_job; exit 130' INT +trap 'cleanup_job; exit 143' TERM +trap 'cleanup_job; exit 129' HUP + +STATE="" +RC=1 +REASON="" +PHASE="submitted" +T_PHASE=$(date +%s) + +while [[ $(date +%s) -lt ${WAIT_DEADLINE} ]]; do + NOW=$(date +%s) + + # (1) Ultimate authority: terminal sentinel (holds rank-0 rc), then the + # explicit accuracy gate line. Honored regardless of phase. + read_verdict && break + + # (2) Scheduler state via scontrol: drives phase transitions and catches + # infra/scheduler failures fast. + ST="$(job_field "${JOB_ID}" JobState)" + case "${ST}" in + RUNNING|COMPLETING) + if [[ "${PHASE}" == "submitted" ]]; then PHASE="bringup"; T_PHASE="${NOW}"; fi + ;; + NODE_FAIL|BOOT_FAIL|CANCELLED|TIMEOUT|OUT_OF_MEMORY|DEADLINE|PREEMPTED) + STATE="infra-${ST}"; RC=1 + REASON="scontrol Reason=$(job_field "${JOB_ID}" Reason)"; break + ;; + FAILED) + # Ambiguous (infra crash vs a legit accuracy exit=1). A gate/sentinel + # line would have won above, so classify by the phase we reached. + case "${PHASE}" in + submitted) STATE="infra-FAILED" ;; + bringup) STATE="server-failed" ;; + *) STATE="workload-failed" ;; + esac + RC=1; REASON="scontrol JobState=FAILED phase=${PHASE}"; break + ;; + COMPLETED) + _vd=$(( NOW + COMPLETED_GRACE_S )) + _found=0 + while :; do + if read_verdict; then _found=1; break; fi + (( $(date +%s) >= _vd )) && break + sleep 5 + done + if (( _found == 1 )); then + echo "[slurm-submit] verdict arrived after JobState=COMPLETED (${REASON})" >&2 + else + STATE="completed-no-verdict"; RC=1 + REASON="COMPLETED with no sentinel/gate within ${COMPLETED_GRACE_S}s (a walltime kill reports COMPLETED here)" + fi + break + ;; + PENDING|CONFIGURING|RESV_DEL_HOLD|REQUEUED) + if (( NOW - T_PHASE > PENDING_MAX_S )); then + STATE="infra-stuck-${ST}"; RC=1; REASON="queued > ${PENDING_MAX_S}s"; break + fi + ;; + "") + # scontrol lost the job (purged past MinJobAge) with no log verdict: + # rely on the terminal markers above + the per-phase deadlines below. + : + ;; + esac + + # (3) Log-driven phase progress + per-phase deadlines. Fast-fails hangs where + # the job is still RUNNING (so scontrol won't help) but bring-up/workload + # is stuck. + case "${PHASE}" in + submitted) + # Job body appearing in the log is a second "it started" signal for + # when scontrol is briefly empty. + if have 'Selected node IPs|health-gate:|\[disagg-pd\]'; then + PHASE="bringup"; T_PHASE="${NOW}" + elif [[ -z "${ST}" ]] && (( NOW - T_PHASE > SUBMIT_GRACE_S )); then + # scontrol can't confirm the job exists (empty state) AND no log + # output within the grace window -> treat as a lost/failed launch. + # NB: a genuinely-queued job reports PENDING via scontrol above and + # is governed by PENDING_MAX_S (default 30m), not this grace window. + STATE="infra-nostart"; RC=1; REASON="no scheduler state or log within ${SUBMIT_GRACE_S}s"; break + fi + ;; + bringup) + if have 'FAIL:|TIMEOUT waiting for|exited while waiting|Traceback \(most recent'; then + STATE="server-failed"; RC=1; REASON="bring-up failure in log"; break + fi + NEED="$(grep -aoE 'waiting on [0-9]+' "${LOG_FILE}" 2>/dev/null | grep -oE '[0-9]+' | tail -n1 || true)" + GOT="$(grep -acE '\] healthy: ' "${LOG_FILE}" 2>/dev/null || true)"; GOT="${GOT:-0}" + ROUTER_OK=1 + if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then + if ! have 'vllm-router healthy on'; then ROUTER_OK=0; fi + fi + if [[ -n "${NEED}" ]] && (( GOT >= NEED )) && (( ROUTER_OK == 1 )); then + PHASE="workload"; T_PHASE="${NOW}" + elif (( NOW - T_PHASE > HEALTH_PHASE_TIMEOUT_S )); then + STATE="bringup-timeout"; RC=1; REASON="no healthy within ${HEALTH_PHASE_TIMEOUT_S}s"; break + fi + ;; + workload) + # PASS/FAIL handled at the top; only enforce the cap here. + if (( NOW - T_PHASE > WORKLOAD_TIMEOUT_S )); then + STATE="workload-timeout"; RC=1; REASON="no verdict within ${WORKLOAD_TIMEOUT_S}s"; break + fi + ;; + esac + + sleep "${POLL_INTERVAL}" +done + +if [[ -z "${STATE}" ]]; then + echo "[slurm-submit] WARN: no verdict before ${TIME_LIMIT}; failing" >&2 + STATE="deadline"; RC=1; REASON="poll deadline" +fi + +# Pre-flight logic. +PF_LINES=() +if compgen -G "${LOG_DIR}/preflight_NODE*.log" >/dev/null 2>&1; then + cat -- "${LOG_DIR}"/preflight_NODE*.log >&2 || true + mapfile -t PF_LINES < <(grep -h '^PREFLIGHT-REJECTED: ' "${LOG_DIR}"/preflight_NODE*.log 2>/dev/null || true) +fi +if (( ${#PF_LINES[@]} > 0 )) && (( RC != 0 )); then + STATE="preflight-rejected" + REASON="${PF_LINES[0]#PREFLIGHT-REJECTED: }" + if (( ${#PF_LINES[@]} > 1 )); then + REASON="${REASON} [+$(( ${#PF_LINES[@]} - 1 )) more node(s), see ${LOG_DIR}/]" + fi +fi + +# Surface the accuracy gate verdict (if any) from the job log — to stderr. +GATE_LINE=$(grep -aE '(PASS|FAIL): ' "${LOG_FILE}" 2>/dev/null | tail -n1 || true) +[[ -n "${GATE_LINE}" ]] && echo "[slurm-submit] gate: ${GATE_LINE}" >&2 + +echo "[slurm-submit] job ${JOB_ID} finished: state=${STATE:-unknown} phase=${PHASE} exit=${RC} reason=${REASON:-}" >&2 +exit "${RC}" diff --git a/.buildkite/amd-disagg/run_xPyD_disagg.slurm b/.buildkite/amd-disagg/run_xPyD_disagg.slurm new file mode 100644 index 000000000000..de50b40f52cd --- /dev/null +++ b/.buildkite/amd-disagg/run_xPyD_disagg.slurm @@ -0,0 +1,536 @@ +#!/bin/bash +#SBATCH --job-name=vllm-disagg-pd +#SBATCH --nodes=2 +#SBATCH --ntasks-per-node=1 +#SBATCH --gres=gpu:8 +#SBATCH --exclusive +#SBATCH --time=02:00:00 +#SBATCH --chdir=/tmp +#SBATCH --output=/tmp/spur-%j.out +#SBATCH --error=/tmp/spur-%j.err + +set -euo pipefail + + +# ---- inputs (env > default; forwarded by run-slurm-disagg-test.sh via --export) +IMAGE="${IMAGE:-vllm/vllm-openai-rocm:nightly}" +WIDE_EP_MODE="${WIDE_EP_MODE:-0}" +xP="${xP:-1}" +yD="${yD:-1}" +GPUS_PER_NODE="${GPUS_PER_NODE:-8}" + +SHARED_MOUNT="${SHARED_MOUNT:-/data}" +LOG_ROOT="${LOG_ROOT:-/data/${USER:-$(id -un)}/disagg_logs}" +RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy +DRY_RUN="${DRY_RUN:-0}" +MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" # 0 -> prefill writes (default); 1 -> decode reads KV (read_mode) + +# ---- router / gateway (host-side; also forwarded into the container) -------- +# vllm-router -> external `vllm/vllm-router` container started on the rank-0 node (default) +# proxy -> in-container MoRIIO proxy (no extra container) +ROUTER_TYPE="${ROUTER_TYPE:-vllm-router}" +ROUTER_PORT="${ROUTER_PORT:-30000}" +ROUTER_POLICY="${ROUTER_POLICY:-round_robin}" +VLLM_ROUTER_IMAGE="${VLLM_ROUTER_IMAGE:-vllm/vllm-router:nightly}" +PROXY_PING_PORT="${PROXY_PING_PORT:-36367}" # must equal vllm-router --vllm-discovery-address port +PROXY_PORT="${PROXY_PORT:-10001}" +# Intra-node data-parallel size the router uses to route the MoRIIO KV-notify to +# the right DP rank. Required for wideEP DP (else decode hangs: "remote blocks +# never arrived"). EP -> GPUS_PER_NODE (matches --data-parallel-size-local); TP -> 1. +if [[ "${WIDE_EP_MODE}" == "1" ]]; then + ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-${GPUS_PER_NODE}}" +else + ROUTER_DP_LOCAL="${ROUTER_DP_LOCAL:-1}" +fi +if [[ "${ROUTER_TYPE}" == "vllm-router" ]]; then + GATEWAY_PORT="${GATEWAY_PORT:-${ROUTER_PORT}}" +else + GATEWAY_PORT="${GATEWAY_PORT:-${PROXY_PORT}}" +fi + +# This line is a recommended template from infra side. +exec > "$LOG_ROOT/${SLURM_JOB_NAME}-${SLURM_JOB_ID}.log" 2>&1 + +# Where the disagg scripts live ON THE HOST (bind-mounted into the container). +# Spur runs the job body from /tmp, so $(pwd)/SLURM_SUBMIT_DIR aren't reliable; +# default to the known host scripts dir. Override via env if you relocate it. +DISAGG_SCRIPTS_DIR="${DISAGG_SCRIPTS_DIR:-${SLURM_SUBMIT_DIR:-$(pwd)}}" +echo "DISAGG SCRIPTS DIR is ----" +echo "$DISAGG_SCRIPTS_DIR" + +# In-container path the scripts are mounted to (and sourced from). +CONTAINER_SCRIPTS="${CONTAINER_SCRIPTS:-/vllm-workspace/.buildkite/amd-disagg}" +CLUSTER_ENV_IN="${CLUSTER_ENV_IN:-${CONTAINER_SCRIPTS}/cluster.sh}" + + +# Health endpoints (TP-mode ports from cluster.sh defaults). +PREFILL_PORT="${PREFILL_PORT:-8100}" +DECODE_PORT="${DECODE_PORT:-8200}" +HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}" + +# EP-mode ports, mirrored from cluster.sh purely so the preflight below can test +# them host-side; the in-container values still come from cluster.sh. +SERVE_PORT="${SERVE_PORT:-20005}" +RPC_PORT="${RPC_PORT:-13345}" +KV_PORT="${KV_PORT:-9711}" +LOCAL_PING_PORT="${LOCAL_PING_PORT:-61555}" +# vllm-router binds this for Prometheus in addition to ROUTER_PORT, and panics +# at startup if it is taken. +ROUTER_METRICS_PORT="${ROUTER_METRICS_PORT:-29000}" + +# Node preflight (see preflight_node): 0 disables, VRAM% above which a GPU counts +# as busy, and how long to let a reaped container's VRAM drain before judging. +PREFLIGHT="${PREFLIGHT:-1}" +PREFLIGHT_MAX_VRAM_PCT="${PREFLIGHT_MAX_VRAM_PCT:-10}" +PREFLIGHT_GPU_SETTLE_S="${PREFLIGHT_GPU_SETTLE_S:-30}" + + +NUM_NODES=$((xP + yD)) +echo "Calculated NUM_NODES: $NUM_NODES (xP=$xP + yD=$yD) WIDE_EP_MODE=$WIDE_EP_MODE" + +# ------------------------ +# Extract NUM_NODES from SLURM allocation and select nodes +# ------------------------ +echo "Original SLURM allocation:" +echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" +echo "SLURM_NNODES: $SLURM_NNODES" +echo "SLURM_NTASKS: $SLURM_NTASKS" + +# scontrol is available on classic Slurm but not in the Spur job body. Use it to +# expand compressed ranges when present, else split the (already comma-separated, +# full-hostname) SLURM_JOB_NODELIST on commas. +if command -v scontrol >/dev/null 2>&1; then + FULL_NODELIST=$(scontrol show hostnames "$SLURM_JOB_NODELIST") +else + FULL_NODELIST=$(echo "$SLURM_JOB_NODELIST" | tr ',' '\n') +fi +AVAILABLE_COUNT=$(echo "$FULL_NODELIST" | wc -l) + +if (( AVAILABLE_COUNT < NUM_NODES )); then + echo "ERROR: Need ${NUM_NODES} nodes (xP=${xP} + yD=${yD}) but only ${AVAILABLE_COUNT} allocated." >&2 + exit 1 +fi + +# Pick the first NUM_NODES nodes and order them to match how IPs get assigned: +# * classic Slurm: sort alphabetically -- srun assigns PROCID=0 to the +# alphabetically-first hostname, so IPS[] (built per node below) lines up. +# * Spur: preserve RANK order (SPUR_NODELIST), so the hostnames pair 1:1 with +# the rank-ordered IPs in SPUR_PEER_NODES; an alphabetical sort here would +# desync MASTER_NODE / MASTER_ADDR / IPADDRS from the actual ranks. +# IPADDRS[0..xP-1] = prefill, IPADDRS[xP..] = decode, so this order must match. +if command -v srun >/dev/null 2>&1; then + SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES" | sort) +elif [[ -n "${SPUR_NODELIST:-}" ]]; then + SELECTED_NODES=$(echo "${SPUR_NODELIST}" | tr ',' '\n' | sed '/^$/d' | head -n "$NUM_NODES") +else + SELECTED_NODES=$(echo "$FULL_NODELIST" | head -n "$NUM_NODES") +fi +SELECTED_NODELIST_STR=$(echo "$SELECTED_NODES" | tr '\n' ',' | sed 's/,$//') + +# Create new nodelist in SLURM format +# This is a simplified approach - for complex ranges, you might need more sophisticated parsing +NEW_SLURM_NODELIST=$(echo "$SELECTED_NODES" | paste -sd, | sed 's/,/,/g') + +# Update SLURM environment variables +export SLURM_NNODES=$NUM_NODES +export SLURM_NTASKS=$NUM_NODES +export SLURM_JOB_NUM_NODES=$NUM_NODES +export SLURM_NPROCS=$NUM_NODES +export SLURM_JOB_NODELIST="$NEW_SLURM_NODELIST" +export SLURM_NODELIST="$NEW_SLURM_NODELIST" + +# Keep other SLURM variables as they were or set defaults +export SLURM_TASKS_PER_NODE="1(x$NUM_NODES)" +export SLURM_SUBMIT_DIR="${SLURM_SUBMIT_DIR:-${HOME}}" +export SLURM_CLUSTER_NAME="${SLURM_CLUSTER_NAME:-vllm-di-ci}" +export SLURM_JOB_CPUS_PER_NODE="${SLURM_JOB_CPUS_PER_NODE:-1}" +export SLURM_JOB_PARTITION="${SLURM_JOB_PARTITION:-amd-rccl}" +export SLURM_JOBID="${SLURM_JOBID:-$SLURM_JOB_ID}" +export SLURM_JOB_QOS="${SLURM_JOB_QOS:-normal}" +export SLURM_JOB_ACCOUNT="${SLURM_JOB_ACCOUNT:-amd-rccl}" +export SLURM_NTASKS_PER_NODE=1 +export SLURM_SUBMIT_HOST="${SLURM_SUBMIT_HOST:-$(hostname)}" +export SLURM_JOB_ID="${SLURM_JOB_ID:-}" +export SLURM_CONF="${SLURM_CONF:-/etc/slurm/slurm.conf}" +export SLURM_JOB_NAME="${SLURM_JOB_NAME:-${xP}p${yD}d_bench-serving}" + +echo "" +echo "Updated SLURM Environment Variables:" +echo "SLURM_JOB_ID: $SLURM_JOB_ID" +echo "SLURM_JOB_NODELIST: $SLURM_JOB_NODELIST" +echo "SLURM_NNODES: $SLURM_NNODES" +echo "SLURM_NTASKS: $SLURM_NTASKS" +echo "SLURM_TASKS_PER_NODE: $SLURM_TASKS_PER_NODE" +echo "SLURM_JOB_CPUS_PER_NODE: $SLURM_JOB_CPUS_PER_NODE" +echo "SLURM_JOB_PARTITION: $SLURM_JOB_PARTITION" +echo "SLURM_JOB_NUM_NODES: $SLURM_JOB_NUM_NODES" +echo "SLURM_JOBID: $SLURM_JOBID" +echo "SLURM_JOB_QOS: $SLURM_JOB_QOS" +echo "SLURM_NODELIST: $SLURM_NODELIST" +echo "SLURM_JOB_ACCOUNT: $SLURM_JOB_ACCOUNT" +echo "SLURM_NPROCS: $SLURM_NPROCS" +echo "SLURM_SUBMIT_HOST: $SLURM_SUBMIT_HOST" +echo "SLURM_CONF: $SLURM_CONF" +echo "SLURM_JOB_NAME: $SLURM_JOB_NAME" +echo "SLURM_NTASKS_PER_NODE: $SLURM_NTASKS_PER_NODE" +echo "SLURM_SUBMIT_DIR: $SLURM_SUBMIT_DIR" +echo "SLURM_CLUSTER_NAME: $SLURM_CLUSTER_NAME" +echo "ulimit: $(ulimit -a)" +echo "" +echo "Selected nodes for execution:" +echo "$SELECTED_NODES" +echo "" + +# Node information +MASTER_NODE=$(echo "$SELECTED_NODES" | head -n 1) +MASTER_PORT=39566 # Choose an open port + +IPS=() +if command -v srun >/dev/null 2>&1; then + # Classic Slurm: resolve each node's IP via srun (original behavior). + for NODE in $SELECTED_NODES; do + IP=$(srun --nodes=1 --ntasks=1 --time=00:20:00 --nodelist="$NODE" bash -c 'hostname -I' | awk 'NR==1 {print $1}') + IPS+=("$IP") + done +else + # Spur job body (no srun): use the rank-ordered IPs Spur provides in + # SPUR_PEER_NODES ("ip:port,ip:port,..."). These are the cross-node-routable + # bootstrap addresses; hostnames resolve to public IPs and the local node to + # 127.0.1.1, so DNS is not usable here. + mapfile -t IPS < <(echo "${SPUR_PEER_NODES:-}" | tr ',' '\n' | cut -d: -f1 | sed '/^$/d') + if (( ${#IPS[@]} == 0 )); then + echo "ERROR: no srun and SPUR_PEER_NODES is empty; cannot determine node IPs." >&2 + exit 1 + fi +fi +MASTER_ADDR="${IPS[0]}" + +IPADDRS="${IPS[*]}" +IPADDRS="${IPADDRS// /,}" +echo "Selected node IPs: ${IPADDRS}" + +NNODES=$NUM_NODES + +echo "MASTER_NODE is ${MASTER_NODE}" +echo "MASTER_ADDR is ${MASTER_ADDR}" +echo "MASTER_PORT is ${MASTER_PORT}" +echo "NNODES is ${NNODES}" + +# Per-job log folder (MAD-style): all per-role server/proxy/bench logs +export LOG_PATH="${LOG_ROOT}/${SLURM_JOB_ID}" +mkdir -p "$LOG_PATH" +echo "[disagg-pd] LOG_PATH=$LOG_PATH" + +# ---------------------------------------------------------------------------- +# Launch one container per node via a single srun and hand off to +# vllm_disagg.sh node (rank-based self-select). +# ---------------------------------------------------------------------------- +# IPADDRS is the ordered node IP list (prefill IPs first, then decode). IPS[] +# was built in SELECTED_NODES (alphabetical) order — the same order srun assigns +# SLURM_PROCID — so the rank->role mapping lines up: +# rank 0 -> prefill master + MoRIIO proxy + orchestrator +# rank 1 .. xP-1 -> prefill +# rank xP .. end -> decode +export IPADDRS + +export WIDE_EP_MODE xP yD GPUS_PER_NODE RUN_AFTER_HEALTH HEALTH_TIMEOUT_S DRY_RUN MORIIO_READ_MODE +export IMAGE SHARED_MOUNT CONTAINER_SCRIPTS DISAGG_SCRIPTS_DIR +export CLUSTER_ENV="${CLUSTER_ENV_IN}" +export DOCKER_CONT_NAME="vllm_disagg_${MODEL_NAME:-model}_${SLURM_JOB_ID}" + +# Router env: forwarded into the main container (cluster.sh honors these) and used +# host-side to launch the external router container on rank 0. +export ROUTER_TYPE ROUTER_PORT ROUTER_POLICY VLLM_ROUTER_IMAGE PROXY_PING_PORT GATEWAY_PORT ROUTER_DP_LOCAL +export ROUTER_CONT_NAME="vllm_router_${SLURM_JOB_ID}" + +# Consumed by preflight_node on each node (srun forwards the environment). +export PREFILL_PORT DECODE_PORT PROXY_PORT ROUTER_METRICS_PORT +export SERVE_PORT RPC_PORT KV_PORT LOCAL_PING_PORT +export PREFLIGHT PREFLIGHT_MAX_VRAM_PCT PREFLIGHT_GPU_SETTLE_S + +echo "[disagg-pd] IPADDRS=${IPADDRS}" +echo "[disagg-pd] launching ${NUM_NODES} containers (1/node) -> vllm_disagg.sh node" + +# Refuse a node that is not actually free. $1 = this node's global rank. +# +# `#SBATCH --exclusive` only keeps out other SLURM jobs. A plain `docker run` on +# a node is invisible to the scheduler, so we get allocated nodes whose GPUs and +# ports are already in use. Without this the job pulls the image, loads the model +# and dies minutes later inside vLLM ("Free memory on device cuda:N ... is less +# than desired GPU memory utilization"); worse, on rank 0 a foreign router +# already listening on ROUTER_PORT can satisfy the health gate and take the +# accuracy traffic. Called after the stale-container reap, with a settle window +# so VRAM the reap is still releasing is not counted against us. +preflight_node() { + local rank="$1" + if [ "${PREFLIGHT:-1}" != "1" ]; then + echo "Rank $rank preflight: skipped (PREFLIGHT=0)" + return 0 + fi + + local host max_pct settle problems="" + local -a issues=() # one short tag per fault, for the one-line summary + host="$(hostname)" + max_pct="${PREFLIGHT_MAX_VRAM_PCT:-10}" + settle="${PREFLIGHT_GPU_SETTLE_S:-30}" + + # GPUs: vLLM asks for a fraction of TOTAL, so memory already resident is + # subtracted from what it can get and pushes it under its own threshold. + if command -v rocm-smi >/dev/null 2>&1; then + local csv busy nbusy ncards deadline + deadline=$(( $(date +%s) + settle )) + while :; do + csv="$(rocm-smi --showmeminfo vram --csv 2>/dev/null || true)" + busy="$(printf '%s\n' "$csv" | + awk -F, -v m="$max_pct" '/^card/ && $2+0 > 0 { + pct = $3 * 100 / $2 + if (pct > m) printf "%s=%d%%(%.0fGiB) ", $1, pct, $3/1073741824 + }' || true)" + if [ -z "$busy" ]; then break; fi + if [ "$(date +%s)" -ge "$deadline" ]; then break; fi + sleep 5 + done + if [ -n "$busy" ]; then + nbusy="$(printf '%s' "$busy" | tr ' ' '\n' | grep -c 'card' || true)" + ncards="$(printf '%s\n' "$csv" | grep -c '^card' || true)" + issues+=("${nbusy}/${ncards} GPUs hold VRAM above ${max_pct}%") + problems="${problems} GPU VRAM in use above ${max_pct}%: ${busy}"$'\n' + fi + else + echo "Rank $rank preflight: no rocm-smi on $host, skipping GPU check" + fi + + # Ports this rank will bind. Every container here is --network host, so a + # foreign listener either collides or silently serves our traffic. + local -a want=() + if [ "${WIDE_EP_MODE:-0}" = "1" ]; then + want+=("$SERVE_PORT" "$RPC_PORT" "$KV_PORT" "$LOCAL_PING_PORT") + elif [ "$rank" -lt "${xP:-1}" ]; then + want+=("$PREFILL_PORT") + else + want+=("$DECODE_PORT") + fi + if [ "$rank" = "0" ]; then + want+=("$PROXY_PING_PORT") + if [ "$ROUTER_TYPE" = "vllm-router" ]; then + want+=("$ROUTER_PORT" "$ROUTER_METRICS_PORT") + else + want+=("$PROXY_PORT") + fi + fi + + if command -v ss >/dev/null 2>&1; then + local bound taken="" p + bound="$(ss -lnt 2>/dev/null | awk '{ n = split($4, a, ":"); print a[n] }' | sort -u || true)" + for p in "${want[@]}"; do + if printf '%s\n' "$bound" | grep -qx -- "$p"; then + taken="${taken}${p} " + fi + done + if [ -n "$taken" ]; then + issues+=("ports already bound: ${taken% }") + problems="${problems} ports already bound: ${taken}"$'\n' + fi + else + echo "Rank $rank preflight: no ss on $host, skipping port check" + fi + + if [ -n "$problems" ]; then + mkdir -p "$LOG_PATH" 2>/dev/null || true + local summary + summary="$(printf '%s; ' "${issues[@]}")" + summary="${summary%; }" + # Every node writes the shared job log concurrently, so this diagnostic can + # be overwritten there. Keep a single-writer copy per node as well, or the + # job fails in seconds with no surviving explanation. The submitter reads + # the PREFLIGHT-REJECTED line back out of it to use as the failure reason, + # so keep that prefix in sync with run-slurm-disagg-test.sh. It must not + # contain "FAIL:" — that is the accuracy gate's pattern. + { + echo "PREFLIGHT-REJECTED: rank ${rank} on ${host}: ${summary}" + printf '%s' "$problems" + echo " containers here: $(docker ps --format '{{.Names}} ({{.Image}})' 2>/dev/null | paste -sd'; ' - || true)" + echo " --exclusive does not exclude non-SLURM workloads; clear the node or exclude it from the allocation." + echo " Set PREFLIGHT=0 to submit anyway." + } | tee -a "$LOG_PATH/preflight_NODE${rank}.log" 2>/dev/null || true + # Unblock the other ranks: they park on this sentinel, and the submitter + # polls it. Without it a peer waits out HEALTH_TIMEOUT_S for an endpoint + # that is never coming. + echo 1 > "$LOG_PATH/.disagg_done" 2>/dev/null || true + return 1 + fi + + echo "Rank $rank preflight ok on $host: GPUs idle, ports free (${want[*]})" + return 0 +} + +# Per-node container launch, auto-dispatched by srun availability: +# Launch this node's disagg container and hand off to vllm_disagg.sh. $1 = this +# node's global rank. Relies on the env exported above (forwarded by srun on +# classic Slurm; already in scope when called directly on Spur). +run_node_container() { + local rank="$1" + set -o pipefail # propagate docker run exit code through the tee pipe + echo "Rank $rank on $(hostname)" + local -a stale=() + mapfile -t stale < <(docker ps -aq --filter name=vllm_disagg_ --filter name=vllm_router_ 2>/dev/null || true) + if (( ${#stale[@]} > 0 )); then + echo "Reaping stale containers on $(hostname): ${stale[*]}" + docker stop -t 30 "${stale[@]}" >/dev/null 2>&1 || true + docker rm "${stale[@]}" >/dev/null 2>&1 || true + fi + + preflight_node "$rank" || return 1 + + # Bind-mount the host AMD Pensando (ionic) RoCE userspace provider into the + # container IF present: the image ships an ABI-incompatible copy and the NICs + # are invisible without the host libs. Each entry is added only when the host + # file exists, so this is a no-op on non-ionic (e.g. mlx5) clusters. + local IONIC_MOUNTS=() + local host_ionic + host_ionic=$(readlink -f /usr/lib/x86_64-linux-gnu/libionic.so.1 2>/dev/null || true) + [ -n "$host_ionic" ] && IONIC_MOUNTS+=(-v "$host_ionic:/usr/lib/x86_64-linux-gnu/libionic.so.1:ro") + [ -f /usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so ] && \ + IONIC_MOUNTS+=(-v "/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:/usr/lib/x86_64-linux-gnu/libibverbs/libionic-rdmav34.so:ro") + [ -f /etc/libibverbs.d/ionic.driver ] && \ + IONIC_MOUNTS+=(-v "/etc/libibverbs.d/ionic.driver:/etc/libibverbs.d/ionic.driver:ro") + echo "Rank $rank ionic mounts: ${IONIC_MOUNTS[*]:-}" + + # Pull latest container. + if [ "$DRY_RUN" != "1" ]; then + echo "Rank $rank pulling image on $(hostname): $IMAGE" + docker pull "$IMAGE" || echo "WARN: docker pull failed on $(hostname) for $IMAGE (using local copy if present)" + if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then + echo "Rank $rank pulling router image on $(hostname): $VLLM_ROUTER_IMAGE" + docker pull "$VLLM_ROUTER_IMAGE" || echo "WARN: docker pull failed on $(hostname) for $VLLM_ROUTER_IMAGE (using local copy if present)" + fi + fi + + # Start the external vLLM router as a SEPARATE container on the rank-0 node. + if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ] && [ "$DRY_RUN" = "1" ]; then + echo "DRY_RUN: would start vllm-router on $(hostname): $VLLM_ROUTER_IMAGE (discovery :$PROXY_PING_PORT, port :$ROUTER_PORT, policy $ROUTER_POLICY, dp_local=$ROUTER_DP_LOCAL)" + elif [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then + mkdir -p "$LOG_PATH" 2>/dev/null || true + { + echo "Starting vllm-router on $(hostname): $VLLM_ROUTER_IMAGE (discovery :$PROXY_PING_PORT, port :$ROUTER_PORT, policy $ROUTER_POLICY, dp_local=$ROUTER_DP_LOCAL)" + docker run -d \ + --name "$ROUTER_CONT_NAME" \ + --network host \ + --user "$(id -u):$(id -g)" \ + -v "$LOG_PATH":/run_logs \ + -v "$SHARED_MOUNT":"$SHARED_MOUNT" \ + -v "/data/$(id -un):/workspace" \ + -e HOME=/workspace \ + -e USER="$(id -un)" \ + -e LOGNAME="$(id -un)" \ + "$VLLM_ROUTER_IMAGE" \ + bash -lc "exec vllm-router --vllm-pd-disaggregation --kv-connector moriio --vllm-discovery-address 0.0.0.0:$PROXY_PING_PORT --port $ROUTER_PORT --host 0.0.0.0 --intra-node-data-parallel-size $ROUTER_DP_LOCAL --policy $ROUTER_POLICY --prefill-policy $ROUTER_POLICY --decode-policy $ROUTER_POLICY --log-level info >> $LOG_PATH/vllm_router.log 2>&1" \ + && echo "vllm-router container started ok" \ + || echo "WARN: failed to start vllm-router container on $(hostname)" + } 2>&1 | tee -a "$LOG_PATH/vllm_router_start.log" + fi + + docker run --rm \ + --user "$(id -u):$(id -g)" \ + --device /dev/dri \ + --device /dev/kfd \ + --device /dev/infiniband \ + --network host \ + --ipc host \ + --group-add "$(getent group video | cut -d: -f3)" \ + --group-add "$(getent group render | cut -d: -f3)" \ + --cap-add SYS_PTRACE \ + --cap-add IPC_LOCK \ + --security-opt seccomp=unconfined \ + --shm-size 64G \ + --ulimit nofile=1048576:1048576 \ + -v "$SHARED_MOUNT":"$SHARED_MOUNT" \ + -v "$DISAGG_SCRIPTS_DIR":"$CONTAINER_SCRIPTS" \ + -v "/data/$(id -un):/workspace" \ + "${IONIC_MOUNTS[@]}" \ + --ulimit memlock=-1:-1 \ + -e HF_TOKEN \ + -e HOME=/workspace \ + -e USER="$(id -un)" \ + -e LOGNAME="$(id -un)" \ + -e TRITON_CACHE_DIR=/tmp/triton_cache \ + -e TORCHINDUCTOR_CACHE_DIR=/tmp/inductor_cache \ + -e VLLM_CACHE_ROOT=/tmp/vllm_cache \ + -e AITER_JIT_DIR=/tmp/aiter_jit \ + -e FLYDSL_RUNTIME_CACHE_DIR=/tmp/flydsl_cache \ + -e SLURM_JOB_ID="$SLURM_JOB_ID" \ + -e SLURM_JOB_NODELIST="$SLURM_JOB_NODELIST" \ + -e NODE_RANK="$rank" \ + -e IPADDRS="$IPADDRS" \ + -e WIDE_EP_MODE="$WIDE_EP_MODE" \ + -e xP="$xP" -e yD="$yD" \ + -e GPUS_PER_NODE="$GPUS_PER_NODE" \ + -e CLUSTER_ENV="$CLUSTER_ENV" \ + -e MODEL_NAME="${MODEL_NAME:-}" \ + -e MODEL_DIR="${MODEL_DIR:-}" \ + -e LOG_PATH="$LOG_PATH" \ + -e RUN_AFTER_HEALTH="$RUN_AFTER_HEALTH" \ + -e HEALTH_TIMEOUT_S="$HEALTH_TIMEOUT_S" \ + -e DRY_RUN="$DRY_RUN" \ + -e MORIIO_READ_MODE="$MORIIO_READ_MODE" \ + -e ROUTER_TYPE="$ROUTER_TYPE" \ + -e ROUTER_PORT="$ROUTER_PORT" \ + -e ROUTER_POLICY="$ROUTER_POLICY" \ + -e GATEWAY_PORT="$GATEWAY_PORT" \ + --name "$DOCKER_CONT_NAME" \ + --entrypoint /bin/bash \ + "$IMAGE" -c "bash $CONTAINER_SCRIPTS/vllm_disagg.sh node" \ + 2>&1 | tee "$LOG_PATH/main_NODE${rank}.log" + local node_rc=${PIPESTATUS[0]} # docker run exit code (not tee), for rank-0 pass/fail + # Stop the external router (rank-0 node) once the main container returns. + if [ "$ROUTER_TYPE" = "vllm-router" ] && [ "$rank" = "0" ]; then + docker stop "$ROUTER_CONT_NAME" 2>/dev/null || true; docker rm "$ROUTER_CONT_NAME" 2>/dev/null || true; + fi + return "$node_rc" +} + +# Teardown of the main + external router containers on every node ($1 = docker +# stop timeout). `docker run` containers outlive the client that started them, +# so on scancel the srun tasks die but the containers keep the GPUs until the +# walltime expires unless they are stopped explicitly. +teardown_containers() { + export TEARDOWN_STOP_T="${1:-30}" + if command -v srun >/dev/null 2>&1; then + # shellcheck disable=SC2016 # must expand on the remote node, not here + srun --nodelist="$SELECTED_NODELIST_STR" \ + --nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \ + bash -c 'docker stop -t "${TEARDOWN_STOP_T:-30}" "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true; docker rm "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true + [ "$ROUTER_TYPE" = "vllm-router" ] && { docker stop -t "${TEARDOWN_STOP_T:-30}" "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true; docker rm "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true; }' || true + else + docker stop -t "$TEARDOWN_STOP_T" "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true + docker rm "$DOCKER_CONT_NAME" >/dev/null 2>&1 || true + if [ "$ROUTER_TYPE" = "vllm-router" ]; then + docker stop -t "$TEARDOWN_STOP_T" "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true + docker rm "$ROUTER_CONT_NAME" >/dev/null 2>&1 || true + fi + fi + return 0 +} + +# SLURM sends SIGTERM then SIGKILL after KillWait (30s by default), so the trap +# uses a short stop timeout to fit inside that window. +trap 'echo "[disagg-pd] signal received, tearing down containers"; teardown_containers 10; exit 143' TERM INT + +RC=0 +if command -v srun >/dev/null 2>&1; then + # classic Slurm: fan out one task/node; ship the fn and call it with the + # srun-supplied per-task SLURM_PROCID (escaped so the remote shell expands it). + srun --label \ + --nodelist="$SELECTED_NODELIST_STR" \ + --nodes="$NUM_NODES" --ntasks="$NUM_NODES" --ntasks-per-node=1 \ + bash -c "$(declare -f preflight_node run_node_container); run_node_container \"\$SLURM_PROCID\"" || RC=$? +else + # Spur: the body already runs once per node; this node's rank is + # SPUR_TASK_OFFSET (fall back to NODE_RANK, then 0 for a single-node run). + run_node_container "${SPUR_TASK_OFFSET:-${NODE_RANK:-0}}" || RC=$? +fi + +teardown_containers 30 + +echo "[disagg-pd] launch finished rc=${RC}; logs in ${LOG_PATH}" +exit "${RC}" diff --git a/.buildkite/amd-disagg/vllm_disagg.sh b/.buildkite/amd-disagg/vllm_disagg.sh new file mode 100644 index 000000000000..0974862a38b4 --- /dev/null +++ b/.buildkite/amd-disagg/vllm_disagg.sh @@ -0,0 +1,635 @@ +#!/bin/bash +# ============================================================================= +# vLLM disaggregated (P/D) launcher (rank-based / SLURM-native). +# ----------------------------------------------------------------------------- +# +# rank 0 .. xP-1 -> prefill (rank 0 is also the orchestrator + proxy) +# rank xP .. xP+yD-1 -> decode +# +# Rank 0 additionally: starts the MoRIIO proxy, health-gates every server, +# runs the post-health workload (RUN_AFTER_HEALTH=bench|accuracy|none), then +# drops a shared-FS completion sentinel so the other ranks shut down and the +# single `srun` returns. Rank 0's exit code is the run's pass/fail. +# +# Explicit per-role invocation is also supported (granular for manual use): +# proxy | prefill | decode | bench | accuracy. +# +# ---- LAUNCH IT (both supported) -------------------------------------------- +# 1) Single srun (SLURM provides the rank) — what run_xPyD_disagg.slurm does: +# srun --nodes=$((xP+yD)) --ntasks-per-node=1 \ +# bash .../vllm_disagg.sh node +# (IPADDRS must be ordered to match ranks: prefill IPs first, then decode, +# in the SAME order srun assigns PROCID — sort nodes alphabetically.) +# +# 2) Manual, one shell per node (rank from NODE_RANK; share LOG_PATH/FS): +# IPADDRS=ipP0,ipD0 xP=1 yD=1 NODE_RANK=0 bash vllm_disagg.sh node +# IPADDRS=ipP0,ipD0 xP=1 yD=1 NODE_RANK=1 bash vllm_disagg.sh node +# +# 3) Manual, fully granular (per-role): +# bash vllm_disagg.sh proxy +# bash vllm_disagg.sh prefill +# NODE_RANK=1 bash vllm_disagg.sh decode +# bash vllm_disagg.sh bench +# +# Cluster config: cluster.sh (sourced). Model flags: models.yaml. +# Parallelism: WIDE_EP_MODE=0 tp (independent TP servers) | 1 ep (DP+EP groups). +# ============================================================================= +set -euo pipefail + +SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +_PROXY_IP_OVERRIDE="${PROXY_IP:-}" + +log() { echo "[vllm_disagg] $*"; } +die() { echo "ERROR: $*" >&2; exit 1; } + +usage() { + echo "Usage: $0 [node|proxy|prefill|decode|bench|accuracy] [--wide-ep-mode 0|1]" >&2 + echo " (no role => 'node': derive role from \$SLURM_PROCID/\$NODE_RANK)" >&2 + exit 1 +} + +# ----------------------------------------------------------------- arg parsing +# Sets globals: ROLE (default 'node'), _WIDE_EP_MODE_OVERRIDE. +parse_args() { + ROLE="${1:-node}" + shift || true + _WIDE_EP_MODE_OVERRIDE="" + while [[ $# -gt 0 ]]; do + case "$1" in + --wide-ep-mode) _WIDE_EP_MODE_OVERRIDE="${2:-}"; shift 2 ;; + --wide-ep-mode=*) _WIDE_EP_MODE_OVERRIDE="${1#*=}"; shift ;; + -h|--help) usage ;; + *) die "Unknown argument: $1" ;; + esac + done +} + +# ----------------------------------------------------------------- load config +# Sources cluster.sh and resolves WIDE_EP_MODE -> PARALLEL_MODE (tp|ep). +load_config() { + CLUSTER_ENV="${CLUSTER_ENV:-${SCRIPT_DIR}/cluster.sh}" + [[ -f "${CLUSTER_ENV}" ]] || die "cluster env file not found: ${CLUSTER_ENV}" + # shellcheck disable=SC1090 + source "${CLUSTER_ENV}" + + # Topology fanout normally arrives from cluster.sh; restate the defaults so + # the rank arithmetic is defined even when it does not. + : "${xP:=1}" "${yD:=1}" + + [[ -n "${_WIDE_EP_MODE_OVERRIDE}" ]] && WIDE_EP_MODE="${_WIDE_EP_MODE_OVERRIDE}" + WIDE_EP_MODE="${WIDE_EP_MODE:-0}" + case "${WIDE_EP_MODE}" in + 0) PARALLEL_MODE="tp" ;; + 1) PARALLEL_MODE="ep" ;; + *) die "WIDE_EP_MODE must be 0 (tp) or 1 (ep); got '${WIDE_EP_MODE}'" ;; + esac + + # node-mode orchestration knobs + RUN_AFTER_HEALTH="${RUN_AFTER_HEALTH:-accuracy}" # bench | accuracy | none + HEALTH_TIMEOUT_S="${HEALTH_TIMEOUT_S:-3600}" + + # MoRIIO KV transfer direction. 0 (default) + MORIIO_READ_MODE="${MORIIO_READ_MODE:-0}" +} + +# Emits the `read_mode` KV-config fragment (leading comma + newline) when read +# mode is enabled, else nothing (leaving MoRIIO's write-mode default). +moriio_read_mode_kv() { + case "${MORIIO_READ_MODE:-0}" in + 1|true|True|TRUE|yes|on) printf ',\n "read_mode": true' ;; + *) : ;; + esac +} + +# ----------------------------------------------------------------- topology +# Resolves MODEL_PATH, MODELS_YAML, the IP array, and the master/proxy addresses. +# IP_ARRAY[0..xP-1] = prefill nodes, [xP..xP+yD-1] = decode nodes. +resolve_topology() { + MODEL_PATH="${MODEL_DIR%/}/${MODEL_NAME}" + MODELS_YAML="${MODELS_YAML:-${SCRIPT_DIR}/models.yaml}" + mkdir -p "${LOG_PATH}" 2>/dev/null || true + + [[ -z "${IPADDRS}" ]] && IPADDRS="${PREFILL_IP},${DECODE_IP}" # 1P1D fallback + IFS=',' read -ra IP_ARRAY <<< "${IPADDRS}" + PREFILL_MASTER_ADDR="${IP_ARRAY[0]}" + DECODE_MASTER_ADDR="${IP_ARRAY[$xP]:-${IP_ARRAY[-1]}}" + # Proxy is co-located on the prefill master (rank 0) unless the CALLER + # explicitly set PROXY_IP. We use the pre-cluster.sh override (captured at the + # top) — not the live PROXY_IP, which cluster.sh always populates with the + # static PREFILL_IP fallback and would otherwise mask the real master. + PROXY_IP="${_PROXY_IP_OVERRIDE:-${PREFILL_MASTER_ADDR}}" +} + +# ============================================================================ +# Non-server roles +# ============================================================================ + +run_proxy() { + [[ -f "${PROXY_SCRIPT}" ]] || die "proxy script not found: ${PROXY_SCRIPT}" + log "proxy: http=:${PROXY_PORT} discovery=:${PROXY_PING_PORT} (${PROXY_SCRIPT})" + exec python3 "${PROXY_SCRIPT}" --port "${PROXY_PORT}" +} + +# Start the proxy in the BACKGROUND (node-mode rank 0). Sets PROXY_PID. +start_proxy_bg() { + [[ -f "${PROXY_SCRIPT}" ]] || die "proxy script not found: ${PROXY_SCRIPT}" + local plog="${LOG_PATH}/proxy_$(date +%Y%m%d_%H%M%S).log" + log "proxy(bg): http=:${PROXY_PORT} discovery=:${PROXY_PING_PORT} log=${plog}" + python3 "${PROXY_SCRIPT}" --port "${PROXY_PORT}" >"${plog}" 2>&1 & + PROXY_PID=$! +} + +run_bench() { + local base_url="http://127.0.0.1:${GATEWAY_PORT:-${PROXY_PORT}}" + local result_dir="${LOG_PATH}/bench_$(date +%Y%m%d_%H%M%S)" + mkdir -p "${result_dir}" + log "bench -> ${base_url} model=${MODEL_PATH} wide_ep_mode=${WIDE_EP_MODE}(${PARALLEL_MODE})" + log "combinations='${BENCHMARK_COMBINATIONS}' concurrency='${BENCHMARK_CON}'" + + local combo isl osl con nump logf + for combo in ${BENCHMARK_COMBINATIONS}; do + isl="${combo%%/*}"; osl="${combo##*/}" + for con in ${BENCHMARK_CON}; do + nump=$(( NUM_PROMPTS_FACTOR * con )) + (( nump < BENCHMARK_MIN_PROMPTS )) && nump="${BENCHMARK_MIN_PROMPTS}" + logf="${result_dir}/bench_${isl}_${osl}_con${con}.log" + log "ISL=${isl} OSL=${osl} CON=${con} num_prompts=${nump} -> ${logf}" + vllm bench serve \ + --backend openai \ + --base-url "${base_url}" \ + --endpoint /v1/completions \ + --model "${MODEL_PATH}" \ + --dataset-name random \ + --random-input-len "${isl}" \ + --random-output-len "${osl}" \ + --num-prompts "${nump}" \ + --max-concurrency "${con}" \ + --percentile-metrics ttft,tpot,itl,e2el \ + --ignore-eos \ + 2>&1 | tee "${logf}" + done + done + log "bench complete. Results in ${result_dir}" +} + +run_accuracy() { + # lm_eval isn't in the stock vLLM image; install on demand (rank 0 only, and + # only when accuracy is actually selected). Override with ACCURACY_PIP_SPEC. + if ! command -v lm_eval >/dev/null 2>&1; then + log "lm_eval not found — installing ${ACCURACY_PIP_SPEC:-lm_eval[api]}" + python3 -m pip install --no-cache-dir "${ACCURACY_PIP_SPEC:-lm_eval[api]}" + fi + local base_url="http://127.0.0.1:${GATEWAY_PORT:-${PROXY_PORT}}/v1/completions" + local ts; ts="$(date +%Y%m%d_%H%M%S)" + local logf="${LOG_PATH}/accuracy_${MODEL_NAME}_${PARALLEL_MODE}_${ts}.log" + local outdir="${LOG_PATH}/lm_eval_${MODEL_NAME}_${PARALLEL_MODE}_${ts}" + log "accuracy -> ${base_url} model=${MODEL_PATH} tasks=${ACCURACY_TASKS} wide_ep_mode=${WIDE_EP_MODE}(${PARALLEL_MODE})" + log "log=${logf} results=${outdir}" + + # --output_path makes lm_eval persist a results_*.json for scraping. + local eval_rc=0 + ( set -o pipefail + # Model weights are already served (loaded offline in a separate process); + # allow the Hub only here so lm_eval can fetch the eval dataset (e.g. gsm8k). + export HF_HUB_OFFLINE=0 HF_DATASETS_OFFLINE=0 + python3 -m lm_eval --model local-completions \ + --tasks "${ACCURACY_TASKS}" \ + --model_args "model=${MODEL_PATH},base_url=${base_url},num_concurrent=${ACCURACY_NUM_CONCURRENT},max_retries=${ACCURACY_MAX_RETRIES},tokenized_requests=False,trust_remote_code=True,timeout=${ACCURACY_TIMEOUT:-3600}" \ + --output_path "${outdir}" \ + 2>&1 | tee "${logf}" + ) || eval_rc=$? + if (( eval_rc != 0 )); then + log "FAIL: lm_eval exited rc=${eval_rc} (no usable score) — Log: ${logf}" + return "${eval_rc}" + fi + + # Gate on the score. Find the results JSON lm_eval just wrote, pull the max + # value of ACCURACY_METRIC across filters/tasks, and compare to the threshold. + local results_json="" + results_json="$(find "${outdir}" -name 'results_*.json' -type f 2>/dev/null | sort | tail -n1)" + if [[ -z "${results_json}" ]]; then + log "FAIL: no results_*.json under ${outdir} — cannot evaluate threshold" + return 1 + fi + + local score="" + score="$(python3 - "${results_json}" "${ACCURACY_METRIC}" <<'PY' +import json, sys +path, metric = sys.argv[1], sys.argv[2] +with open(path) as f: + data = json.load(f) +best = None +for task, metrics in (data.get("results") or {}).items(): + for key, val in metrics.items(): + if isinstance(val, bool) or not isinstance(val, (int, float)): + continue + if key.split(",")[0] == metric: + best = val if best is None else max(best, val) +print("" if best is None else "%.6f" % best) +PY +)" || score="" + + if [[ -z "${score}" ]]; then + log "FAIL: could not parse metric '${ACCURACY_METRIC}' from ${results_json}" + return 1 + fi + + if awk -v s="${score}" -v t="${ACCURACY_THRESHOLD}" 'BEGIN{exit !(s+0 >= t+0)}'; then + log "PASS: ${ACCURACY_METRIC}=${score} >= threshold=${ACCURACY_THRESHOLD} (tasks=${ACCURACY_TASKS}). Log: ${logf}" + return 0 + fi + log "FAIL: ${ACCURACY_METRIC}=${score} < threshold=${ACCURACY_THRESHOLD} (tasks=${ACCURACY_TASKS}). Log: ${logf}" + return 1 +} + +# ============================================================================ +# Server roles (prefill / decode) and their shared helpers +# ============================================================================ + +# Read model-specific flags from models.yaml (mode + role aware) into globals: +# MODEL_BASE_FLAGS / MODEL_ROLE_FLAGS / MODEL_EXPERIMENTAL_FLAGS -> MODEL_CONFIG +# Also exports the model's env: block (only if not already set; caller env wins). +load_model_flags() { + [[ -f "${MODELS_YAML}" ]] || die "models.yaml not found: ${MODELS_YAML}" + export MODELS_YAML MODEL_NAME PARALLEL_MODE ROLE + eval "$(python3 - <<'PY' +import os, shlex, sys, yaml +path = os.environ["MODELS_YAML"]; name = os.environ["MODEL_NAME"] +mode = os.environ["PARALLEL_MODE"]; role = os.environ["ROLE"] +with open(path, "r", encoding="utf-8") as f: + doc = yaml.safe_load(f) or {} +# Preferred form: top-level `models:` list of {model: , ...} entries. +# Back-compat: a top-level mapping keyed by model name. +entries = doc.get("models") if isinstance(doc, dict) else doc +cfg = None +if isinstance(entries, list): + cfg = next((e for e in entries if isinstance(e, dict) and e.get("model") == name), None) +elif isinstance(doc, dict): + cfg = doc.get(name) +if cfg is None: + print(f'echo "ERROR: model {name} not found in {path}" >&2; exit 1'); sys.exit(0) +role_cfg = cfg.get(role, {}) or {} +def q(v): return shlex.quote(str(v if v is not None else "")) +exports = { + "MODEL_BASE_FLAGS": cfg.get("base_flags", "") or "", + "MODEL_ROLE_FLAGS": role_cfg.get(mode, "") or "", + "MODEL_EXPERIMENTAL_FLAGS": cfg.get("experimental_flags", "") or "", +} +for k, v in exports.items(): print(f"{k}={q(v)}") +# Model-specific env: exported only if not already set (caller env wins). +for k, v in (cfg.get("env", {}) or {}).items(): + if k not in os.environ: + print(f"export {k}={q(v)}") +PY +)" + MODEL_CONFIG="${MODEL_BASE_FLAGS} ${MODEL_ROLE_FLAGS} ${MODEL_EXPERIMENTAL_FLAGS}" +} + +# Build the TP-mode (WIDE_EP_MODE=0) serve command into the global CMD array. +# Each node is an independent TP server registered to the proxy. +build_tp_cmd() { + local port="${TP_PORT}" + [[ -n "${HOST_IP}" ]] && CMD+=(--host "${HOST_IP}") + CMD+=(--port "${port}" --tensor-parallel-size "${TP_SIZE}") + local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}") + local kv_cfg read_mode_kv + read_mode_kv="$(moriio_read_mode_kv)" + kv_cfg=$(cat < mori_high_throughput (paired with --enforce-eager in models.yaml) + # decode -> mori_low_latency (paired with CUDA graphs in models.yaml) + local default_backend="mori_low_latency" + [[ "${ROLE}" == "prefill" ]] && default_backend="mori_high_throughput" + local backend="${ALL2ALL_BACKEND:-${default_backend}}" + export VLLM_ALL2ALL_BACKEND="${VLLM_ALL2ALL_BACKEND:-${backend}}" + local port="${SERVE_PORT}" + CMD+=( + -tp 1 + --data-parallel-size "${DP_GROUP_SIZE}" + --data-parallel-size-local "${GPUS_PER_NODE}" + --data-parallel-address "${DP_MASTER_ADDR}" + --data-parallel-rpc-port "${RPC_PORT}" + --enable-expert-parallel + --all2all-backend "${backend}" + --port "${port}" + --no-enable-prefix-caching + --distributed-timeout-seconds "${DISTRIBUTED_TIMEOUT_SECONDS}" + ) + local _mc; read -ra _mc <<< "${MODEL_CONFIG}"; CMD+=("${_mc[@]}") + + if [[ "${IS_MASTER}" == "1" ]]; then + CMD+=(--api-server-count="${GPUS_PER_NODE}") + local kv_cfg read_mode_kv + read_mode_kv="$(moriio_read_mode_kv)" + kv_cfg=$(cat <&1 | tee "${LOGF}" +} + +run_prefill() { configure_prefill; run_server_fg; } +run_decode() { configure_decode; run_server_fg; } + +# ============================================================================ +# node mode — SLURM-native rank-based self-select + rank-0 orchestration +# ============================================================================ + +# Endpoints to health-check, one per line as "ip:port", depending on mode. +# tp: every prefill node (:PREFILL_PORT) and every decode node (:DECODE_PORT) +# ep: only the masters expose an API (prefill rank 0, decode rank xP) :SERVE_PORT +health_endpoints() { + local i + if [[ "${WIDE_EP_MODE}" == "0" ]]; then + for (( i=0; i/dev/null 2>&1; do + (( $(date +%s) >= deadline )) && { log "TIMEOUT waiting for ${ep}"; return 1; } + kill -0 "${SERVER_PID}" 2>/dev/null || { log "local server (pid ${SERVER_PID}) exited while waiting"; return 1; } + sleep 10 + done + log "healthy: ${ep}" + done + return 0 +} + +run_workload() { + case "${RUN_AFTER_HEALTH}" in + bench) run_bench ;; + accuracy) run_accuracy ;; + none|"") log "RUN_AFTER_HEALTH=none — skipping workload" ;; + *) die "unknown RUN_AFTER_HEALTH='${RUN_AFTER_HEALTH}'" ;; + esac +} + +# rank 0: proxy + health-gate + workload, then write the completion sentinel. +orchestrate_master() { + local sentinel="$1" rc=0 + # Front door: the proxy runs in-container (started here); the vllm-router + # runs as a SEPARATE container started by the SLURM job on this (rank-0) node, + # so in that mode we don't start anything here. + PROXY_PID="" + if [[ "${ROUTER_TYPE:-vllm-router}" == "vllm-router" ]]; then + log "ROUTER_TYPE=vllm-router: external router expected on gateway :${GATEWAY_PORT:-${ROUTER_PORT}} (not starting proxy)" + else + start_proxy_bg + fi + if wait_all_healthy; then + # If using the external vllm-router, wait for it to become ready before + # running the workload. The router starts as a separate container on this + # (rank-0) node; give it up to 300s to bind port ROUTER_PORT. + if [[ "${ROUTER_TYPE:-vllm-router}" == "vllm-router" ]]; then + local router_deadline=$(( $(date +%s) + 300 )) + log "waiting for vllm-router on :${ROUTER_PORT} (up to 300s)" + until /usr/bin/curl -sf "http://127.0.0.1:${ROUTER_PORT}/health" >/dev/null 2>&1; do + if (( $(date +%s) >= router_deadline )); then + log "FAIL: vllm-router not ready on :${ROUTER_PORT} after 300s (bring-up)" + rc=1 + break + fi + sleep 5 + done + (( rc == 0 )) && log "vllm-router healthy on :${ROUTER_PORT}" + fi + if (( rc == 0 )); then + set +e + run_workload + rc=$? + set -e + fi + else + log "FAIL: health-gate did not pass — server bring-up failed" + rc=1 + fi + echo "${rc}" > "${sentinel}" 2>/dev/null || true + log "master: workload rc=${rc}; sentinel=${sentinel}; tearing down local proxy+server" + [[ -n "${PROXY_PID:-}" ]] && kill "${PROXY_PID}" 2>/dev/null || true + kill "${SERVER_PID}" 2>/dev/null || true + wait "${SERVER_PID}" 2>/dev/null || true + return "${rc}" +} + +# non-master ranks: keep serving until the master signals completion (sentinel) +# or our local server dies unexpectedly. +watch_until_done() { + local sentinel="$1" rc=0 + log "rank ${NODE_RANK} serving; waiting for completion sentinel (${sentinel})" + while :; do + [[ -f "${sentinel}" ]] && { log "completion sentinel seen; shutting down"; break; } + if ! kill -0 "${SERVER_PID}" 2>/dev/null; then + log "local server (pid ${SERVER_PID}) exited before completion"; rc=1; break + fi + sleep 10 + done + kill "${SERVER_PID}" 2>/dev/null || true + wait "${SERVER_PID}" 2>/dev/null || true + return "${rc}" +} + +run_node() { + NODE_RANK="${NODE_RANK:-${SLURM_PROCID:-0}}" + local total=$(( xP + yD )) + (( NODE_RANK >= 0 && NODE_RANK < total )) \ + || die "NODE_RANK=${NODE_RANK} out of range [0,${total}) (xP=${xP} yD=${yD})" + + if (( NODE_RANK < xP )); then ROLE="prefill"; configure_prefill + else ROLE="decode"; configure_decode + fi + + # Shared-FS completion sentinel (LOG_PATH is shared & per-run/per-job). + local sentinel="${LOG_PATH}/.disagg_done" + + log "node mode: rank=${NODE_RANK}/${total} role=${ROLE} master=${IS_MASTER} wide_ep=${WIDE_EP_MODE}(${PARALLEL_MODE})" + log "topology: xP=${xP} yD=${yD} gpus/node=${GPUS_PER_NODE} IPADDRS=${IPADDRS}" + + # DRY_RUN: resolve+print the plan (role, serve cmd, endpoints) and exit 0. + # Validates wiring without launching vLLM / the proxy / any workload. + if [[ "${DRY_RUN:-0}" == "1" ]]; then + build_server_cmd + log "DRY_RUN: server cmd: ${CMD[*]}" + if (( NODE_RANK == 0 )); then + local eps; mapfile -t eps < <(health_endpoints) + log "DRY_RUN: health endpoints: ${eps[*]}" + log "DRY_RUN: post-health workload: ${RUN_AFTER_HEALTH}" + log "DRY_RUN: sentinel: ${sentinel}" + fi + log "DRY_RUN complete (rank ${NODE_RANK}) — no processes launched" + exit 0 + fi + + # aiter and flydsl keep runtime JIT/kernel caches inside the aiter install + # tree (…/aiter/jit and …/aiter/jit/flydsl_cache), which is read-only for the + # non-root (--user) container. AITER_JIT_DIR / FLYDSL_RUNTIME_CACHE_DIR + # redirect those to node-local /tmp (both are sanctioned overrides: aiter only + # points flydsl at its bundled cache when FLYDSL_RUNTIME_CACHE_DIR is unset). + # Seed each once from the image's bundled copy so we keep the prebuilt + # modules/kernels (fast, no NFS, no from-source rebuild / MLIR recompile). + local _aiter_src + _aiter_src="$(python3 -c 'import os,aiter; print(os.path.join(os.path.dirname(aiter.__file__),"jit"))' 2>/dev/null || true)" + if [[ -n "${AITER_JIT_DIR:-}" && ! -e "${AITER_JIT_DIR}/.seeded" ]]; then + if [[ -n "${_aiter_src}" && -d "${_aiter_src}" ]]; then + mkdir -p "${AITER_JIT_DIR}" + cp -a "${_aiter_src}/." "${AITER_JIT_DIR}/" 2>/dev/null || true + touch "${AITER_JIT_DIR}/.seeded" 2>/dev/null || true + log "seeded AITER_JIT_DIR=${AITER_JIT_DIR} from ${_aiter_src}" + else + log "WARN: could not locate aiter jit dir to seed AITER_JIT_DIR=${AITER_JIT_DIR}; aiter may rebuild from source" + fi + fi + if [[ -n "${FLYDSL_RUNTIME_CACHE_DIR:-}" && ! -e "${FLYDSL_RUNTIME_CACHE_DIR}/.seeded" ]]; then + mkdir -p "${FLYDSL_RUNTIME_CACHE_DIR}" + if [[ -n "${_aiter_src}" && -d "${_aiter_src}/flydsl_cache" ]]; then + cp -a "${_aiter_src}/flydsl_cache/." "${FLYDSL_RUNTIME_CACHE_DIR}/" 2>/dev/null || true + log "seeded FLYDSL_RUNTIME_CACHE_DIR=${FLYDSL_RUNTIME_CACHE_DIR} from ${_aiter_src}/flydsl_cache" + else + log "FLYDSL_RUNTIME_CACHE_DIR=${FLYDSL_RUNTIME_CACHE_DIR} set; no bundled flydsl_cache to seed (will JIT-compile at runtime)" + fi + touch "${FLYDSL_RUNTIME_CACHE_DIR}/.seeded" 2>/dev/null || true + fi + + # mori JIT-compiles its shmem/all2all kernels to ~/.mori/jit/_/ on + # first EP use. HOME here is /workspace, bind-mounted to persistent+shared NFS + # (/data/$USER), so a wrong-arch cache from an earlier run/build survives and + # gets reloaded even after MORI_GPU_ARCHS changes -- an existing cache dir + # wins over the arch setting. Reloading a gfx942 shmem_kernels.hsaco on gfx950 + # hardware fails with "device kernel image is invalid" and every EP worker + # dies at init (surfacing async as a torch.tensor HIP error). Guard: on EP + # runs, if the cache has no build for our target arch, purge it so mori + # recompiles for MORI_GPU_ARCHS. A matching arch dir is kept (fast reload). + if [[ "${WIDE_EP_MODE}" == "1" && -n "${MORI_GPU_ARCHS:-}" ]]; then + local _mori_jit="${HOME:-/workspace}/.mori/jit" + if [[ -d "${_mori_jit}" ]] && ! compgen -G "${_mori_jit}/${MORI_GPU_ARCHS}_*" >/dev/null 2>&1; then + log "purging stale mori jit cache ${_mori_jit} (no ${MORI_GPU_ARCHS}_* build present; forcing recompile for ${MORI_GPU_ARCHS})" + rm -rf "${_mori_jit}" 2>/dev/null || true + fi + fi + + # Start this node's server in the background (full server log -> LOGF; this + # shell's stdout stays free for orchestration logs the srun/CI captures). + build_server_cmd + "${CMD[@]}" >"${LOGF}" 2>&1 & + SERVER_PID=$! + log "server started pid=${SERVER_PID} log=${LOGF}" + + local rc=0 + if (( NODE_RANK == 0 )); then + orchestrate_master "${sentinel}" || rc=$? + # Canonical, human-readable end-of-run verdict for the login-node poller + # and CI logs (the .disagg_done sentinel remains the machine authority). + log "VERDICT: $( (( rc == 0 )) && echo PASS || echo FAIL) rc=${rc}" + else + watch_until_done "${sentinel}" || rc=$? + fi + log "node rank=${NODE_RANK} exiting rc=${rc}" + exit "${rc}" +} + +# ============================================================================ +main() { + parse_args "$@" + load_config + resolve_topology + case "${ROLE}" in + node) run_node ;; + proxy) run_proxy ;; + bench) run_bench ;; + accuracy) run_accuracy ;; + prefill) run_prefill ;; + decode) run_decode ;; + *) usage ;; + esac +} + +main "$@" From d40c3e3c00ab904928ab51c9c62b7f514c39647d Mon Sep 17 00:00:00 2001 From: xijiaat <52128022+xijiaat@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:12:19 +0800 Subject: [PATCH 05/30] Fix DSpark warmup without sparse index buffer (#50693) Signed-off-by: xijiade.aihemaiti <3146335281@qq.com> Signed-off-by: xijiaat <52128022+xijiaat@users.noreply.github.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Co-authored-by: OpenAI Codex --- vllm/models/deepseek_v4/nvidia/flashmla.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/vllm/models/deepseek_v4/nvidia/flashmla.py b/vllm/models/deepseek_v4/nvidia/flashmla.py index 29ec2d6f2abf..f0b00163d1a6 100644 --- a/vllm/models/deepseek_v4/nvidia/flashmla.py +++ b/vllm/models/deepseek_v4/nvidia/flashmla.py @@ -96,8 +96,11 @@ def forward_mqa( // self.compress_ratio ) M = N + self.window_size + self.max_num_batched_tokens - assert self.topk_indices_buffer is not None - top_k = 0 if swa_only else self.topk_indices_buffer.shape[-1] + if swa_only: + top_k = 0 + else: + assert self.topk_indices_buffer is not None + top_k = self.topk_indices_buffer.shape[-1] combined_topk = round_up(top_k + self.window_size, 128) current_workspace_manager().get_simultaneous( ((self.PREFILL_CHUNK_SIZE, M, q.shape[-1]), torch.bfloat16), From 3f142bd85e962a795c6c484f0249f9c92aed37b2 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Mon, 10 Aug 2026 15:30:30 -0400 Subject: [PATCH 06/30] [Bugfix] Align deepseek v4 parser thinking default with tokenizer (#51296) Signed-off-by: sfeng33 <4florafeng@gmail.com> --- tests/parser/engine/test_deepseek_v4.py | 59 +++++++++++++++++++++---- tests/parser/engine/trace_builder.py | 2 +- vllm/parser/deepseek_v4.py | 8 ++-- 3 files changed, 57 insertions(+), 12 deletions(-) diff --git a/tests/parser/engine/test_deepseek_v4.py b/tests/parser/engine/test_deepseek_v4.py index e5e7b075bac2..ba5ebf76ed67 100644 --- a/tests/parser/engine/test_deepseek_v4.py +++ b/tests/parser/engine/test_deepseek_v4.py @@ -243,15 +243,29 @@ def test_thinking_false_starts_in_content(self): cfg = deepseek_v4_config(thinking=False) assert cfg.initial_state.name == "CONTENT" - def test_enable_thinking_kwarg(self, mock_tokenizer): - p = DeepSeekV4Parser( - mock_tokenizer, chat_template_kwargs={"enable_thinking": True} + @pytest.mark.parametrize( + ("chat_template_kwargs", "expected_state"), + [ + ({}, "REASONING"), + ({"thinking": True}, "REASONING"), + ({"enable_thinking": True}, "REASONING"), + ({"reasoning_effort": "high"}, "REASONING"), + ({"thinking": False}, "CONTENT"), + ({"enable_thinking": False}, "CONTENT"), + ( + {"enable_thinking": True, "reasoning_effort": "none"}, + "CONTENT", + ), + ], + ) + def test_parser_thinking_mode_matches_tokenizer_default( + self, mock_tokenizer, chat_template_kwargs, expected_state + ): + parser = DeepSeekV4Parser( + mock_tokenizer, + chat_template_kwargs=chat_template_kwargs, ) - assert p.parser_engine_config.initial_state.name == "REASONING" - - def test_no_thinking_kwarg_defaults_to_content(self, mock_tokenizer): - p = DeepSeekV4Parser(mock_tokenizer) - assert p.parser_engine_config.initial_state.name == "CONTENT" + assert parser.parser_engine_config.initial_state.name == expected_state def test_thinking_mode_reasoning_without_tags(self, mock_tokenizer): parser = DeepSeekV4Parser( @@ -838,6 +852,35 @@ def test_tool_calls_extracted_at_all_chunk_sizes( args = json.loads(output.tool_calls[0]["arguments"]) assert args == {"location": "Berlin", "units": "celsius"} + def test_default_thinking_extracts_tool_call_without_think_end(self, dsv4_tokens): + tokens = [ + token + for token in dsv4_tokens + if token[0] != _DSV4_FULL_VOCAB[DSML_THINK_END] + ] + tokenizer = MockTokenizer( + vocab=dict(_DSV4_FULL_VOCAB), + tokens=tokens, + ) + parser = _DeepSeekV4Delegating(tokenizer) + + deltas = replay_streaming( + parser, + tokens, + chunk_size=1, + finished_on_last=True, + tools=DUMMY_TOOLS, + prompt_token_ids=[_DSV4_FULL_VOCAB[DSML_THINK_START]], + ) + output = collect_output(deltas) + + assert "The user wants" in output.reasoning + assert output.content == "" + assert len(output.tool_calls) == 1 + assert output.tool_calls[0]["name"] == "get_weather" + args = json.loads(output.tool_calls[0]["arguments"]) + assert args == {"location": "Berlin", "units": "celsius"} + def test_eos_drop_token_does_not_swallow_tool_calls(self): """Tool calls must survive when an EOS DROP token's ID is in delta_token_ids but its text is absent from delta_text. diff --git a/tests/parser/engine/trace_builder.py b/tests/parser/engine/trace_builder.py index f24e3318f975..6d5415dfdaba 100644 --- a/tests/parser/engine/trace_builder.py +++ b/tests/parser/engine/trace_builder.py @@ -722,7 +722,7 @@ def _dsv4_segments(scenario: Scenario, thinking: bool) -> list[tuple[str, bool]] def _build_deepseek_v4(scenario: Scenario, validate: bool = True) -> Sample: thinking = scenario.reasoning is not None - chat_kwargs = {"thinking": True} if thinking else None + chat_kwargs = {"thinking": thinking} if thinking: expected_reasoning: str | None = scenario.reasoning or "" diff --git a/vllm/parser/deepseek_v4.py b/vllm/parser/deepseek_v4.py index c65668708202..9955d938545f 100644 --- a/vllm/parser/deepseek_v4.py +++ b/vllm/parser/deepseek_v4.py @@ -217,10 +217,12 @@ def __init__( **kwargs, ) -> None: chat_kwargs = kwargs.pop("chat_template_kwargs", None) or {} - thinking = ( - bool(chat_kwargs.get("thinking") or chat_kwargs.get("enable_thinking")) - and chat_kwargs.get("reasoning_effort") != "none" + thinking = bool( + chat_kwargs.get("thinking") or chat_kwargs.get("enable_thinking") ) + if "thinking" not in chat_kwargs and "enable_thinking" not in chat_kwargs: + thinking = True + thinking = thinking and chat_kwargs.get("reasoning_effort") != "none" super().__init__( tokenizer, tools, From 63ac04a61e6272b746d82e180e7f6276b3ee1ede Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Mon, 10 Aug 2026 12:54:26 -0700 Subject: [PATCH 07/30] [Kimi-K3] DCP support (#50484) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: girasoley Signed-off-by: Summer Yang Signed-off-by: 云挚 Co-authored-by: Claude Fable 5 Co-authored-by: 云挚 Co-authored-by: foraxe <73625538+foraxe@users.noreply.github.com> Co-authored-by: Woosuk Kwon --- CMakeLists.txt | 5 + .../dcp_utils/dcp_direct_a2a_lse_reduce.cu | 381 ++++++ .../attention/dcp_utils/dcp_direct_common.cuh | 92 ++ .../dcp_utils/dcp_direct_kv_gather.cu | 209 ++++ .../dcp_utils/dcp_direct_q_gather.cu | 197 +++ tests/distributed/test_dcp_a2a.py | 108 +- .../test_dcp_direct_a2a_lse_reduce.py | 1077 +++++++++++++++++ .../test_kimi_linear_context_parallel.py | 128 ++ .../test_kimi_k3_mla_key_concat_kv_cache.py | 246 ++++ tests/v1/attention/test_flashinfer_mla_dcp.py | 80 ++ vllm/envs.py | 13 + .../layers/attention/mla_attention.py | 135 ++- .../layers/attention/sparse_mla_attention.py | 25 +- vllm/models/kimi_k3/nvidia/mla.py | 80 +- .../attention/backends/mla/flashinfer_mla.py | 6 +- .../attention/backends/mla/tokenspeed_mla.py | 3 +- vllm/v1/attention/ops/common.py | 55 +- vllm/v1/attention/ops/dcp_alltoall.py | 31 +- vllm/v1/attention/ops/dcp_utils.py | 740 +++++++++++ 19 files changed, 3529 insertions(+), 82 deletions(-) create mode 100644 csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu create mode 100644 csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh create mode 100644 csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu create mode 100644 csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu create mode 100644 tests/distributed/test_dcp_direct_a2a_lse_reduce.py create mode 100644 tests/distributed/test_kimi_linear_context_parallel.py create mode 100644 tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py create mode 100644 tests/v1/attention/test_flashinfer_mla_dcp.py create mode 100644 vllm/v1/attention/ops/dcp_utils.py diff --git a/CMakeLists.txt b/CMakeLists.txt index bcc0fc53d500..baa4cf7dfae0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -474,6 +474,11 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") endif() if(VLLM_GPU_LANG STREQUAL "CUDA") + list(APPEND VLLM_STABLE_EXT_SRC + "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu" + "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu" + "csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu") + SET(CUTLASS_ENABLE_HEADERS_ONLY ON CACHE BOOL "Enable only the header library") # Set CUTLASS_REVISION. Used for FetchContent. Also fixes some bogus messages when building. diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu new file mode 100644 index 000000000000..bfc8e6c42a11 --- /dev/null +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_a2a_lse_reduce.cu @@ -0,0 +1,381 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Direct symmetric-memory DCP A2A LSE reduction. + +#include +#include +#include +#include +#include + +#include +#include +#include + +#include "dcp_direct_common.cuh" + +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::store_release_system; +using vllm::direct_dcp::wait_for_epoch; + +__device__ __forceinline__ int64_t find_sequence(const int32_t* query_start_loc, + int64_t token_idx, + int64_t num_seqs) { + int64_t left = 0; + int64_t right = num_seqs; + while (left < right) { + int64_t mid = (left + right) / 2; + if (query_start_loc[mid] <= token_idx) { + left = mid + 1; + } else { + right = mid; + } + } + return left - 1; +} + +template +__device__ __forceinline__ float to_float(scalar_t value) { + if constexpr (std::is_same_v) { + return value; + } else if constexpr (std::is_same_v) { + return __half2float(value); + } else { + return __bfloat162float(value); + } +} + +template +__device__ __forceinline__ scalar_t from_float(float value) { + if constexpr (std::is_same_v) { + return __float2half_rn(value); + } else { + return __float2bfloat16_rn(value); + } +} + +// Send each destination's head slice and LSE to its symmetric buffer. Mask +// undefined empty-shard rows at the source so peers give them zero weight. +template +__global__ void dispatch_output_lse_kernel( + const uint4* partial_output, const lse_t* partial_lse, + const int32_t* seq_lens, const int32_t* query_start_loc, + const int64_t* peer_output_ptrs, const int64_t* peer_lse_ptrs, + const int64_t* epoch_ptr, int64_t world_size, int64_t rank, + int64_t num_tokens, int64_t max_num_tokens, int64_t num_seqs, + int64_t heads_per_rank, int64_t head_dim, int64_t output_token_stride, + int64_t lse_token_stride) { + int64_t item = static_cast(blockIdx.x); + int64_t destination_rank = item / num_tokens; + int64_t token_idx = item - destination_rank * num_tokens; + + __shared__ bool empty_kv; + if (threadIdx.x == 0) { + empty_kv = false; + if (seq_lens != nullptr) { + int64_t seq_idx = find_sequence(query_start_loc, token_idx, num_seqs); + empty_kv = seq_lens[seq_idx] == 0; + } + } + __syncthreads(); + + uint32_t epoch = static_cast(epoch_ptr[0]); + int64_t parity = static_cast(epoch & 1u); + int64_t destination_item = + ((parity * world_size + rank) * max_num_tokens + token_idx) * + heads_per_rank; + int64_t source_head = destination_rank * heads_per_rank; + + if (!empty_kv) { + uint4* peer_output = + get_peer_ptr(peer_output_ptrs, destination_rank); + int64_t vectors_per_item = heads_per_rank * head_dim / 8; + int64_t source_vector = + (token_idx * output_token_stride + source_head * head_dim) / 8; + int64_t destination_vector = destination_item * head_dim / 8; + for (int64_t vector_idx = threadIdx.x; vector_idx < vectors_per_item; + vector_idx += blockDim.x) { + peer_output[destination_vector + vector_idx] = + partial_output[source_vector + vector_idx]; + } + } + + float* peer_lse = get_peer_ptr(peer_lse_ptrs, destination_rank); + int64_t source_lse = token_idx * lse_token_stride + source_head; + for (int64_t head_idx = threadIdx.x; head_idx < heads_per_rank; + head_idx += blockDim.x) { + peer_lse[destination_item + head_idx] = + empty_kv ? -CUDART_INF_F : to_float(partial_lse[source_lse + head_idx]); + } +} + +// Publish completion after the stream-ordered payload writes. +__global__ void signal_kernel(const int64_t* peer_signal_ptrs, + const int64_t* epoch_ptr, int64_t world_size, + int64_t rank) { + int64_t destination_rank = + static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (destination_rank >= world_size) { + return; + } + + uint32_t epoch = static_cast(epoch_ptr[0]); + int64_t parity = static_cast(epoch & 1u); + uint32_t* peer_signal = + get_peer_ptr(peer_signal_ptrs, destination_rank); + int64_t signal_item = parity * world_size + rank; + store_release_system(peer_signal + signal_item, epoch); +} + +// Form stable base-2 LSE weights and combine in FP32. +template +__global__ void wait_lse_combine_kernel( + const scalar_t* received_output, const float* received_lse, + const uint32_t* received_signal, const int64_t* epoch_ptr, + scalar_t* combined_output, int64_t world_size, int64_t num_tokens, + int64_t max_num_tokens, int64_t heads_per_rank, int64_t head_dim, + bool is_lse_base_on_e) { + extern __shared__ float weights[]; + + int64_t item = static_cast(blockIdx.x); + int64_t token_idx = item / heads_per_rank; + int64_t head_idx = item - token_idx * heads_per_rank; + uint32_t epoch = static_cast(epoch_ptr[0]); + int64_t parity = static_cast(epoch & 1u); + + if (threadIdx.x == 0) { + float lse_max = -CUDART_INF_F; + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { + int64_t source_item = + ((parity * world_size + source_rank) * max_num_tokens + token_idx) * + heads_per_rank + + head_idx; + int64_t signal_item = parity * world_size + source_rank; + if (!wait_for_epoch(received_signal + signal_item, epoch)) { + printf( + "direct DCP A2A timeout source=%lld token=%lld head=%lld " + "epoch=%u\n", + static_cast(source_rank), + static_cast(token_idx), static_cast(head_idx), + epoch); + asm volatile("trap;"); + } + float value = received_lse[source_item]; + if (isnan(value) || value == CUDART_INF_F) { + value = -CUDART_INF_F; + } + if (is_lse_base_on_e) { + value *= CUDART_L2E_F; + } + weights[source_rank] = value; + lse_max = fmaxf(lse_max, value); + } + if (lse_max == -CUDART_INF_F) { + lse_max = 0.0f; + } + + float lse_sum = 0.0f; + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { + lse_sum += exp2f(weights[source_rank] - lse_max); + } + float inverse_lse_sum = lse_sum > 0.0f ? 1.0f / lse_sum : 0.0f; + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { + weights[source_rank] = + exp2f(weights[source_rank] - lse_max) * inverse_lse_sum; + } + } + __syncthreads(); + + for (int64_t dim_idx = threadIdx.x; dim_idx < head_dim; + dim_idx += blockDim.x) { + float accumulator = 0.0f; + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { + // Empty shards leave stale payloads, so never read zero-weight sources. + float weight = weights[source_rank]; + if (weight == 0.0f) { + continue; + } + int64_t source_item = + ((parity * world_size + source_rank) * max_num_tokens + token_idx) * + heads_per_rank + + head_idx; + accumulator += + to_float(received_output[source_item * head_dim + dim_idx]) * weight; + } + combined_output[item * head_dim + dim_idx] = + from_float(accumulator); + } +} + +void direct_dcp_a2a_lse_reduce( + const torch::stable::Tensor& partial_output, + const torch::stable::Tensor& partial_lse, + const std::optional& seq_lens, + const std::optional& query_start_loc, + const torch::stable::Tensor& peer_output_ptrs, + const torch::stable::Tensor& peer_lse_ptrs, + const torch::stable::Tensor& peer_signal_ptrs, + torch::stable::Tensor& received_output, torch::stable::Tensor& received_lse, + torch::stable::Tensor& received_signal, torch::stable::Tensor& epoch, + torch::stable::Tensor& combined_output, int64_t world_size, int64_t rank, + int64_t max_num_tokens, bool is_lse_base_on_e) { + STD_TORCH_CHECK(partial_output.is_cuda() && partial_lse.is_cuda(), + "partial output and LSE must be CUDA tensors"); + auto output_dtype = partial_output.scalar_type(); + STD_TORCH_CHECK(output_dtype == torch::headeronly::ScalarType::Half || + output_dtype == torch::headeronly::ScalarType::BFloat16, + "direct DCP A2A only supports FP16 and BF16 output"); + auto lse_dtype = partial_lse.scalar_type(); + STD_TORCH_CHECK(lse_dtype == torch::headeronly::ScalarType::Float || + lse_dtype == torch::headeronly::ScalarType::Half || + lse_dtype == torch::headeronly::ScalarType::BFloat16, + "partial LSE must be FP32, FP16, or BF16"); + STD_TORCH_CHECK(partial_output.dim() == 3 && partial_lse.dim() == 2, + "expected output [T,H,D] and LSE [T,H]"); + STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1"); + STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank"); + + int64_t num_tokens = partial_output.size(0); + int64_t total_heads = partial_output.size(1); + int64_t head_dim = partial_output.size(2); + int64_t output_token_stride = partial_output.stride(0); + int64_t lse_token_stride = partial_lse.stride(0); + STD_TORCH_CHECK( + partial_output.stride(2) == 1 && partial_output.stride(1) == head_dim && + output_token_stride >= total_heads * head_dim && + output_token_stride % 8 == 0, + "partial output must have packed heads and an aligned token stride"); + STD_TORCH_CHECK(partial_lse.stride(1) == 1 && lse_token_stride >= total_heads, + "partial LSE must have packed heads"); + STD_TORCH_CHECK(num_tokens > 0 && num_tokens <= max_num_tokens, + "token count exceeds symmetric buffer capacity"); + STD_TORCH_CHECK(total_heads % world_size == 0, + "attention heads must divide evenly across DCP ranks"); + STD_TORCH_CHECK( + partial_lse.size(0) == num_tokens && partial_lse.size(1) == total_heads, + "LSE shape must match attention output"); + STD_TORCH_CHECK(head_dim % 8 == 0, + "head_dim must be divisible by 8 for 16-byte stores"); + int64_t heads_per_rank = total_heads / world_size; + STD_TORCH_CHECK(combined_output.scalar_type() == output_dtype && + combined_output.is_contiguous() && + combined_output.is_cuda(), + "combined output must match the contiguous CUDA input"); + STD_TORCH_CHECK(combined_output.size(0) == num_tokens && + combined_output.size(1) == heads_per_rank && + combined_output.size(2) == head_dim, + "combined output has the wrong shape"); + STD_TORCH_CHECK( + received_output.scalar_type() == output_dtype && + received_lse.scalar_type() == torch::headeronly::ScalarType::Float && + received_signal.scalar_type() == torch::headeronly::ScalarType::Int, + "invalid symmetric staging dtypes"); + STD_TORCH_CHECK( + peer_output_ptrs.scalar_type() == torch::headeronly::ScalarType::Long && + peer_lse_ptrs.scalar_type() == torch::headeronly::ScalarType::Long && + peer_signal_ptrs.scalar_type() == torch::headeronly::ScalarType::Long, + "peer pointer tables must be int64"); + const int32_t* seq_lens_ptr = nullptr; + const int32_t* query_start_loc_ptr = nullptr; + int64_t num_seqs = 0; + STD_TORCH_CHECK(seq_lens.has_value() == query_start_loc.has_value(), + "seq_lens and query_start_loc must be provided together"); + if (seq_lens.has_value() && query_start_loc.has_value()) { + STD_TORCH_CHECK( + seq_lens->is_cuda() && seq_lens->dim() == 1 && + seq_lens->stride(0) == 1 && + seq_lens->scalar_type() == torch::headeronly::ScalarType::Int, + "seq_lens must be a contiguous 1-D int32 CUDA tensor"); + STD_TORCH_CHECK( + query_start_loc->is_cuda() && query_start_loc->dim() == 1 && + query_start_loc->stride(0) == 1 && + query_start_loc->scalar_type() == + torch::headeronly::ScalarType::Int, + "query_start_loc must be a contiguous 1-D int32 CUDA tensor"); + num_seqs = seq_lens->size(0); + STD_TORCH_CHECK(num_seqs > 0 && query_start_loc->size(0) == num_seqs + 1, + "query_start_loc must contain one boundary per sequence"); + seq_lens_ptr = seq_lens->const_data_ptr(); + query_start_loc_ptr = query_start_loc->const_data_ptr(); + } + + const torch::stable::accelerator::DeviceGuard device_guard( + partial_output.get_device_index()); + cudaStream_t stream = get_current_cuda_stream(); + constexpr int kExchangeThreads = 256; + constexpr int kCombineThreads = 128; + + increment_epoch_kernel<<<1, 1, 0, stream>>>( + epoch.mutable_data_ptr()); + check_cuda_launch("direct DCP A2A"); + int64_t dispatch_blocks = world_size * num_tokens; + auto launch_dispatch = [&]() { + dispatch_output_lse_kernel + <<>>( + reinterpret_cast(partial_output.data_ptr()), + reinterpret_cast(partial_lse.data_ptr()), + seq_lens_ptr, query_start_loc_ptr, + peer_output_ptrs.const_data_ptr(), + peer_lse_ptrs.const_data_ptr(), + epoch.const_data_ptr(), world_size, rank, num_tokens, + max_num_tokens, num_seqs, heads_per_rank, head_dim, + output_token_stride, lse_token_stride); + }; + if (lse_dtype == torch::headeronly::ScalarType::Float) { + launch_dispatch.operator()(); + } else if (lse_dtype == torch::headeronly::ScalarType::BFloat16) { + launch_dispatch.operator()<__nv_bfloat16>(); + } else { + launch_dispatch.operator()<__half>(); + } + check_cuda_launch("direct DCP A2A"); + int64_t signal_items = world_size; + int64_t signal_blocks = + (signal_items + kExchangeThreads - 1) / kExchangeThreads; + signal_kernel<<>>( + peer_signal_ptrs.const_data_ptr(), + epoch.const_data_ptr(), world_size, rank); + check_cuda_launch("direct DCP A2A"); + int64_t combine_blocks = num_tokens * heads_per_rank; + size_t shared_memory_bytes = world_size * sizeof(float); + auto launch_combine = [&]() { + wait_lse_combine_kernel + <<>>( + reinterpret_cast(received_output.data_ptr()), + received_lse.const_data_ptr(), + reinterpret_cast( + received_signal.const_data_ptr()), + epoch.const_data_ptr(), + reinterpret_cast(combined_output.mutable_data_ptr()), + world_size, num_tokens, max_num_tokens, heads_per_rank, head_dim, + is_lse_base_on_e); + }; + if (output_dtype == torch::headeronly::ScalarType::BFloat16) { + launch_combine.operator()<__nv_bfloat16>(); + } else { + launch_combine.operator()<__half>(); + } + check_cuda_launch("direct DCP A2A"); +} + +} // namespace + +STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_a2a_ops) { + direct_dcp_a2a_ops.def( + "direct_dcp_a2a_lse_reduce(" + "Tensor partial_output, Tensor partial_lse, Tensor? seq_lens, " + "Tensor? query_start_loc, Tensor peer_output_ptrs, " + "Tensor peer_lse_ptrs, Tensor peer_signal_ptrs, Tensor! received_output, " + "Tensor! received_lse, Tensor! received_signal, Tensor! epoch, " + "Tensor! combined_output, int world_size, int rank, " + "int max_num_tokens, bool is_lse_base_on_e) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, direct_dcp_a2a_ops) { + direct_dcp_a2a_ops.impl("direct_dcp_a2a_lse_reduce", + TORCH_BOX(&direct_dcp_a2a_lse_reduce)); +} diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh new file mode 100644 index 000000000000..968f8c6d7319 --- /dev/null +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_common.cuh @@ -0,0 +1,92 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +#pragma once + +#include + +#include +#include + +#include "../../torch_utils.h" + +namespace vllm::direct_dcp { + +constexpr uint64_t kSpinLimit = 100000000; + +// Advance the invocation ID; its low bit selects one of two staging slots. +static __global__ void increment_epoch_kernel(int64_t* epoch) { + if (blockIdx.x == 0 && threadIdx.x == 0) { + epoch[0] += 1; + } +} + +template +__device__ __forceinline__ T* get_peer_ptr(const int64_t* peer_ptrs, + int64_t peer) { + return reinterpret_cast(static_cast(peer_ptrs[peer])); +} + +// Replicate one 16-byte payload to every symmetric-buffer replica. +__device__ __forceinline__ void multimem_store_16(uint4* mc_ptr, uint4 value) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("multimem.st.relaxed.sys.global.v4.f32 [%0], {%1,%2,%3,%4};" + : + : "l"(mc_ptr), "r"(value.x), "r"(value.y), "r"(value.z), + "r"(value.w) + : "memory"); +#else + asm volatile("trap;"); +#endif +} + +// Publish prior system-scope writes and signal every replica. +__device__ __forceinline__ void multimem_store_release_system(uint32_t* mc_ptr, + uint32_t value) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 900 + asm volatile("multimem.st.release.sys.global.u32 [%0], %1;" + : + : "l"(mc_ptr), "r"(value) + : "memory"); +#else + asm volatile("trap;"); +#endif +} + +__device__ __forceinline__ void store_release_system(uint32_t* ptr, + uint32_t value) { + uint64_t address = reinterpret_cast(ptr); + asm volatile("st.global.release.sys.u32 [%0], %1;" + : + : "l"(address), "r"(value) + : "memory"); +} + +__device__ __forceinline__ uint32_t load_acquire_system(const uint32_t* ptr) { + uint32_t value; + uint64_t address = reinterpret_cast(ptr); + asm volatile("ld.global.acquire.sys.u32 %0, [%1];" + : "=r"(value) + : "l"(address) + : "memory"); + return value; +} + +__device__ __forceinline__ bool wait_for_epoch(const uint32_t* ptr, + uint32_t epoch) { + for (uint64_t spins = 0; spins < kSpinLimit; ++spins) { + if (load_acquire_system(ptr) == epoch) { + return true; + } + } + return false; +} + +inline void check_cuda_launch(const char* operation) { + cudaError_t error = cudaGetLastError(); + STD_TORCH_CHECK(error == cudaSuccess, + std::string(operation) + + " kernel launch failed: " + cudaGetErrorString(error)); +} + +} // namespace vllm::direct_dcp 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 new file mode 100644 index 000000000000..809f759b6acc --- /dev/null +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_kv_gather.cu @@ -0,0 +1,209 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Direct symmetric-memory DCP KV gather. + +#include +#include + +#include + +#include "dcp_direct_common.cuh" + +namespace { + +using vllm::direct_dcp::check_cuda_launch; +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::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. +__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) { + 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; + + 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]); + } + + // Publish all multicast writes before incrementing completion. + __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); + + 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", + 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) { + 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"); + STD_TORCH_CHECK(local_kv.dim() == 2 && local_kv.is_contiguous(), + "local kv must be a contiguous [T,D] tensor"); + STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1"); + STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank"); + + 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(num_tokens * world_size <= max_gathered_tokens, + "gathered kv exceeds symmetric kv-gather buffer capacity"); + + STD_TORCH_CHECK(received_kv.is_cuda() && received_kv.scalar_type() == dtype && + received_kv.is_contiguous() && received_kv.dim() == 3 && + received_kv.size(0) == 2 && + received_kv.size(1) == max_gathered_tokens && + received_kv.size(2) == token_dim, + "received kv has the wrong symmetric buffer layout"); + STD_TORCH_CHECK( + received_signal.is_cuda() && received_signal.is_contiguous() && + received_signal.scalar_type() == ScalarType::Int && + received_signal.dim() == 2 && received_signal.size(0) == 2 && + received_signal.size(1) == world_size, + "received signal has the wrong symmetric buffer layout"); + STD_TORCH_CHECK(completion.is_cuda() && completion.is_contiguous() && + completion.scalar_type() == ScalarType::Int && + completion.numel() == 2, + "completion counter must be a two-element CUDA int32 tensor"); + STD_TORCH_CHECK(epoch.is_cuda() && epoch.is_contiguous() && + 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; + 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 && + slot_stride_bytes % sizeof(uint4) == 0; + STD_TORCH_CHECK( + vectorized, + "direct DCP kv-gather requires 16-byte-aligned pointers and strides"); + STD_TORCH_CHECK(kv_mc_ptr != 0 && signal_mc_ptr != 0, + "direct DCP 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"); + + int64_t item_count = slice_bytes / sizeof(uint4); + 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()), + 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"); +} + +} // namespace + +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, " + "int world_size, int rank, int max_gathered_tokens, " + "int kv_mc_ptr, int signal_mc_ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, direct_dcp_kv_gather_ops) { + direct_dcp_kv_gather_ops.impl("direct_dcp_kv_gather", + TORCH_BOX(&direct_dcp_kv_gather)); +} diff --git a/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu new file mode 100644 index 000000000000..c6d36e2b540d --- /dev/null +++ b/csrc/libtorch_stable/attention/dcp_utils/dcp_direct_q_gather.cu @@ -0,0 +1,197 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// Direct symmetric-memory DCP query gather into the consumer-final buffer. + +#include +#include + +#include + +#include "dcp_direct_common.cuh" + +namespace { + +using vllm::direct_dcp::check_cuda_launch; +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::wait_for_epoch; + +// Multicast each rank's head slice directly into every consumer's final query +// buffer. Reuse is ordered by the downstream DCP output synchronization. +__global__ void direct_dcp_q_gather_multimem_kernel( + const uint4* local_query, uint4* mc_final_query, uint32_t* mc_signal, + const uint32_t* received_signal, int64_t* epoch_ptr, uint32_t* completion, + int64_t world_size, int64_t rank, int64_t num_tokens, + int64_t bytes_per_token, int64_t query_token_stride_bytes, + int64_t destination_token_stride_bytes) { + // The common one-token decode case uses one block. Fold the epoch update + // into that publication kernel to avoid an otherwise separate tiny launch. + if (gridDim.x == 1 && threadIdx.x == 0) { + epoch_ptr[0] += 1; + } + __syncthreads(); + uint32_t epoch = static_cast(epoch_ptr[0]); + int64_t signal_slot = static_cast(epoch & 1u); + + int64_t items_per_token = bytes_per_token / sizeof(uint4); + int64_t source_token_stride = query_token_stride_bytes / sizeof(uint4); + int64_t destination_token_stride = + destination_token_stride_bytes / sizeof(uint4); + int64_t destination_head_offset = rank * items_per_token; + for (int64_t token_idx = blockIdx.x; token_idx < num_tokens; + token_idx += gridDim.x) { + for (int64_t token_item = threadIdx.x; token_item < items_per_token; + token_item += blockDim.x) { + multimem_store_16( + mc_final_query + token_idx * destination_token_stride + + destination_head_offset + token_item, + local_query[token_idx * source_token_stride + token_item]); + } + } + + // Publish all multicast writes before incrementing completion. + __threadfence_system(); + __syncthreads(); + if (threadIdx.x != 0) { + return; + } + + if (gridDim.x > 1) { + uint32_t completed = atomicAdd(completion, 1u); + if (completed + 1u != gridDim.x) { + return; + } + atomicExch(completion, 0u); + } + + multimem_store_release_system(mc_signal + signal_slot * world_size + rank, + epoch); + + for (int64_t source_rank = 0; source_rank < world_size; ++source_rank) { + int64_t signal_item = signal_slot * world_size + source_rank; + if (!wait_for_epoch(received_signal + signal_item, epoch)) { + printf("direct DCP q-gather multimem timeout source=%lld epoch=%u\n", + static_cast(source_rank), epoch); + asm volatile("trap;"); + } + } +} + +void direct_dcp_q_gather(const torch::stable::Tensor& local_query, + torch::stable::Tensor& final_query, + torch::stable::Tensor& received_signal, + torch::stable::Tensor& completion, + torch::stable::Tensor& epoch, int64_t world_size, + int64_t rank, int64_t max_num_tokens, + int64_t padded_num_heads, int64_t query_mc_ptr, + int64_t signal_mc_ptr) { + using torch::headeronly::ScalarType; + + STD_TORCH_CHECK(local_query.is_cuda(), "local query must be a CUDA tensor"); + ScalarType dtype = local_query.scalar_type(); + STD_TORCH_CHECK(local_query.dim() == 3, + "local query must have shape [T,H,D]"); + STD_TORCH_CHECK(world_size > 1, "world_size must be greater than 1"); + STD_TORCH_CHECK(rank >= 0 && rank < world_size, "invalid rank"); + + int64_t num_tokens = local_query.size(0); + int64_t heads_per_rank = local_query.size(1); + int64_t head_dim = local_query.size(2); + int64_t gathered_num_heads = world_size * heads_per_rank; + int64_t element_size = local_query.element_size(); + STD_TORCH_CHECK(num_tokens > 0 && num_tokens <= max_num_tokens, + "token count exceeds symmetric q-gather buffer capacity"); + STD_TORCH_CHECK(heads_per_rank > 0 && head_dim > 0, + "query head dimensions must be positive"); + STD_TORCH_CHECK(padded_num_heads >= gathered_num_heads, + "padded query heads must cover all gathered heads"); + STD_TORCH_CHECK(local_query.stride(2) == 1 && + local_query.stride(1) == head_dim && + local_query.stride(0) >= heads_per_rank * head_dim, + "local query must have packed heads"); + + STD_TORCH_CHECK( + final_query.is_cuda() && final_query.scalar_type() == dtype && + final_query.is_contiguous() && final_query.dim() == 3 && + final_query.size(0) == num_tokens && + final_query.size(1) == gathered_num_heads && + final_query.size(2) == head_dim, + "final query must be contiguous with shape [T,world_size*H,D]"); + STD_TORCH_CHECK( + received_signal.is_cuda() && received_signal.is_contiguous() && + received_signal.scalar_type() == ScalarType::Int && + received_signal.dim() == 2 && received_signal.size(0) == 2 && + received_signal.size(1) == world_size, + "received signal has the wrong symmetric buffer layout"); + STD_TORCH_CHECK(completion.is_cuda() && completion.is_contiguous() && + completion.scalar_type() == ScalarType::Int && + completion.numel() == 1, + "completion counter must be one CUDA int32 tensor"); + STD_TORCH_CHECK(epoch.is_cuda() && epoch.is_contiguous() && + epoch.scalar_type() == ScalarType::Long && + epoch.numel() == 1, + "epoch must be a one-element CUDA int64 tensor"); + int64_t device_index = local_query.get_device_index(); + STD_TORCH_CHECK( + final_query.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 q-gather tensors must be on the same CUDA device"); + + int64_t query_token_stride_bytes = local_query.stride(0) * element_size; + int64_t bytes_per_token = heads_per_rank * head_dim * element_size; + int64_t gathered_token_stride_bytes = + gathered_num_heads * head_dim * element_size; + bool vectorized = + reinterpret_cast(local_query.data_ptr()) % alignof(uint4) == + 0 && + reinterpret_cast(final_query.data_ptr()) % alignof(uint4) == + 0 && + query_token_stride_bytes % sizeof(uint4) == 0 && + bytes_per_token % sizeof(uint4) == 0; + STD_TORCH_CHECK( + vectorized, + "direct DCP q-gather requires 16-byte-aligned pointers and strides"); + STD_TORCH_CHECK(query_mc_ptr != 0 && signal_mc_ptr != 0, + "direct DCP q-gather requires multicast pointers"); + + const torch::stable::accelerator::DeviceGuard device_guard(device_index); + cudaStream_t stream = get_current_cuda_stream(); + constexpr int kThreads = 256; + int64_t blocks = num_tokens < world_size ? num_tokens : world_size; + if (blocks > 1) { + increment_epoch_kernel<<<1, 1, 0, stream>>>( + epoch.mutable_data_ptr()); + check_cuda_launch("direct DCP q-gather epoch"); + } + + direct_dcp_q_gather_multimem_kernel<<>>( + reinterpret_cast(local_query.data_ptr()), + reinterpret_cast(static_cast(query_mc_ptr)), + reinterpret_cast(static_cast(signal_mc_ptr)), + reinterpret_cast( + received_signal.const_data_ptr()), + epoch.mutable_data_ptr(), + reinterpret_cast(completion.mutable_data_ptr()), + world_size, rank, num_tokens, bytes_per_token, query_token_stride_bytes, + gathered_token_stride_bytes); + check_cuda_launch("direct DCP q-gather"); +} + +} // namespace + +STABLE_TORCH_LIBRARY_FRAGMENT(_C, direct_dcp_q_gather_ops) { + direct_dcp_q_gather_ops.def( + "direct_dcp_q_gather(" + "Tensor local_query, Tensor! final_query, Tensor! received_signal, " + "Tensor! completion, Tensor! epoch, " + "int world_size, int rank, int max_num_tokens, " + "int padded_num_heads, int query_mc_ptr, int signal_mc_ptr) -> ()"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, direct_dcp_q_gather_ops) { + direct_dcp_q_gather_ops.impl("direct_dcp_q_gather", + TORCH_BOX(&direct_dcp_q_gather)); +} diff --git a/tests/distributed/test_dcp_a2a.py b/tests/distributed/test_dcp_a2a.py index 5ab0f3de97b5..6dc39a166c16 100644 --- a/tests/distributed/test_dcp_a2a.py +++ b/tests/distributed/test_dcp_a2a.py @@ -222,6 +222,49 @@ def test_dominant_rank(self): assert result.shape == (B, H, D) torch.testing.assert_close(result, outputs[1], atol=1e-5, rtol=1e-5) + def test_empty_shard_ignores_undefined_output(self): + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + + outputs = torch.tensor([[[[float("nan")]]], [[[3.0]]]]) + lses = torch.tensor([[[-float("inf")]], [[0.0]]]) + + result = _lse_weighted_combine(outputs, lses) + + torch.testing.assert_close(result, outputs[1]) + + def test_ag_rs_masks_empty_shard_and_padded_lse(self, monkeypatch): + import vllm.v1.attention.ops.common as common + + class FakeGroup: + world_size = 2 + rank_in_group = 0 + + def all_gather(self, tensor, dim): + assert dim == 0 + return torch.cat((tensor, tensor), dim=dim) + + monkeypatch.setattr( + common, + "correct_attn_out", + lambda output, lses, *args, **kwargs: (output, lses[0]), + ) + output = torch.ones(7, 1, 1) + lse = torch.ones(7, 1) + seq_lens = torch.tensor([0, 2], dtype=torch.int32) + query_start_loc = torch.tensor([0, 1, 5], dtype=torch.int32) + + _, masked_lse = common._cp_lse_common( + output, + lse, + FakeGroup(), + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + + assert torch.isneginf(masked_lse[:1]).all() + torch.testing.assert_close(masked_lse[1:5], torch.ones_like(masked_lse[1:5])) + assert torch.isneginf(masked_lse[5:]).all() + def test_mathematically_correct(self): """Verify mathematical correctness of LSE combination.""" from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine @@ -344,7 +387,7 @@ def test_pack_unpack_combine_matches_reference( world_size, B, h_per_rank, D = 4, 7, 2, 32 H = world_size * h_per_rank cp_attn_out = torch.randn(B, H, D, device=device, dtype=dtype) - cp_attn_lse = torch.randn(B, H, device=device, dtype=torch.float32) + cp_attn_lse = torch.randn(B, H, device=device, dtype=dtype) lse_pack_dim = _dcp_a2a_lse_pack_dim(dtype) send_buffer = torch.empty( (world_size, B, h_per_rank, D + lse_pack_dim), @@ -375,6 +418,67 @@ def test_pack_unpack_combine_matches_reference( else: _assert_packed_a2a_close(actual, expected_out, dtype) + @pytest.mark.skipif( + torch.accelerator.device_count() < 1, reason="CUDA is required." + ) + def test_empty_seq_lens_ignore_undefined_output(self): + from vllm.v1.attention.ops.dcp_alltoall import ( + _dcp_a2a_lse_pack_dim, + _dcp_a2a_pack_send, + _dcp_a2a_unpack_combine, + ) + + device = torch.device("cuda") + world_size, num_tokens, h_per_rank, head_dim = 2, 5, 1, 32 + num_heads = world_size * h_per_rank + output = torch.randn( + num_tokens, + num_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + ) + output[:1] = float("nan") + lse = torch.randn(num_tokens, num_heads, device=device, dtype=output.dtype) + seq_lens = torch.tensor([0, 2], device=device, dtype=torch.int32) + query_start_loc = torch.tensor([0, 1, 5], device=device, dtype=torch.int32) + lse_pack_dim = _dcp_a2a_lse_pack_dim(output.dtype) + send_buffer = torch.empty( + ( + world_size, + num_tokens, + h_per_rank, + head_dim + lse_pack_dim, + ), + device=device, + dtype=output.dtype, + ) + + _dcp_a2a_pack_send( + output, + lse, + send_buffer, + world_size, + h_per_rank, + head_dim, + lse_pack_dim, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + actual_output, actual_lse = _dcp_a2a_unpack_combine( + send_buffer, + head_dim, + lse_pack_dim, + return_lse=True, + is_lse_base_on_e=True, + ) + + torch.testing.assert_close( + actual_output[:1], torch.zeros_like(actual_output[:1]) + ) + assert torch.isneginf(actual_lse[:1]).all() + assert torch.isfinite(actual_output[1:]).all() + def _distributed_packed_a2a_worker(env: dict[str, str]) -> None: update_environment_variables(env) @@ -417,7 +521,7 @@ def _distributed_packed_a2a_worker(env: dict[str, str]) -> None: B, H, device=f"cuda:{local_rank}", - dtype=torch.float32, + dtype=dtype, generator=generator, ) actual = dcp_a2a_lse_reduce( diff --git a/tests/distributed/test_dcp_direct_a2a_lse_reduce.py b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py new file mode 100644 index 000000000000..b05a54825d91 --- /dev/null +++ b/tests/distributed/test_dcp_direct_a2a_lse_reduce.py @@ -0,0 +1,1077 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for direct symmetric-memory DCP collectives.""" + +import functools +from unittest.mock import MagicMock + +import multiprocess as mp +import pytest +import torch +import torch.distributed as dist + +import vllm.v1.attention.ops.dcp_utils as dcp_utils +from vllm.utils.network_utils import get_open_port +from vllm.utils.system_utils import update_environment_variables + +mp.set_start_method("spawn", force=True) + + +def _has_multicast_support() -> bool: + if not torch.cuda.is_available(): + return False + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + return _SymmetricMemory.has_multicast_support(DeviceType.CUDA, 0) + except Exception: + return False + + +def _dtype_from_name(dtype_name: str) -> torch.dtype: + return { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float8_e4m3fn": torch.float8_e4m3fn, + "float32": torch.float32, + }[dtype_name] + + +def _assert_close( + actual: torch.Tensor, + expected: torch.Tensor, + dtype: torch.dtype, +) -> None: + torch.testing.assert_close(actual.float(), expected.float(), rtol=3e-2, atol=3e-2) + + +def _q_gather_reference( + local_query: torch.Tensor, + world_size: int, + padded_num_heads: int | None, +) -> torch.Tensor: + num_tokens, heads_per_rank, head_dim = local_query.shape + gathered = torch.empty( + (world_size * num_tokens, heads_per_rank, head_dim), + dtype=local_query.dtype, + device=local_query.device, + ) + dist.all_gather_into_tensor(gathered, local_query.contiguous()) + expected = ( + gathered.view(world_size, num_tokens, heads_per_rank, head_dim) + .movedim(0, 1) + .reshape(num_tokens, world_size * heads_per_rank, head_dim) + ) + if padded_num_heads is not None: + reserved = expected.new_empty((num_tokens, padded_num_heads, head_dim)) + reserved.resize_(expected.shape) + reserved.copy_(expected) + expected = reserved + return expected + + +def _assert_q_gather_matches_reference( + actual: torch.Tensor, + local_query: torch.Tensor, + world_size: int, + padded_num_heads: int | None, +) -> None: + expected = _q_gather_reference(local_query, world_size, padded_num_heads) + assert actual.shape == expected.shape + assert actual.dtype == expected.dtype + assert actual.is_contiguous() + assert actual.stride() == expected.stride() + assert torch.equal(actual.view(torch.uint8), expected.view(torch.uint8)) + + storage_num_heads = padded_num_heads or world_size * local_query.shape[1] + remaining_storage_bytes = ( + actual.untyped_storage().nbytes() + - actual.storage_offset() * actual.element_size() + ) + required_storage_bytes = ( + local_query.shape[0] + * storage_num_heads + * local_query.shape[2] + * actual.element_size() + ) + assert remaining_storage_bytes >= required_storage_bytes + + +def _distributed_run(fn, world_size: int, extra_env: dict[str, str]) -> None: + port = str(get_open_port()) + processes: list[mp.Process] = [] + for rank in range(world_size): + env = { + "RANK": str(rank), + "LOCAL_RANK": str(rank), + "WORLD_SIZE": str(world_size), + "LOCAL_WORLD_SIZE": str(world_size), + "MASTER_ADDR": "localhost", + "MASTER_PORT": port, + **extra_env, + } + process = mp.Process(target=fn, args=(env,)) + processes.append(process) + process.start() + + for process in processes: + process.join(timeout=120) + + for process in processes: + if process.is_alive(): + process.kill() + process.join() + assert process.exitcode == 0 + + +class _FakeGroupCoordinator: + device_group = None + cpu_group = None + world_size = 4 + + +class _FakeProcessGroup: + def size(self) -> int: + return 4 + + def rank(self) -> int: + return 0 + + +class TestDirectDCPGating: + def test_env_disabled_returns_none(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0") + dcp_utils.get_direct_dcp_a2a_workspace.cache_clear() + workspace = dcp_utils.get_direct_dcp_a2a_workspace( + _FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.bfloat16, 1 + ) + assert workspace is None + + def test_forced_with_unsupported_dtype_raises(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1") + dcp_utils.get_direct_dcp_a2a_workspace.cache_clear() + with pytest.raises(ValueError, match="does not support"): + dcp_utils.get_direct_dcp_a2a_workspace( + _FakeGroupCoordinator(), + torch.device("cpu"), + 16, + 2, + 32, + torch.float32, + 1, + ) + + def test_zero_ubatches_raises(self): + with pytest.raises(ValueError, match="ubatch"): + dcp_utils.DirectDCPA2AWorkspace( + None, torch.device("cpu"), 16, 2, 32, torch.bfloat16, num_ubatches=0 + ) + + def test_auto_with_unsupported_dtype_returns_none(self, monkeypatch): + monkeypatch.delenv("VLLM_USE_DIRECT_DCP_A2A", raising=False) + dcp_utils.get_direct_dcp_a2a_workspace.cache_clear() + workspace = dcp_utils.get_direct_dcp_a2a_workspace( + _FakeGroupCoordinator(), torch.device("cpu"), 16, 2, 32, torch.float32, 1 + ) + assert workspace is None + + def test_q_gather_env_disabled_returns_none(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "0") + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "1") + dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear() + workspace = dcp_utils.get_direct_dcp_q_gather_workspace( + _FakeGroupCoordinator(), + torch.device("cpu"), + 16, + 2, + 32, + torch.bfloat16, + 1, + ) + assert workspace is None + + def test_q_gather_flag_is_independent(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_Q_GATHER", "1") + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0") + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True) + dcp_utils.get_direct_dcp_q_gather_workspace.cache_clear() + workspace = object() + init_workspace = MagicMock(return_value=workspace) + monkeypatch.setattr( + dcp_utils, + "DirectDCPQGatherWorkspace", + init_workspace, + ) + + result = dcp_utils.get_direct_dcp_q_gather_workspace( + _FakeGroupCoordinator(), + torch.device("cpu"), + 16, + 2, + 32, + torch.float32, + 1, + ) + + assert result is workspace + assert init_workspace.call_args.args[5] == torch.float32 + + def test_kv_gather_env_disabled_returns_none(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "0") + 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 + ) + assert workspace is None + + def test_kv_gather_flag_is_independent(self, monkeypatch): + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_KV_GATHER", "1") + monkeypatch.setenv("VLLM_USE_DIRECT_DCP_A2A", "0") + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: True) + 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, torch.bfloat16, 1 + ) + + 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, 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") + monkeypatch.setattr(dcp_utils, "_symm_mem_spans_group", lambda group: False) + factory.cache_clear() + + assert ( + factory( + _FakeGroupCoordinator(), + torch.device("cpu"), + *factory_args, + ) + is None + ) + + 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 + ) + with pytest.raises(ValueError, match="divide evenly"): + dcp_utils.DirectDCPKVGatherWorkspace( + _FakeProcessGroup(), torch.device("cpu"), 63, 576 + ) + with pytest.raises(ValueError, match="16-byte"): + dcp_utils.DirectDCPKVGatherWorkspace( + _FakeProcessGroup(), torch.device("cpu"), 64, 3 + ) + + def test_q_gather_rejects_invalid_workspace_geometry(self): + with pytest.raises(ValueError, match="ubatch"): + dcp_utils.DirectDCPQGatherWorkspace( + None, torch.device("cpu"), 16, 2, 32, num_ubatches=0 + ) + with pytest.raises(ValueError, match="padded heads"): + dcp_utils.DirectDCPQGatherWorkspace( + _FakeProcessGroup(), + torch.device("cpu"), + 16, + 2, + 32, + padded_num_heads=7, + ) + with pytest.raises(ValueError, match="16-byte"): + dcp_utils.DirectDCPQGatherWorkspace( + _FakeProcessGroup(), + torch.device("cpu"), + 16, + 1, + 3, + ) + + +def _manager_config(dcp_comm_backend: str = "a2a"): + config = MagicMock() + config.parallel_config.num_ubatches = 1 + config.parallel_config.dcp_comm_backend = dcp_comm_backend + config.scheduler_config.max_num_batched_tokens = 16 + config.scheduler_config.max_num_seqs = 4 + config.num_speculative_tokens = 0 + config.speculative_config = None + config.compilation_config.max_cudagraph_capture_size = 0 + return config + + +def test_mla_dcp_manager_selects_direct_backends(monkeypatch): + import vllm.v1.attention.ops.dcp_utils as dcp_manager + + group = MagicMock(world_size=2) + monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: group) + direct_a2a = MagicMock() + direct_query = MagicMock() + direct_kv = MagicMock() + monkeypatch.setattr( + dcp_manager, "get_direct_dcp_a2a_workspace", MagicMock(return_value=direct_a2a) + ) + monkeypatch.setattr( + dcp_manager, + "get_direct_dcp_q_gather_workspace", + MagicMock(return_value=direct_query), + ) + monkeypatch.setattr( + dcp_manager, + "get_direct_dcp_kv_gather_workspace", + MagicMock(return_value=direct_kv), + ) + + manager = dcp_manager.MLADCPManager( + vllm_config=_manager_config(), + device=torch.device("cpu"), + num_heads=2, + query_head_dim=8, + output_head_dim=4, + query_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + padded_num_heads=None, + 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) + 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) + manager.combine( + output, + lse, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + direct_a2a.lse_reduce.assert_called_once_with( + output, + lse, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + is_lse_base_on_e=False, + ) + + +def test_mla_dcp_manager_selects_fallback_backends(monkeypatch): + import vllm.v1.attention.ops.dcp_utils as dcp_manager + + group = MagicMock(world_size=2) + gathered_query = torch.empty(1, 4, 8) + group.all_gather.return_value = gathered_query + monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: group) + monkeypatch.setattr( + dcp_manager, "get_direct_dcp_a2a_workspace", MagicMock(return_value=None) + ) + monkeypatch.setattr( + dcp_manager, "get_direct_dcp_q_gather_workspace", MagicMock(return_value=None) + ) + monkeypatch.setattr( + dcp_manager, "get_direct_dcp_kv_gather_workspace", MagicMock(return_value=None) + ) + fallback_combine = MagicMock(return_value=torch.empty(1)) + monkeypatch.setattr(dcp_manager, "dcp_a2a_lse_reduce", fallback_combine) + + manager = dcp_manager.MLADCPManager( + vllm_config=_manager_config(), + device=torch.device("cpu"), + num_heads=2, + query_head_dim=8, + output_head_dim=4, + query_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + padded_num_heads=None, + is_lse_base_on_e=True, + use_pcp=False, + ) + + all_gather = MagicMock() + monkeypatch.setattr(torch.distributed, "all_gather_into_tensor", all_gather) + workspace = torch.empty(96, 8) + manager.init_kv_gather(workspace, 64) + 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) + + query = torch.empty(1, 2, 8) + assert manager.query_gather is not None + assert manager.query_gather(query) is gathered_query + group.all_gather.assert_called_once_with(query, dim=1) + + partial_output, partial_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) + manager.combine( + partial_output, + partial_lse, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + fallback_combine.assert_called_once_with( + partial_output, + partial_lse, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + cp_group=group, + is_lse_base_on_e=True, + ) + + +def test_dcp_workspace_covers_parallel_drafting(): + config = _manager_config() + config.scheduler_config.max_num_batched_tokens = 128 + config.num_speculative_tokens = 3 + config.speculative_config = MagicMock(parallel_drafting=True) + + assert dcp_utils.get_dcp_workspace_max_num_tokens(config) == 28 + + +def test_mla_dcp_manager_selects_pcp_combine(monkeypatch): + import vllm.v1.attention.ops.dcp_utils as dcp_manager + + monkeypatch.setattr(dcp_manager, "get_dcp_group", lambda: MagicMock(world_size=2)) + manager = dcp_manager.MLADCPManager( + vllm_config=_manager_config(dcp_comm_backend="ag_rs"), + device=torch.device("cpu"), + num_heads=2, + query_head_dim=8, + output_head_dim=4, + query_dtype=torch.bfloat16, + output_dtype=torch.bfloat16, + padded_num_heads=None, + is_lse_base_on_e=True, + use_pcp=True, + ) + + assert isinstance(manager.combine, functools.partial) + assert manager.combine.func is dcp_manager.cp_lse_ag_out_ar + assert manager.query_gather is None + + +def test_dcp_chunk_workspace_alignment_covers_interleave(): + from vllm.model_executor.layers.attention.mla_attention import ( + align_mla_chunked_context_workspace_size, + ) + + config = MagicMock() + config.cache_config.block_size = 32 + config.parallel_config.decode_context_parallel_size = 8 + config.parallel_config.cp_kv_cache_interleave_size = 8 + config.scheduler_config.max_num_seqs = 3 + + assert align_mla_chunked_context_workspace_size(config, 100) == 192 + + +def test_sparse_mla_builder_initializes_dcp_manager(monkeypatch): + import vllm.model_executor.layers.attention.sparse_mla_attention as sparse_mla + + monkeypatch.setattr( + sparse_mla.AttentionMetadataBuilder, + "__init__", + lambda *args, **kwargs: None, + ) + monkeypatch.setattr( + sparse_mla, + "get_dcp_group", + lambda: MagicMock(world_size=2), + ) + monkeypatch.setattr( + sparse_mla, + "get_mla_dims", + lambda _: MagicMock(kv_lora_rank=8, qk_rope_head_dim=4), + ) + + manager = object.__new__(dcp_utils.MLADCPManager) + manager.init_kv_gather = MagicMock() + layer = MagicMock(dcp_manager=manager) + config = MagicMock() + config.model_config.dtype = torch.bfloat16 + config.model_config.max_model_len = 64 + config.model_config.hf_config.index_topk = 8 + config.scheduler_config.max_num_batched_tokens = 64 + config.scheduler_config.max_num_seqs = 2 + config.cache_config.block_size = 4 + config.parallel_config.prefill_context_parallel_size = 1 + config.parallel_config.decode_context_parallel_size = 2 + config.parallel_config.cp_kv_cache_interleave_size = 1 + config.compilation_config.static_forward_context = {"layer": layer} + + builder = sparse_mla.SparseMLACommonMetadataBuilder( + MagicMock(), + ["layer"], + config, + torch.device("cpu"), + ) + + assert builder.dcp_manager is manager + manager.init_kv_gather.assert_called_once_with( + builder.chunked_prefill_workspace, + builder.chunked_prefill_workspace_size, + ) + + +def test_sparse_mla_workspace_preserves_non_dcp_size(): + from vllm.model_executor.layers.attention.sparse_mla_attention import ( + SparseMLACommonMetadataBuilder, + ) + + config = MagicMock() + config.model_config.max_model_len = 1 + config.model_config.hf_config.index_topk = 7 + config.scheduler_config.max_num_seqs = 3 + config.cache_config.block_size = 4 + config.parallel_config.decode_context_parallel_size = 1 + + assert ( + SparseMLACommonMetadataBuilder.determine_chunked_prefill_workspace_size(config) + == 21 + ) + + +def _distributed_direct_q_gather_worker(env: dict[str, str]) -> None: + update_environment_variables(env) + local_rank = int(env["LOCAL_RANK"]) + device = torch.device(f"cuda:{local_rank}") + torch.accelerator.set_device_index(local_rank) + dist.init_process_group(backend="nccl") + try: + rank = dist.get_rank() + world_size = dist.get_world_size() + heads_per_rank, head_dim, max_num_tokens = 6, 576, 128 + padded_num_heads = 128 if world_size == 4 else None + active_ubatch = [0] + dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] + for dtype_idx, dtype_name in enumerate( + ("bfloat16", "float8_e4m3fn", "float32") + ): + dtype = _dtype_from_name(dtype_name) + workspace = dcp_utils.DirectDCPQGatherWorkspace( + dist.group.WORLD, + device, + max_num_tokens, + heads_per_rank, + head_dim, + dtype, + num_ubatches=2, + padded_num_heads=padded_num_heads, + ) + + cases = ( + ((1, False), (128, True), (17, True), (5, False)) + if dtype == torch.bfloat16 + else ((17, True),) + ) + for iteration, (num_tokens, noncontiguous) in enumerate(cases): + generator = torch.Generator(device=device) + generator.manual_seed(9000 + rank * 101 + dtype_idx * 1009 + iteration) + source_num_heads = ( + heads_per_rank + 2 if noncontiguous else heads_per_rank + ) + query_storage = torch.randn( + num_tokens, + source_num_heads, + head_dim, + device=device, + dtype=torch.bfloat16, + generator=generator, + ).to(dtype) + local_query = query_storage[:, :heads_per_rank] + assert local_query.is_contiguous() is not noncontiguous + + active_ubatch[0] = iteration % 2 + actual = workspace.gather(local_query) + torch.accelerator.synchronize() + + assert ( + actual.data_ptr() + == workspace.final_query[active_ubatch[0]].data_ptr() + ) + assert workspace.completion[active_ubatch[0]].numel() == 1 + _assert_q_gather_matches_reference( + actual, local_query, world_size, padded_num_heads + ) + + if env.get("TEST_CUDA_GRAPH") != "1" or dtype != torch.bfloat16: + continue + + capture_num_tokens = 17 + capture_storage = torch.empty( + capture_num_tokens, + heads_per_rank + 2, + head_dim, + dtype=dtype, + device=device, + ) + captured_input = capture_storage[:, :heads_per_rank] + assert not captured_input.is_contiguous() + input_pattern = ( + torch.arange(captured_input.numel(), device=device, dtype=torch.int32) + .remainder(31) + .view(captured_input.shape) + .to(dtype) + ) + + active_ubatch[0] = 1 + torch.accelerator.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured_output = workspace.gather(captured_input) + torch.accelerator.synchronize() + dist.barrier() + + eager_input = torch.full( + captured_input.shape, + 96 + rank, + dtype=dtype, + device=device, + ) + eager_output = workspace.gather(eager_input) + torch.accelerator.synchronize() + _assert_q_gather_matches_reference( + eager_output, eager_input, world_size, padded_num_heads + ) + epoch_before_replays = int(workspace.epoch[1].item()) + + for replay in range(2): + captured_input.copy_(input_pattern + rank * 32 + replay) + torch.accelerator.synchronize() + graph.replay() + torch.accelerator.synchronize() + assert int(workspace.epoch[1].item()) == ( + epoch_before_replays + replay + 1 + ) + _assert_q_gather_matches_reference( + captured_output, + captured_input, + world_size, + padded_num_heads, + ) + + interleaved_eager_input = torch.full( + captured_input.shape, + 160 + rank, + dtype=dtype, + device=device, + ) + interleaved_eager_output = workspace.gather(interleaved_eager_input) + torch.accelerator.synchronize() + _assert_q_gather_matches_reference( + interleaved_eager_output, + interleaved_eager_input, + world_size, + padded_num_heads, + ) + + captured_input.copy_(input_pattern + rank * 32 + 2) + torch.accelerator.synchronize() + epoch_before_mixed_replay = int(workspace.epoch[1].item()) + graph.replay() + torch.accelerator.synchronize() + assert int(workspace.epoch[1].item()) == epoch_before_mixed_replay + 1 + _assert_q_gather_matches_reference( + captured_output, captured_input, world_size, padded_num_heads + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "world_size", + [ + 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.", + ), + ), + ], +) +def test_distributed_direct_q_gather_cuda_graph_replay(world_size: int): + _distributed_run( + _distributed_direct_q_gather_worker, + world_size=world_size, + extra_env={"TEST_CUDA_GRAPH": "1"}, + ) + + +def _distributed_direct_kv_gather_worker(env: dict[str, str]) -> None: + update_environment_variables(env) + local_rank = int(env["LOCAL_RANK"]) + device = torch.device(f"cuda:{local_rank}") + torch.accelerator.set_device_index(local_rank) + dist.init_process_group(backend="nccl") + try: + rank = dist.get_rank() + world_size = dist.get_world_size() + token_dim = 576 + max_gathered_tokens = 128 * world_size + active_ubatch = [0] + dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] + + for dtype_idx, dtype_name in enumerate(("bfloat16", "float16")): + dtype = _dtype_from_name(dtype_name) + workspace = dcp_utils.DirectDCPKVGatherWorkspace( + dist.group.WORLD, + device, + max_gathered_tokens, + token_dim, + dtype, + num_ubatches=2, + ) + + # 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) + torch.accelerator.synchronize() + + 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) + ) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "world_size", + [ + 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.", + ), + ), + ], +) +def test_distributed_direct_kv_gather_matches_reference(world_size: int): + _distributed_run( + _distributed_direct_kv_gather_worker, + world_size=world_size, + extra_env={}, + ) + + +def _distributed_direct_a2a_worker(env: dict[str, str]) -> None: + update_environment_variables(env) + local_rank = int(env["LOCAL_RANK"]) + device = torch.device(f"cuda:{local_rank}") + torch.accelerator.set_device_index(local_rank) + dist.init_process_group(backend="nccl") + try: + from vllm.v1.attention.ops.dcp_alltoall import _lse_weighted_combine + + rank = dist.get_rank() + world_size = dist.get_world_size() + dtype = _dtype_from_name(env["TEST_DTYPE"]) + lse_dtype = _dtype_from_name(env["TEST_LSE_DTYPE"]) + is_lse_base_on_e = env["LSE_BASE_E"] == "1" + # Match Kimi-K3's six heads per DCP rank. + heads_per_rank, head_dim, max_num_tokens = 6, 512, 128 + total_heads = world_size * heads_per_rank + active_ubatch = [0] + dcp_utils.dbo_current_ubatch_id = lambda: active_ubatch[0] + workspace = dcp_utils.DirectDCPA2AWorkspace( + dist.group.WORLD, + device, + max_num_tokens, + heads_per_rank, + head_dim, + dtype, + num_ubatches=2, + ) + + def check(num_tokens: int, iteration: int, padded: bool) -> None: + generator = torch.Generator(device=device) + generator.manual_seed(1234 + rank + iteration * world_size) + storage_heads = 128 if padded else total_heads + partial_output_storage = torch.randn( + num_tokens, + storage_heads, + head_dim, + device=device, + dtype=dtype, + generator=generator, + ) + partial_lse_storage = torch.randn( + num_tokens, + storage_heads, + device=device, + dtype=lse_dtype, + generator=generator, + ) + partial_output = partial_output_storage[:, :total_heads, :] + partial_lse = partial_lse_storage[:, :total_heads] + if padded: + assert not partial_output.is_contiguous() + assert not partial_lse.is_contiguous() + active_ubatch[0] = iteration % 2 + actual = workspace.lse_reduce(partial_output, partial_lse, is_lse_base_on_e) + torch.accelerator.synchronize() + + reference_output = partial_output.contiguous() + reference_lse = partial_lse.contiguous() + gathered_output = [ + torch.empty_like(reference_output) for _ in range(world_size) + ] + gathered_lse = [torch.empty_like(reference_lse) for _ in range(world_size)] + dist.all_gather(gathered_output, reference_output) + dist.all_gather(gathered_lse, reference_lse) + outputs = torch.stack( + [ + value[ + :, + rank * heads_per_rank : (rank + 1) * heads_per_rank, + :, + ] + for value in gathered_output + ] + ).float() + lses = torch.stack( + [ + value[:, rank * heads_per_rank : (rank + 1) * heads_per_rank] + for value in gathered_lse + ] + ) + expected = _lse_weighted_combine( + outputs, lses, is_lse_base_on_e=is_lse_base_on_e + ) + _assert_close(actual, expected, dtype) + + def check_empty_shards(query_lens: list[int], iteration: int) -> None: + """Verify empty local shards contribute zero weight.""" + num_seqs = len(query_lens) + num_tokens = sum(query_lens) + generator = torch.Generator(device=device) + generator.manual_seed(777 + rank + iteration * world_size) + partial_output = torch.randn( + num_tokens, + total_heads, + head_dim, + device=device, + dtype=dtype, + generator=generator, + ) + partial_lse = torch.randn( + num_tokens, + total_heads, + device=device, + dtype=lse_dtype, + generator=generator, + ) + + # Cover globally empty, rank-local empty, and non-empty sequences. + def is_empty(seq_idx: int, source_rank: int) -> bool: + return seq_idx == 0 or seq_idx == 1 + source_rank + + seq_lens = torch.tensor( + [ + 0 if is_empty(seq_idx, rank) else seq_idx + 3 + for seq_idx in range(num_seqs) + ], + dtype=torch.int32, + device=device, + ) + query_lens_tensor = torch.tensor( + query_lens, + dtype=torch.int32, + device=device, + ) + query_start_loc = torch.cat( + ( + query_lens_tensor.new_zeros(1), + query_lens_tensor.cumsum(0), + ) + ) + empty_rows = torch.repeat_interleave(seq_lens == 0, query_lens_tensor) + # Model undefined attention rows with NaNs. + partial_output[empty_rows] = float("nan") + partial_lse[empty_rows] = float("nan") + + active_ubatch[0] = iteration % 2 + actual = workspace.lse_reduce( + partial_output, + partial_lse, + is_lse_base_on_e, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) + torch.accelerator.synchronize() + + gathered_output = [ + torch.empty_like(partial_output) for _ in range(world_size) + ] + gathered_lse = [torch.empty_like(partial_lse) for _ in range(world_size)] + dist.all_gather(gathered_output, partial_output.contiguous()) + dist.all_gather(gathered_lse, partial_lse.contiguous()) + head_slice = slice(rank * heads_per_rank, (rank + 1) * heads_per_rank) + outputs = torch.stack( + [value[:, head_slice, :] for value in gathered_output] + ).float() + lses = torch.stack([value[:, head_slice] for value in gathered_lse]) + for source_rank in range(world_size): + source_empty = torch.tensor( + [is_empty(seq_idx, source_rank) for seq_idx in range(num_seqs)], + device=device, + ).repeat_interleave(query_lens_tensor) + outputs[source_rank][source_empty] = 0.0 + lses[source_rank][source_empty] = float("-inf") + expected = _lse_weighted_combine( + outputs, lses, is_lse_base_on_e=is_lse_base_on_e + ) + all_empty = torch.tensor( + [ + all(is_empty(seq_idx, r) for r in range(world_size)) + for seq_idx in range(num_seqs) + ], + device=device, + ).repeat_interleave(query_lens_tensor) + assert torch.equal(actual[all_empty], torch.zeros_like(actual[all_empty])) + assert not torch.isnan(actual.float()).any() + _assert_close(actual, expected, dtype) + + cases = ((1, False), (17, True), (128, True)) + for iteration, (num_tokens, padded) in enumerate(cases): + check(num_tokens, iteration, padded) + check_empty_shards( + query_lens=[1, 3, 2, *([1] * (world_size - 1))], + iteration=len(cases), + ) + generator = torch.Generator(device=device) + generator.manual_seed(4321 + rank) + partial_output_storage = torch.randn( + 128, + 128, + head_dim, + device=device, + dtype=dtype, + generator=generator, + ) + partial_lse_storage = torch.randn( + 128, + 128, + device=device, + dtype=lse_dtype, + generator=generator, + ) + partial_output = partial_output_storage[:, :total_heads, :] + partial_lse = partial_lse_storage[:, :total_heads] + assert not partial_output.is_contiguous() + assert not partial_lse.is_contiguous() + torch.accelerator.synchronize() + active_ubatch[0] = 1 + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + actual = workspace.lse_reduce(partial_output, partial_lse, is_lse_base_on_e) + for _ in range(3): + graph.replay() + torch.accelerator.synchronize() + + reference_output = partial_output.contiguous() + reference_lse = partial_lse.contiguous() + gathered_output = [ + torch.empty_like(reference_output) for _ in range(world_size) + ] + gathered_lse = [torch.empty_like(reference_lse) for _ in range(world_size)] + dist.all_gather(gathered_output, reference_output) + dist.all_gather(gathered_lse, reference_lse) + head_slice = slice(rank * heads_per_rank, (rank + 1) * heads_per_rank) + outputs = torch.stack( + [value[:, head_slice, :] for value in gathered_output] + ).float() + lses = torch.stack([value[:, head_slice] for value in gathered_lse]) + expected = _lse_weighted_combine( + outputs, lses, is_lse_base_on_e=is_lse_base_on_e + ) + _assert_close(actual, expected, dtype) + finally: + dist.destroy_process_group() + + +@pytest.mark.parametrize( + "world_size", + [ + pytest.param( + 2, + marks=pytest.mark.skipif( + torch.accelerator.device_count() < 2, reason="Need at least 2 GPUs." + ), + ), + pytest.param( + 4, + marks=pytest.mark.skipif( + torch.accelerator.device_count() < 4, reason="Need at least 4 GPUs." + ), + ), + ], +) +def test_distributed_direct_a2a_matches_reference(world_size: int): + _distributed_run( + _distributed_direct_a2a_worker, + world_size=world_size, + extra_env={ + "TEST_DTYPE": "bfloat16", + "TEST_LSE_DTYPE": "bfloat16", + "LSE_BASE_E": "0", + }, + ) + + +if __name__ == "__main__": + pytest.main([__file__, "-v"]) diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py new file mode 100644 index 000000000000..d37e19d0bd3b --- /dev/null +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -0,0 +1,128 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import math + +import pytest + +from tests.conftest import VllmRunner +from tests.utils import create_new_process_for_each_test +from vllm import SamplingParams, TokensPrompt + +MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" + + +def _make_tiny_overrides() -> dict: + linear_attn_config = { + "full_attn_layers": [4], + "head_dim": 128, + "kda_layers": [1, 2, 3], + "num_heads": 8, + "short_conv_kernel_size": 4, + "use_full_rank_gate": True, + } + + return { + "bos_token_id": 1, + "eos_token_id": 2, + "pad_token_id": 0, + "vocab_size": 256, + "hidden_size": 256, + "head_dim": 32, + "intermediate_size": 512, + "num_hidden_layers": 4, + "num_attention_heads": 8, + "num_key_value_heads": 8, + "num_experts": None, + "num_experts_per_token": None, + "num_shared_experts": 0, + "q_lora_rank": 128, + "kv_lora_rank": 512, + "qk_nope_head_dim": 128, + "qk_rope_head_dim": 64, + "v_head_dim": 128, + "mla_use_nope": True, + "mla_use_output_gate": True, + "linear_attn_config": linear_attn_config, + "max_position_embeddings": 512, + "model_max_length": 512, + } + + +def _run_tiny_model( + vllm_runner: type[VllmRunner], dcp_size: int +) -> list[tuple[list[int], list[float]]]: + with vllm_runner( + model_name=MODEL, + skip_tokenizer_init=True, + load_format="dummy", + hf_overrides=_make_tiny_overrides(), + tensor_parallel_size=2, + decode_context_parallel_size=dcp_size, + cp_kv_cache_interleave_size=1, + distributed_executor_backend="mp", + dtype="bfloat16", + seed=0, + enforce_eager=True, + max_model_len=512, + max_num_seqs=8, + max_num_batched_tokens=1024, + gpu_memory_utilization=0.25, + enable_chunked_prefill=True, + enable_prefix_caching=False, + ) as runner: + lengths = [1, 32, 33, 129] + prompts = [ + TokensPrompt( + prompt_token_ids=[ + 3 + ((request_idx + token_idx) % 251) for token_idx in range(length) + ] + ) + for request_idx, length in enumerate(lengths) + ] + outputs = runner.llm.generate( + prompts, + SamplingParams( + temperature=0, + max_tokens=8, + seed=0, + ignore_eos=True, + logprobs=20, + ), + use_tqdm=False, + ) + + results = [] + for request_output in outputs: + completion = request_output.outputs[0] + token_ids = list(completion.token_ids) + assert completion.logprobs is not None + selected_logprobs = [ + step_logprobs[token_id].logprob + for token_id, step_logprobs in zip(token_ids, completion.logprobs) + ] + results.append((token_ids, selected_logprobs)) + return results + + +@create_new_process_for_each_test() +@pytest.mark.distributed(num_gpus=2) +def test_kimi_linear_dcp_tiny( + vllm_runner: type[VllmRunner], + num_gpus_available: int, +) -> None: + if num_gpus_available < 2: + pytest.skip("Need at least 2 GPUs") + + baseline = _run_tiny_model(vllm_runner, dcp_size=1) + dcp = _run_tiny_model(vllm_runner, dcp_size=2) + + assert [tokens for tokens, _ in dcp] == [tokens for tokens, _ in baseline] + logprob_drifts = [] + for (_, baseline_logprobs), (_, dcp_logprobs) in zip(baseline, dcp): + for baseline_logprob, dcp_logprob in zip(baseline_logprobs, dcp_logprobs): + assert math.isfinite(baseline_logprob) + assert math.isfinite(dcp_logprob) + logprob_drifts.append(abs(baseline_logprob - dcp_logprob)) + + assert max(logprob_drifts) <= 1e-2 diff --git a/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py new file mode 100644 index 000000000000..e1582b7621a4 --- /dev/null +++ b/tests/kernels/attention/test_kimi_k3_mla_key_concat_kv_cache.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest +import torch + +import vllm._custom_ops as ops +from vllm.models.kimi_k3.nvidia.ops.fused_mla_key_concat_kv_cache import ( + fused_mla_decode_q_concat_kv_cache_insert, + fused_mla_key_concat_ds_mla_insert, + fused_mla_key_concat_kv_cache_insert, + fused_mla_qkv_quant_kv_cache_fp8_insert, +) +from vllm.platforms import current_platform + +NUM_TOKENS = 4 +NUM_HEADS = 2 +KV_LORA_RANK = 512 +NOPE_HEAD_DIM = 128 +ROPE_HEAD_DIM = 64 +QK_HEAD_DIM = NOPE_HEAD_DIM + ROPE_HEAD_DIM +V_HEAD_DIM = 128 +CACHE_ENTRY = KV_LORA_RANK + ROPE_HEAD_DIM +DS_MLA_CACHE_ENTRY = 656 +BLOCK_SIZE = 8 +OWNED_TOKEN_INDICES = [1, 3] + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda(), reason="Kimi-K3 fused MLA ops require CUDA" +) + + +def _inputs() -> dict[str, torch.Tensor]: + torch.manual_seed(0) + + def make(*shape: int) -> torch.Tensor: + return torch.randn(*shape, device="cuda", dtype=torch.bfloat16) + + return { + "q": make(NUM_TOKENS, NUM_HEADS, QK_HEAD_DIM), + "k_nope": make(NUM_TOKENS, NUM_HEADS, NOPE_HEAD_DIM), + "k_pe": make(NUM_TOKENS, ROPE_HEAD_DIM), + "kv_c": make(NUM_TOKENS, KV_LORA_RANK), + "v": make(NUM_TOKENS, NUM_HEADS, V_HEAD_DIM), + "ql_nope": make(NUM_TOKENS, NUM_HEADS, KV_LORA_RANK), + "q_pe": make(NUM_TOKENS, NUM_HEADS, ROPE_HEAD_DIM), + } + + +def _slot_mapping() -> torch.Tensor: + return torch.tensor([-1, 0, -1, 5], device="cuda", dtype=torch.int64) + + +def _owned(tensor: torch.Tensor) -> torch.Tensor: + return tensor[OWNED_TOKEN_INDICES] + + +def _full_key(k_nope: torch.Tensor, k_pe: torch.Tensor) -> torch.Tensor: + shared_pe = k_pe[:, None, :].expand(-1, k_nope.shape[1], -1) + return torch.cat((k_nope, shared_pe), dim=-1) + + +def _latent_query(ql_nope: torch.Tensor, q_pe: torch.Tensor) -> torch.Tensor: + return torch.cat((ql_nope, q_pe), dim=-1) + + +def _make_cache(cache_format: str) -> torch.Tensor: + if cache_format == "bf16": + return torch.full( + (1, BLOCK_SIZE, CACHE_ENTRY), + -7, + device="cuda", + dtype=torch.bfloat16, + ) + if cache_format == "fp8": + return torch.full( + (1, BLOCK_SIZE, CACHE_ENTRY), + 1.0, + device="cuda", + dtype=torch.float32, + ).to(torch.float8_e4m3fn) + assert cache_format == "fp8_ds_mla" + return torch.full( + (1, BLOCK_SIZE, DS_MLA_CACHE_ENTRY), + 165, + device="cuda", + dtype=torch.uint8, + ) + + +def _reference_cache( + cache_format: str, + inputs: dict[str, torch.Tensor], + initial_cache: torch.Tensor, + slots: torch.Tensor, +) -> torch.Tensor: + reference_cache = initial_cache.clone() + valid_slots = _owned(slots) + if cache_format == "fp8_ds_mla": + ops.concat_and_cache_mla( + _owned(inputs["kv_c"]), + _owned(inputs["k_pe"]), + reference_cache, + valid_slots, + cache_format, + torch.ones(1, device="cuda", dtype=torch.float32), + ) + return reference_cache + + latent = torch.cat( + (_owned(inputs["kv_c"]), _owned(inputs["k_pe"])), + dim=-1, + ) + flat_reference_cache = reference_cache.view(-1, CACHE_ENTRY) + for row, slot in zip(latent, valid_slots.tolist()): + flat_reference_cache[slot].copy_(row.to(reference_cache.dtype)) + return reference_cache + + +def _assert_cache_matches_reference( + cache: torch.Tensor, + reference_cache: torch.Tensor, + initial_cache: torch.Tensor, +) -> None: + assert torch.equal(cache, reference_cache) + for slot in (0, 5): + assert not torch.equal(cache[0, slot], initial_cache[0, slot]) + for slot in (1, 2, 3, 4, 6, 7): + assert torch.equal(cache[0, slot], initial_cache[0, slot]) + + +@pytest.mark.parametrize("cache_format", ["bf16", "fp8_ds_mla"]) +def test_prefill_concat_ignores_negative_slots_for_cache(cache_format: str) -> None: + inputs = _inputs() + slots = _slot_mapping() + initial_cache = _make_cache(cache_format) + mixed_cache = initial_cache.clone() + reference_cache = _reference_cache(cache_format, inputs, initial_cache, slots) + + op = ( + fused_mla_key_concat_kv_cache_insert + if cache_format == "bf16" + else fused_mla_key_concat_ds_mla_insert + ) + output = op( + inputs["q"], + inputs["k_nope"], + inputs["k_pe"], + inputs["kv_c"], + mixed_cache, + slots, + ) + + torch.testing.assert_close( + output, _full_key(inputs["k_nope"], inputs["k_pe"]), rtol=0, atol=0 + ) + _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache) + + +def test_fp8_prefill_outputs_ignore_negative_slots_for_cache() -> None: + inputs = _inputs() + slots = _slot_mapping() + initial_cache = _make_cache("fp8") + mixed_cache = initial_cache.clone() + reference_cache = _reference_cache("fp8", inputs, initial_cache, slots) + scale_inv = torch.ones(1, device="cuda", dtype=torch.float32) + + q, k, v = fused_mla_qkv_quant_kv_cache_fp8_insert( + inputs["q"], + inputs["k_nope"], + inputs["k_pe"], + inputs["kv_c"], + inputs["v"], + mixed_cache, + slots, + scale_inv, + scale_inv, + scale_inv, + scale_inv, + ) + + expected = ( + inputs["q"].to(torch.float8_e4m3fn), + _full_key(inputs["k_nope"], inputs["k_pe"]).to(torch.float8_e4m3fn), + inputs["v"].to(torch.float8_e4m3fn), + ) + for actual, reference in zip((q, k, v), expected): + torch.testing.assert_close(actual.float(), reference.float(), rtol=0, atol=0) + _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache) + + +@pytest.mark.parametrize("cache_format", ["bf16", "fp8", "fp8_ds_mla"]) +def test_decode_concat_ignores_negative_slots_for_cache(cache_format: str) -> None: + inputs = _inputs() + slots = _slot_mapping() + initial_cache = _make_cache(cache_format) + mixed_cache = initial_cache.clone() + reference_cache = _reference_cache(cache_format, inputs, initial_cache, slots) + scale_inv = torch.ones(1, device="cuda", dtype=torch.float32) + kwargs = { + "ds_mla": cache_format == "fp8_ds_mla", + "q_scale_inv": scale_inv if cache_format == "fp8" else None, + "cache_scale_inv": scale_inv if cache_format == "fp8" else None, + } + + output = fused_mla_decode_q_concat_kv_cache_insert( + inputs["ql_nope"], + inputs["q_pe"], + inputs["kv_c"], + inputs["k_pe"], + mixed_cache, + slots, + **kwargs, + ) + + expected = _latent_query(inputs["ql_nope"], inputs["q_pe"]) + if cache_format == "fp8": + expected = expected.to(torch.float8_e4m3fn) + torch.testing.assert_close(output.float(), expected.float(), rtol=0, atol=0) + else: + torch.testing.assert_close(output, expected, rtol=0, atol=0) + _assert_cache_matches_reference(mixed_cache, reference_cache, initial_cache) + + +def test_ds_mla_cache_insert_bit_compatible_with_reference() -> None: + """Fused ds_mla insertion must match the reference bit-for-bit.""" + torch.manual_seed(0) + num_tokens, num_blocks = 33, 16 + dt = torch.bfloat16 + kv_c = torch.randn(num_tokens, KV_LORA_RANK, device="cuda", dtype=dt) + k_pe = torch.randn(num_tokens, 1, ROPE_HEAD_DIM, device="cuda", dtype=dt) + k_nope = torch.randn(num_tokens, NUM_HEADS, NOPE_HEAD_DIM, device="cuda", dtype=dt) + q = torch.randn(num_tokens, NUM_HEADS, QK_HEAD_DIM, device="cuda", dtype=dt) + slots = torch.randperm(num_blocks * BLOCK_SIZE, device="cuda")[:num_tokens].long() + slots[torch.rand(num_tokens, device="cuda") < 0.3] = -1 + ref = torch.full( + (num_blocks, BLOCK_SIZE, DS_MLA_CACHE_ENTRY), + 165, + device="cuda", + dtype=torch.uint8, + ) + got = ref.clone() + scale = torch.ones(1, device="cuda", dtype=torch.float32) + ops.concat_and_cache_mla(kv_c, k_pe.squeeze(1), ref, slots, "fp8_ds_mla", scale) + fused_mla_key_concat_ds_mla_insert(q, k_nope, k_pe, kv_c, got, slots) + assert torch.equal(ref, got) diff --git a/tests/v1/attention/test_flashinfer_mla_dcp.py b/tests/v1/attention/test_flashinfer_mla_dcp.py new file mode 100644 index 000000000000..acc38919b802 --- /dev/null +++ b/tests/v1/attention/test_flashinfer_mla_dcp.py @@ -0,0 +1,80 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from unittest.mock import MagicMock + +import pytest +import torch + +from vllm.platforms import current_platform + +requires_flashinfer_mla = pytest.mark.skipif( + not current_platform.has_device_capability(100), + reason="FlashInfer MLA requires compute capability 10 or above.", +) + + +def test_mla_dcp_gathered_query_reserves_backend_head_storage(): + from vllm.v1.attention.ops.dcp_utils import reserve_query_head_storage + + query = torch.randn(3, 24, 576, dtype=torch.bfloat16) + + padded = reserve_query_head_storage(query, 128) + + assert padded.shape == query.shape + assert padded.stride() == query.stride() + assert padded.untyped_storage().nbytes() >= 3 * 128 * 576 * 2 + assert torch.equal(padded, query) + + +@requires_flashinfer_mla +def test_flashinfer_mla_selects_backend_for_gathered_heads(): + from vllm.v1.attention.backends.mla.flashinfer_mla import ( + _select_mla_decode_backend, + ) + + assert _select_mla_decode_backend(6) is None + assert _select_mla_decode_backend(24) == "cute-dsl" + + +@requires_flashinfer_mla +def test_flashinfer_mla_forward_uses_gathered_head_count(monkeypatch): + import vllm.v1.attention.backends.mla.flashinfer_mla as flashinfer_mla + + impl = MagicMock() + impl.bmm1_scale = 1.0 + impl.bmm2_scale = 1.0 + impl.need_to_return_lse_for_decode = True + impl.num_heads = 6 + impl.qk_nope_head_dim = 128 + impl.kv_lora_rank = 512 + impl.qk_rope_head_dim = 64 + impl.kv_cache_dtype = "auto" + + attn_metadata = MagicMock() + attn_metadata.causal = True + attn_metadata.num_decode_tokens = 2 + attn_metadata.num_decodes = 2 + attn_metadata.max_seq_len = 1 + attn_metadata.decode.block_table = torch.zeros(2, 1, dtype=torch.int32) + attn_metadata.decode.seq_lens = torch.ones(2, dtype=torch.int32) + + query = torch.ones(2, 24, 576, dtype=torch.bfloat16) + kv_cache = torch.ones(1, 128, 576, dtype=torch.bfloat16) + kernel = MagicMock( + return_value=( + torch.ones(2, 1, 24, 512, dtype=torch.bfloat16), + torch.ones(2, 24, dtype=torch.float32), + ) + ) + monkeypatch.setattr(flashinfer_mla, "_get_workspace_buffer", MagicMock()) + monkeypatch.setattr(flashinfer_mla, "trtllm_batch_decode_with_kv_cache_mla", kernel) + layer = MagicMock() + + output, lse = flashinfer_mla.FlashInferMLAImpl.forward_mqa( + impl, query, kv_cache, attn_metadata, layer + ) + + assert kernel.call_args.kwargs["backend"] == "cute-dsl" + assert output.shape == (2, 24, 512) + assert lse is not None and lse.shape == (2, 24) diff --git a/vllm/envs.py b/vllm/envs.py index 1460fa1a8946..1b941db97097 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -190,6 +190,9 @@ VLLM_USE_DEEP_GEMM_E8M0: bool = True VLLM_USE_DEEP_GEMM_TMA_ALIGNED_SCALES: bool = True VLLM_DCP_Q_REPLICATE: bool = False + 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 VLLM_DEEP_GEMM_WARMUP: Literal[ "skip", "full", @@ -2057,6 +2060,16 @@ def _resolve_rust_cli_path() -> str | None: "VLLM_USE_SIMPLE_KV_OFFLOAD": lambda: bool( int(os.getenv("VLLM_USE_SIMPLE_KV_OFFLOAD", "0")) ), + # Direct DCP ops default on when applicable; set to 1 to enforce or 0 to disable. + "VLLM_USE_DIRECT_DCP_A2A": lambda: maybe_convert_bool( + os.getenv("VLLM_USE_DIRECT_DCP_A2A") + ), + "VLLM_USE_DIRECT_DCP_Q_GATHER": lambda: maybe_convert_bool( + os.getenv("VLLM_USE_DIRECT_DCP_Q_GATHER") + ), + "VLLM_USE_DIRECT_DCP_KV_GATHER": lambda: maybe_convert_bool( + os.getenv("VLLM_USE_DIRECT_DCP_KV_GATHER") + ), # Whether to enable dual cuda streams for LoRA computation # (used by both BaseLinearLayerWithLoRA and FusedMoEWithLoRA to # overlap the base layer compute with the LoRA fast path). diff --git a/vllm/model_executor/layers/attention/mla_attention.py b/vllm/model_executor/layers/attention/mla_attention.py index 978ee27ed710..b21f96773255 100644 --- a/vllm/model_executor/layers/attention/mla_attention.py +++ b/vllm/model_executor/layers/attention/mla_attention.py @@ -206,6 +206,7 @@ from collections.abc import Callable from dataclasses import dataclass from enum import Enum +from math import lcm from typing import ClassVar, Generic, TypeVar, cast import numpy as np @@ -222,7 +223,6 @@ ModelConfig, VllmConfig, get_current_vllm_config, - get_current_vllm_config_or_none, ) from vllm.config.cache import CacheDType from vllm.distributed.parallel_state import ( @@ -264,7 +264,7 @@ from vllm.model_executor.utils import replace_parameter from vllm.platforms import current_platform from vllm.utils.flashinfer import has_flashinfer -from vllm.utils.math_utils import cdiv, round_down +from vllm.utils.math_utils import cdiv, round_down, round_up from vllm.utils.torch_utils import ( LayerNameType, _encode_layer_name, @@ -292,8 +292,7 @@ get_num_attention_heads_from_layers, split_decodes_and_prefills, ) -from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs -from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce +from vllm.v1.attention.ops.dcp_utils import 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 ( @@ -579,12 +578,27 @@ def __init__( self.use_sparse = use_sparse - _vllm_config = get_current_vllm_config_or_none() - self.dcp_a2a = ( - _vllm_config is not None - and _vllm_config.parallel_config.decode_context_parallel_size > 1 - and _vllm_config.parallel_config.dcp_comm_backend == "a2a" - ) + self.dcp_manager: MLADCPManager | None = None + if self.impl.dcp_world_size > 1: + query_dtype = ( + current_platform.fp8_dtype() + if is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype != "fp8_ds_mla" + and self.impl.supports_quant_query_input + else dtype + ) + self.dcp_manager = MLADCPManager( + vllm_config=vllm_config, + device=next(kv_b_proj.parameters()).device, + num_heads=self.num_heads, + query_head_dim=self.kv_lora_rank + self.qk_rope_head_dim, + output_head_dim=self.kv_lora_rank, + query_dtype=query_dtype, + output_dtype=dtype, + padded_num_heads=self.q_pad_num_heads, + is_lse_base_on_e=self.impl.lse_base_on_e, + use_pcp=self.use_pcp, + ) self.is_aiter_triton_fp8_bmm_enabled = rocm_aiter_ops.is_fp8bmm_enabled() @@ -753,9 +767,6 @@ def forward_impl( return quant_output.fill_(0) return output.fill_(0) - if self.impl.dcp_world_size == -1: - self.impl.dcp_world_size = get_dcp_group().world_size - fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype) num_actual_toks = attn_metadata.num_actual_tokens @@ -914,6 +925,7 @@ def forward_impl( mqa_q = (mqa_ql_nope, mqa_q_pe) # concatenate nope + pe -> (B, N, L + P) (fp8 op above may have fused) if self.impl.dcp_world_size > 1: + assert self.dcp_manager is not None if self.use_pcp: if self.impl.dcp_world_size > self.impl.pcp_world_size: if isinstance(mqa_q, tuple): @@ -924,8 +936,8 @@ def forward_impl( # concatenate mqa_ql_nope and mqa_q_pe -> (B, N, L + P) mqa_q = torch.cat(mqa_q, dim=-1) if not qrep_decode: - # mqa_q do allgather in head dim. - mqa_q = get_dcp_group().all_gather(mqa_q, dim=1) + assert self.dcp_manager.query_gather is not None + mqa_q = self.dcp_manager.query_gather(mqa_q) # call decode attn if not self.impl.is_sparse: @@ -935,27 +947,23 @@ def forward_impl( # correct dcp attn_out with lse. if self.impl.dcp_world_size > 1: assert lse is not None - if self.dcp_a2a: - attn_out = dcp_a2a_lse_reduce( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) - elif self.use_pcp: - attn_out = cp_lse_ag_out_ar( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) - else: - attn_out = cp_lse_ag_out_rs( - attn_out, - lse, - get_dcp_group(), - is_lse_base_on_e=self.impl.lse_base_on_e, - ) + assert self.dcp_manager is not None + seq_lens = ( + attn_metadata.decode.seq_lens + if attn_metadata.decode is not None + else cast(torch.Tensor, attn_metadata.seq_lens)[ # type: ignore[attr-defined] + : attn_metadata.num_decodes + ] + ) + query_start_loc = attn_metadata.query_start_loc[ + : attn_metadata.num_decodes + 1 + ] + attn_out = self.dcp_manager.combine( + attn_out, + lse, + seq_lens=seq_lens, + query_start_loc=query_start_loc, + ) if self.use_pcp: attn_out = finalize_mla_pcp_decode(attn_out, self.num_heads) @@ -1427,6 +1435,7 @@ class ChunkedContextMetadata: chunks: "list[MLACommonPrefillMetadata.ContextChunk]" context_lens_list: list[int] empty_token_slices: list[slice] + dcp_manager: MLADCPManager | None = None block_table: torch.Tensor query_start_loc: torch.Tensor @@ -1684,6 +1693,25 @@ def _flat_int32(values: list[int] | np.ndarray) -> torch.Tensor: return np_to_pinned_tensor(np.asarray(values, dtype=np.int32)) +def align_mla_chunked_context_workspace_size( + vllm_config: VllmConfig, + workspace_size: int, +) -> int: + parallel_config = vllm_config.parallel_config + alignment = vllm_config.cache_config.block_size + if parallel_config.decode_context_parallel_size > 1: + alignment = lcm( + alignment, + parallel_config.decode_context_parallel_size + * parallel_config.cp_kv_cache_interleave_size, + ) + workspace_size = max( + workspace_size, + vllm_config.scheduler_config.max_num_seqs * alignment, + ) + return round_up(workspace_size, alignment) + + def build_mla_chunked_context_metadata( *, context_lens_cpu: torch.Tensor, @@ -1696,6 +1724,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, ) -> "MLACommonPrefillMetadata.ChunkedContextMetadata | None": """Build chunked-context metadata for an MLA prefill. @@ -1714,6 +1743,7 @@ def build_mla_chunked_context_metadata( dcp_world_size: Decode-context-parallel world size (1 if disabled). dcp_local_block_size: Per-rank interleave block size for DCP. dcp_virtual_block_size: ``dcp_local_block_size * dcp_world_size``. + dcp_manager: Shared MLA DCP collective manager. Returns: The chunked-context metadata, or None when no prefill has any context. @@ -1902,6 +1932,7 @@ def padded_rows(rows: int) -> int: chunks=chunks, context_lens_list=context_lens, empty_token_slices=empty_token_slices, + dcp_manager=dcp_manager, ) @@ -1954,7 +1985,10 @@ def determine_chunked_prefill_workspace_size(vllm_config: VllmConfig) -> int: 64 * 1024, ) - return max(chunked_prefill_workspace_size, cache_config.block_size) + return align_mla_chunked_context_workspace_size( + vllm_config, + chunked_prefill_workspace_size, + ) @staticmethod def determine_prefill_query_data_type( @@ -2033,6 +2067,7 @@ def __init__( self.q_data_type = self.determine_prefill_query_data_type( vllm_config, self.model_config.dtype ) + attention_layer = self.compilation_config.static_forward_context[layer_names[0]] try: self.dcp_world_size = get_dcp_group().world_size @@ -2050,6 +2085,7 @@ 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 @@ -2067,6 +2103,12 @@ def __init__( else self.model_config.dtype, device=device, ) + self.dcp_manager = getattr(attention_layer, "dcp_manager", None) + assert isinstance(self.dcp_manager, MLADCPManager) + self.dcp_manager.init_kv_gather( + self.chunked_prefill_workspace, + self.chunked_prefill_workspace_size, + ) else: self.chunked_prefill_workspace = torch.empty( ( @@ -2080,9 +2122,7 @@ def __init__( # Metadata builders are created per ubatch when DBO is enabled. MLA # prefill backends keep the prepared metadata on the backend object, so # each builder needs its own backend instance to avoid cross-ubatch races. - self._prefill_backend = self.compilation_config.static_forward_context[ - layer_names[0] - ].prefill_backend.clone() + self._prefill_backend = attention_layer.prefill_backend.clone() supports_spec_decode = self.query_len_support != QueryLenSupport.SINGLE_ONLY self._init_reorder_batch_threshold( @@ -2223,6 +2263,7 @@ def build( dcp_world_size=self.dcp_world_size, dcp_local_block_size=self.dcp_local_block_size, dcp_virtual_block_size=self.dcp_virtual_block_size, + dcp_manager=self.dcp_manager, ) prefill_metadata = MLACommonPrefillMetadata( @@ -2660,9 +2701,8 @@ def _context_parallel_compute_prefill_context( ] assert toks * dcp_world_size <= cur_allgather_workspace.shape[0] cur_allgather_kvcache = cur_allgather_workspace[: toks * dcp_world_size] - cur_allgather_kvcache.copy_( - get_dcp_group().all_gather(local_gathered_kvcache, dim=0) - ) + dcp_manager = cast(MLADCPManager, chunked_context.dcp_manager) + dcp_manager.kv_gather(cur_allgather_kvcache, local_gathered_kvcache) assert ( cur_allgather_kvcache.shape[-1] == self.kv_lora_rank + self.qk_rope_head_dim @@ -2732,8 +2772,6 @@ def forward_mha( # type: ignore[override] output_scale: torch.Tensor | None = None, ) -> None: assert attn_metadata.prefill is not None - assert self.dcp_world_size != -1 - prefill_metadata = attn_metadata.prefill assert prefill_metadata.prefill_backend is not None use_fp8_prefill = prefill_metadata.q_data_type == current_platform.fp8_dtype() @@ -2877,10 +2915,11 @@ def __init__( and (self.qk_rope_head_dim == 64) ) - self.dcp_world_size: int = -1 - + parallel_config = get_current_vllm_config().parallel_config + # Avoid requiring an initialized DCP group in tests. + self.dcp_world_size: int = parallel_config.decode_context_parallel_size self.cp_kv_cache_interleave_size: int = ( - get_current_vllm_config().parallel_config.cp_kv_cache_interleave_size + parallel_config.cp_kv_cache_interleave_size ) @abstractmethod diff --git a/vllm/model_executor/layers/attention/sparse_mla_attention.py b/vllm/model_executor/layers/attention/sparse_mla_attention.py index ec3a86bbb33e..132bb5f9203b 100644 --- a/vllm/model_executor/layers/attention/sparse_mla_attention.py +++ b/vllm/model_executor/layers/attention/sparse_mla_attention.py @@ -20,6 +20,7 @@ MLACommonMetadata, MLACommonPrefillMetadata, accumulate_mla_context_chunk, + align_mla_chunked_context_workspace_size, build_mla_chunked_context_metadata, get_mla_dims, init_mla_context_partial, @@ -31,6 +32,7 @@ from vllm.v1.attention.backend import AttentionMetadata, AttentionMetadataBuilder from vllm.v1.attention.backends.fa_utils import get_flash_attn_version from vllm.v1.attention.backends.utils import split_decodes_and_prefills +from vllm.v1.attention.ops.dcp_utils import MLADCPManager from vllm.v1.attention.ops.merge_attn_states import merge_attn_states if TYPE_CHECKING: @@ -162,9 +164,19 @@ def __init__( dtype=torch.int32, device=device, ) - layer_prefill_backend = vllm_config.compilation_config.static_forward_context[ + attention_layer = vllm_config.compilation_config.static_forward_context[ layer_names[0] - ].prefill_backend + ] + layer_prefill_backend = attention_layer.prefill_backend + self.dcp_manager: MLADCPManager | None = None + if self.dcp_world_size > 1: + 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, + self.chunked_prefill_workspace_size, + ) self._prefill_backend = ( layer_prefill_backend.clone() if layer_prefill_backend is not None else None ) @@ -184,7 +196,13 @@ def determine_chunked_prefill_workspace_size(vllm_config: "VllmConfig") -> int: 64 * 1024, scheduler_config.max_num_seqs * topk_tokens, ) - return max(workspace_size, cache_config.block_size) + workspace_size = max( + workspace_size, + scheduler_config.max_num_seqs * cache_config.block_size, + ) + if vllm_config.parallel_config.decode_context_parallel_size > 1: + return align_mla_chunked_context_workspace_size(vllm_config, workspace_size) + return workspace_size def _build_req_id_per_token( self, @@ -232,6 +250,7 @@ def _build_chunked_context_fields( dcp_world_size=self.dcp_world_size, dcp_local_block_size=self.dcp_local_block_size, dcp_virtual_block_size=self.dcp_virtual_block_size, + dcp_manager=self.dcp_manager, ) def build( diff --git a/vllm/models/kimi_k3/nvidia/mla.py b/vllm/models/kimi_k3/nvidia/mla.py index c4e3ce90e567..654f83b45243 100644 --- a/vllm/models/kimi_k3/nvidia/mla.py +++ b/vllm/models/kimi_k3/nvidia/mla.py @@ -24,8 +24,8 @@ K3 specifics: optional rotary embedding (disabled for the target model's NoPE layers, enabled for DSpark) and an optional sigmoid output gate (``g_proj``). -Out of scope (extension points, not wired here): context parallelism (DCP/PCP), -sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths. +Out of scope (extension points, not wired here): prefill context parallelism +(PCP), sparse/indexer MLA, and the ROCm/aiter fp8/fp4 BMM fast paths. """ import math @@ -35,7 +35,11 @@ from torch import nn from vllm.compilation.breakable_cudagraph import eager_break_during_capture -from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config +from vllm.config import ( + CacheConfig, + VllmConfig, + get_current_vllm_config, +) from vllm.distributed import get_tensor_model_parallel_world_size from vllm.forward_context import get_forward_context from vllm.logger import init_logger @@ -78,6 +82,7 @@ MLAAttentionImpl, ) from vllm.v1.attention.backends.mla.prefill import get_mla_prefill_backend +from vllm.v1.attention.ops.dcp_utils import 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 KVCacheSpec, MLAAttentionSpec, get_kv_quant_mode @@ -310,10 +315,36 @@ def __init__( vllm_config = get_current_vllm_config() parallel_config = vllm_config.parallel_config - assert ( - parallel_config.decode_context_parallel_size <= 1 - and parallel_config.prefill_context_parallel_size <= 1 - ), "Kimi-K3 MultiHeadLatentAttention does not support context parallelism." + assert parallel_config.prefill_context_parallel_size == 1, ( + "Kimi-K3 MultiHeadLatentAttention does not support prefill context " + "parallelism." + ) + self.dcp_world_size = parallel_config.decode_context_parallel_size + assert self.dcp_world_size <= 1 or self.rotary_emb is None, ( + "Kimi-K3 MultiHeadLatentAttention does not support RoPE with decode " + "context parallelism because gathered queries require gathered " + "positions." + ) + self.dcp_manager: MLADCPManager | None = None + if self.dcp_world_size > 1: + query_dtype = ( + torch.float8_e4m3fn + if is_quantized_kv_cache(self.kv_cache_dtype) + and self.kv_cache_dtype != "fp8_ds_mla" + else dtype + ) + self.dcp_manager = MLADCPManager( + vllm_config=vllm_config, + device=next(self.kv_b_proj.parameters()).device, + num_heads=self.num_local_heads, + query_head_dim=self.head_size, + output_head_dim=self.kv_lora_rank, + query_dtype=query_dtype, + output_dtype=dtype, + padded_num_heads=self.q_pad_num_heads, + is_lse_base_on_e=self.impl.lse_base_on_e, + use_pcp=False, + ) self.prefill_backend = get_mla_prefill_backend(vllm_config)( num_heads=self.num_local_heads, scale=self.scale, @@ -584,9 +615,25 @@ def _attention( cos_sin_cache, slot_mapping[:num_mqa_tokens], ) - latent_out, _lse = self.impl.forward_mqa( # type: ignore[attr-defined] + if self.dcp_world_size > 1: + assert self.dcp_manager is not None + assert self.dcp_manager.query_gather is not None + mqa_q = self.dcp_manager.query_gather(mqa_q) + latent_out, lse = self.impl.forward_mqa( # type: ignore[attr-defined] mqa_q, self._attn_read_kv_cache(), attn_metadata, self ) + if self.dcp_world_size > 1: + assert lse is not None + assert self.dcp_manager is not None + assert attn_metadata.decode is not None + latent_out = self.dcp_manager.combine( + latent_out, + lse, + seq_lens=attn_metadata.decode.seq_lens, + query_start_loc=attn_metadata.query_start_loc[ + : attn_metadata.num_decodes + 1 + ], + ) self._v_up_proj(latent_out, out=attn_out[:num_mqa_tokens]) def _decode_concat_cache( @@ -752,9 +799,20 @@ def _forward_prefill_fused( ) if has_context: - context_output, context_lse = self.impl._compute_prefill_context( # type: ignore[attr-defined] - q, self._attn_read_kv_cache(), attn_metadata, self._k_scale - ) + if self.dcp_world_size > 1: + context_output, context_lse = ( + self.impl._context_parallel_compute_prefill_context( # type: ignore[attr-defined] + q, + self._attn_read_kv_cache(), + attn_metadata, + k_scale=self._k_scale, + dcp_world_size=self.dcp_world_size, + ) + ) + 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 + ) suffix_output, suffix_lse = output_prefill out = out.view(-1, self.num_local_heads, self.v_head_dim) merge_attn_states( diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla.py b/vllm/v1/attention/backends/mla/flashinfer_mla.py index acd3b81816bc..bf92bf692bf0 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla.py @@ -309,9 +309,11 @@ def forward_mqa( return_lse = self.need_to_return_lse_for_decode workspace_buffer = _get_workspace_buffer(return_lse) + # Parallel gathers can change the runtime Q heads from TP-local num_heads. + runtime_num_heads = q.shape[-2] # trtllm-gen rejects MLA head counts it can't tile (e.g. 96); # fall back to cute-dsl for those. - decode_backend = _select_mla_decode_backend(self.num_heads) + decode_backend = _select_mla_decode_backend(runtime_num_heads) extra_kwargs = {} if decode_backend: extra_kwargs["backend"] = decode_backend @@ -326,7 +328,7 @@ def forward_mqa( if self._mla_counter_bytes is None: self._mla_counter_bytes = get_trtllm_gen_multi_ctas_kv_counter_bytes( self._mla_counter_max_batch, - self.num_heads, + runtime_num_heads, get_device_sm_count(q.device), ) extra_kwargs["multi_ctas_kv_counter_buffer"] = ( diff --git a/vllm/v1/attention/backends/mla/tokenspeed_mla.py b/vllm/v1/attention/backends/mla/tokenspeed_mla.py index df4e5c45d13c..8ab81c438f20 100644 --- a/vllm/v1/attention/backends/mla/tokenspeed_mla.py +++ b/vllm/v1/attention/backends/mla/tokenspeed_mla.py @@ -289,8 +289,9 @@ def forward_mqa( self.output_scale = layer._k_scale_float if self._workspace_buffer is None: + # Parallelism can change the runtime query head count. self._workspace_buffer = _get_workspace( - q.device, self.num_heads, self.kv_lora_rank + q.device, q.shape[-2], self.kv_lora_rank ) # vLLM kv_c_and_k_pe_cache is already (num_blocks, block_size, head_size). diff --git a/vllm/v1/attention/ops/common.py b/vllm/v1/attention/ops/common.py index 901d6bb30bd8..cbe24cf87ce4 100644 --- a/vllm/v1/attention/ops/common.py +++ b/vllm/v1/attention/ops/common.py @@ -6,6 +6,34 @@ from vllm.triton_utils import tl, triton +def mask_dcp_empty_shards_( + lse: torch.Tensor, + seq_lens: torch.Tensor | None, + query_start_loc: torch.Tensor | None, +) -> None: + if seq_lens is None and query_start_loc is None: + return + if seq_lens is None or query_start_loc is None: + raise ValueError("seq_lens and query_start_loc must be provided together") + if ( + seq_lens.ndim != 1 + or query_start_loc.ndim != 1 + or query_start_loc.shape[0] != seq_lens.shape[0] + 1 + ): + raise ValueError("query_start_loc must contain one boundary per sequence") + + row_indices = torch.arange( + lse.shape[0], device=lse.device, dtype=query_start_loc.dtype + ) + sequence_indices = torch.searchsorted( + query_start_loc[1:], row_indices, right=True + ).clamp_max(seq_lens.shape[0] - 1) + empty_rows = (row_indices >= query_start_loc[-1]) | ( + seq_lens[sequence_indices] == 0 + ) + lse.masked_fill_(empty_rows[:, None], float("-inf")) + + @triton.jit def _correct_attn_cp_out_kernel( outputs_ptr, @@ -51,7 +79,7 @@ def _correct_attn_cp_out_kernel( ) # calc final lse - lse = tl.load(lses_ptr + lse_offsets) + lse = tl.load(lses_ptr + lse_offsets).to(tl.float32) lse = tl.where((lse != lse) | (lse == float("inf")), -float("inf"), lse) lse_max = tl.max(lse, axis=0) lse_max = tl.where(lse_max == -float("inf"), 0, lse_max) @@ -80,7 +108,7 @@ def _correct_attn_cp_out_kernel( lse_offset = ( lse_idx * lses_stride_N + batch_idx * lses_stride_B + head_idx * lses_stride_H ) - lse_tmp = tl.load(lses_ptr + lse_offset) + lse_tmp = tl.load(lses_ptr + lse_offset).to(tl.float32) lse_finally = lse_tmp - lse lse_finally = tl.where( (lse_finally != lse_finally) | (lse_finally == float("inf")), @@ -185,6 +213,8 @@ def _cp_lse_common( cp_group: GroupCoordinator, ctx: CPTritonContext | None = None, is_lse_base_on_e=True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, ): """ cp_attn_out: [ B, H, D ] @@ -197,6 +227,7 @@ def _cp_lse_common( ctx = CPTritonContext() cp_attn_lse = cp_attn_lse.contiguous() + mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) lses = cp_group.all_gather(cp_attn_lse, dim=0).reshape( (cp_group.world_size,) + cp_attn_lse.shape ) @@ -217,13 +248,21 @@ def cp_lse_ag_out_rs( ctx: CPTritonContext | None = None, return_lse: bool = False, is_lse_base_on_e=True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, ): """ cp_attn_out: [ B, H, D ] cp_attn_lse: [ B, H ] """ out, lse = _cp_lse_common( - cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e + cp_attn_out, + cp_attn_lse, + cp_group, + ctx=ctx, + is_lse_base_on_e=is_lse_base_on_e, + seq_lens=seq_lens, + query_start_loc=query_start_loc, ) out = cp_group.reduce_scatter(out, dim=1) @@ -242,13 +281,21 @@ def cp_lse_ag_out_ar( ctx: CPTritonContext | None = None, return_lse: bool = False, is_lse_base_on_e=True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, ): """ cp_attn_out: [ B, H, D ] cp_attn_lse: [ B, H ] """ out, lse = _cp_lse_common( - cp_attn_out, cp_attn_lse, cp_group, ctx=ctx, is_lse_base_on_e=is_lse_base_on_e + cp_attn_out, + cp_attn_lse, + cp_group, + ctx=ctx, + is_lse_base_on_e=is_lse_base_on_e, + seq_lens=seq_lens, + query_start_loc=query_start_loc, ) out = cp_group.all_reduce(out) diff --git a/vllm/v1/attention/ops/dcp_alltoall.py b/vllm/v1/attention/ops/dcp_alltoall.py index e100dbc79ee2..5939b61a8d35 100644 --- a/vllm/v1/attention/ops/dcp_alltoall.py +++ b/vllm/v1/attention/ops/dcp_alltoall.py @@ -26,6 +26,7 @@ import torch.distributed as dist from vllm.triton_utils import tl, triton +from vllm.v1.attention.ops.common import mask_dcp_empty_shards_ if TYPE_CHECKING: from vllm.distributed.parallel_state import GroupCoordinator @@ -87,7 +88,9 @@ def _lse_weighted_combine( weights = weights / weight_sum.clamp(min=1e-10) # [N, B, H] # Weighted combination: sum over N dimension - result = (outputs * weights.unsqueeze(-1)).sum(dim=0) # [B, H, D] + weights = weights.unsqueeze(-1) + outputs = torch.where(weights == 0, torch.zeros_like(outputs), outputs) + result = (outputs * weights).sum(dim=0) # [B, H, D] if return_lse: if is_lse_base_on_e: @@ -172,7 +175,7 @@ def _dcp_a2a_pack_send_kernel( lse_val = tl.load( lse_ptr + batch_idx * lse_stride_B + src_head_idx * lse_stride_H - ) + ).to(tl.float32) if LSE_PACK_DIM == 1: tl.store( send_ptr + send_base + HEAD_DIM * send_stride_D, @@ -301,10 +304,11 @@ def _dcp_a2a_unpack_combine_kernel( else: weight = tl.exp2(lse_val - global_lse) weight = tl.where(weight != weight, 0.0, weight) - acc += ( - tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to(tl.float32) - * weight + partial = tl.load(recv_ptr + recv_base + d_offsets * recv_stride_D).to( + tl.float32 ) + partial = tl.where(weight == 0.0, 0.0, partial) + acc += partial * weight final_offsets = ( batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D @@ -324,7 +328,10 @@ def _dcp_a2a_pack_send( h_per_rank: int, head_dim: int, lse_pack_dim: int, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, ) -> None: + mask_dcp_empty_shards_(cp_attn_lse, seq_lens, query_start_loc) grid = (cp_attn_out.shape[0], h_per_rank, 1) _dcp_a2a_pack_send_kernel[grid]( cp_attn_out, @@ -396,20 +403,24 @@ def dcp_a2a_lse_reduce( ctx: CPTritonContext | None = None, return_lse: bool = False, is_lse_base_on_e: bool = True, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: """ Combine partial attention outputs across DCP ranks using All-to-All. - The output and fp32 LSE are packed into a single output-dtype buffer, sent + The output and LSE are packed into a single output-dtype buffer, sent with one All-to-All, then unpacked and combined with exact LSE weighting. Args: cp_attn_out: [B, H, D] where B=num_tokens, H=total_heads, D=head_dim - cp_attn_lse: [B, H] log-sum-exp values (fp32) + cp_attn_lse: [B, H] floating-point log-sum-exp values cp_group: GroupCoordinator for DCP communication ctx: CPTritonContext (unused, for signature compatibility) return_lse: If True, also return the combined global LSE is_lse_base_on_e: If True, LSE is base e; if False, base 2 + seq_lens: Local KV lengths. Empty shards contribute zero weight. + query_start_loc: Cumulative query-token offsets for each request. Returns: Combined output [B, H/N, D] (head-scattered) @@ -426,10 +437,6 @@ def dcp_a2a_lse_reduce( if H % world_size != 0: raise ValueError(f"H={H} must be divisible by DCP world size {world_size}.") H_per_rank = H // world_size - # The pack kernel bit-casts the LSE as fp32; some MLA backends return it in - # the activation dtype (bf16/fp16), so enforce the documented fp32 contract. - if cp_attn_lse.dtype != torch.float32: - cp_attn_lse = cp_attn_lse.to(torch.float32) lse_pack_dim = _dcp_a2a_lse_pack_dim(cp_attn_out.dtype) send_buffer, recv_buffer = _dcp_a2a_send_recv_buffers( @@ -446,6 +453,8 @@ def dcp_a2a_lse_reduce( H_per_rank, D, lse_pack_dim, + seq_lens=seq_lens, + query_start_loc=query_start_loc, ) work = dist.all_to_all_single( diff --git a/vllm/v1/attention/ops/dcp_utils.py b/vllm/v1/attention/ops/dcp_utils.py new file mode 100644 index 000000000000..b28d337a1046 --- /dev/null +++ b/vllm/v1/attention/ops/dcp_utils.py @@ -0,0 +1,740 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MLA DCP collective selection and direct symmetric-memory implementations.""" + +from __future__ import annotations + +import functools +from collections.abc import Callable +from typing import TYPE_CHECKING, Any, Protocol + +import torch + +import vllm.envs as envs +from vllm.config import VllmConfig +from vllm.distributed import get_dcp_group +from vllm.distributed.parallel_state import in_the_same_node_as +from vllm.logger import init_logger +from vllm.platforms import current_platform +from vllm.v1.attention.ops.common import cp_lse_ag_out_ar, cp_lse_ag_out_rs +from vllm.v1.attention.ops.dcp_alltoall import dcp_a2a_lse_reduce +from vllm.v1.worker.ubatching import dbo_current_ubatch_id + +logger = init_logger(__name__) + +if TYPE_CHECKING: + from torch.distributed import ProcessGroup + + from vllm.distributed.parallel_state import GroupCoordinator + +try: + import torch.distributed._symmetric_memory as symm_mem + + symm_mem_available = True +except ImportError: + symm_mem = None # type: ignore[assignment] + symm_mem_available = False + + +@functools.cache +def _symm_mem_spans_group(group: GroupCoordinator) -> bool: + """Probe whether the group has NVLS symmetric memory.""" + if not symm_mem_available: + return False + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + device = torch.device("cuda", torch.accelerator.current_device_index()) + if not _SymmetricMemory.has_multicast_support(DeviceType.CUDA, device.index): + return False + probe = symm_mem.empty(8, dtype=torch.uint8, device=device) + probe.zero_() + torch.accelerator.synchronize() + handle = symm_mem.rendezvous(probe, group.device_group.group_name) + spans = handle is not None and handle.multicast_ptr != 0 + except Exception as error: + logger.debug("Direct DCP symmetric-memory probe failed: %s", error) + return False + logger.debug_once( + "Direct DCP symmetric memory across %d ranks: %s", + group.world_size, + "available" if spans else "unavailable", + ) + return spans + + +def _direct_dcp_enabled( + group: GroupCoordinator, + dtype: torch.dtype, + use_direct: bool | None, + supported_dtypes: tuple[torch.dtype, ...] | None = None, +) -> bool: + if use_direct is not None: + return use_direct + return ( + symm_mem_available + and current_platform.is_cuda() + and (supported_dtypes is None or dtype in supported_dtypes) + and ( + all(in_the_same_node_as(group.cpu_group, source_rank=0)) + or _symm_mem_spans_group(group) + ) + ) + + +def _direct_dcp_multicast_enabled( + group: GroupCoordinator, + dtype: torch.dtype, + use_direct: bool | None, + supported_dtypes: tuple[torch.dtype, ...] | None = None, +) -> bool: + return _direct_dcp_enabled( + group, dtype, use_direct, supported_dtypes + ) and _symm_mem_spans_group(group) + + +def get_dcp_workspace_max_num_tokens(vllm_config: VllmConfig) -> int: + scheduler_config = vllm_config.scheduler_config + speculative_config = vllm_config.speculative_config + speculative_tokens = vllm_config.num_speculative_tokens + tokens_per_seq = ( + 1 + + ( + 2 + if speculative_config is not None and speculative_config.parallel_drafting + else 1 + ) + * speculative_tokens + ) + return min( + scheduler_config.max_num_batched_tokens, + max( + scheduler_config.max_num_seqs * tokens_per_seq, + vllm_config.compilation_config.max_cudagraph_capture_size or 0, + ), + ) + + +class _DirectDCPWorkspace: + def __init__( + self, + group: ProcessGroup, + device: torch.device, + num_ubatches: int, + ) -> None: + self.group = group + self.world_size = group.size() + self.rank = group.rank() + self.device = torch.device(device) + self.num_ubatches = num_ubatches + self.epoch = torch.zeros(num_ubatches, dtype=torch.int64, device=self.device) + self._allocations: list[tuple[torch.Tensor, Any, list[torch.Tensor]]] = [] + + def _allocate( + self, shape: tuple[int, ...], dtype: torch.dtype + ) -> tuple[torch.Tensor, torch.Tensor]: + storage = symm_mem.empty(shape, device=self.device, dtype=dtype) + storage.zero_() + torch.accelerator.synchronize() + handle = symm_mem.rendezvous(storage, self.group.group_name) + assert handle is not None, "DCP symmetric memory rendezvous returned None" + handle.barrier() + views = [ + handle.get_buffer(peer, list(shape), dtype, 0) + for peer in range(self.world_size) + ] + self.device = storage.device + peer_ptrs = torch.tensor( + [ + [view[ubatch].data_ptr() for view in views] + for ubatch in range(self.num_ubatches) + ], + dtype=torch.int64, + device=self.device, + ) + self._allocations.append((storage, handle, views)) + return storage, peer_ptrs + + def _multicast_ptrs(self, storage: torch.Tensor) -> list[int]: + disabled = [0] * self.num_ubatches + for allocated, handle, _ in self._allocations: + if allocated is storage: + break + else: + return disabled + try: + from torch._C._autograd import DeviceType + from torch._C._distributed_c10d import _SymmetricMemory + + if not _SymmetricMemory.has_multicast_support( + DeviceType.CUDA, storage.device.index + ): + return disabled + multicast_base = handle.multicast_ptr + except Exception: + return disabled + if not multicast_base: + return disabled + storage_base = storage.data_ptr() + return [ + multicast_base + (storage[ubatch].data_ptr() - storage_base) + for ubatch in range(self.num_ubatches) + ] + + +def reserve_query_head_storage( + query: torch.Tensor, padded_num_heads: int +) -> torch.Tensor: + """Reserve backing storage for fixed-head decode kernels.""" + assert query.ndim == 3 + assert query.shape[1] <= padded_num_heads + padded = query.new_empty((query.shape[0], padded_num_heads, query.shape[2])) + padded.resize_(query.shape) + padded.copy_(query) + return padded + + +_A2A_SUPPORTED_DTYPES = (torch.float16, torch.bfloat16) + + +class DirectDCPA2AWorkspace(_DirectDCPWorkspace): + """Persistent symmetric buffers for direct DCP output exchange.""" + + def __init__( + self, + group: ProcessGroup, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype = torch.bfloat16, + num_ubatches: int = 1, + ) -> None: + if dtype not in _A2A_SUPPORTED_DTYPES: + raise ValueError(f"Direct DCP A2A does not support {dtype}") + if num_ubatches < 1: + raise ValueError( + f"Direct DCP A2A requires at least one ubatch slot, got {num_ubatches}" + ) + super().__init__(group, device, num_ubatches) + self.max_num_tokens = max_num_tokens + self.heads_per_rank = heads_per_rank + self.head_dim = head_dim + + output_shape = ( + num_ubatches, + 2, + self.world_size, + max_num_tokens, + heads_per_rank, + head_dim, + ) + lse_shape = ( + num_ubatches, + 2, + self.world_size, + max_num_tokens, + heads_per_rank, + ) + signal_shape = (num_ubatches, 2, self.world_size) + self.received_output, self.peer_output_ptrs = self._allocate( + output_shape, dtype + ) + self.received_lse, self.peer_lse_ptrs = self._allocate(lse_shape, torch.float32) + self.received_signal, self.peer_signal_ptrs = self._allocate( + signal_shape, torch.int32 + ) + + def lse_reduce( + self, + partial_output: torch.Tensor, + partial_lse: torch.Tensor, + is_lse_base_on_e: bool, + seq_lens: torch.Tensor | None = None, + query_start_loc: torch.Tensor | None = None, + ) -> torch.Tensor: + ubatch = dbo_current_ubatch_id() + num_tokens = partial_output.shape[0] + output = partial_output.new_empty( + (num_tokens, self.heads_per_rank, self.head_dim) + ) + torch.ops._C.direct_dcp_a2a_lse_reduce( + partial_output, + partial_lse, + seq_lens, + query_start_loc, + self.peer_output_ptrs[ubatch], + self.peer_lse_ptrs[ubatch], + self.peer_signal_ptrs[ubatch], + self.received_output[ubatch], + self.received_lse[ubatch], + self.received_signal[ubatch], + self.epoch[ubatch : ubatch + 1], + output, + self.world_size, + self.rank, + self.max_num_tokens, + is_lse_base_on_e, + ) + return output + + +@functools.cache +def get_direct_dcp_a2a_workspace( + group: GroupCoordinator, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype, + num_ubatches: int, +) -> DirectDCPA2AWorkspace | None: + if not _direct_dcp_enabled( + group, dtype, envs.VLLM_USE_DIRECT_DCP_A2A, _A2A_SUPPORTED_DTYPES + ): + return None + return DirectDCPA2AWorkspace( + group.device_group, + device, + max_num_tokens, + heads_per_rank, + head_dim, + dtype, + num_ubatches, + ) + + +def _q_gather_layout_supported( + world_size: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype, + padded_num_heads: int | None, +) -> bool: + element_size = torch.empty((), dtype=dtype).element_size() + gathered_num_heads = world_size * heads_per_rank + storage_num_heads = ( + gathered_num_heads if padded_num_heads is None else padded_num_heads + ) + return ( + heads_per_rank * head_dim * element_size % 16 == 0 + and storage_num_heads * head_dim * element_size % 16 == 0 + ) + + +class DirectDCPQGatherWorkspace(_DirectDCPWorkspace): + """Publish query shards directly into the consumer-final symmetric buffer. + + The final buffer is reusable after the downstream DCP output combine. That + combine orders all ranks after attention has consumed the gathered query. + """ + + def __init__( + self, + group: ProcessGroup, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype = torch.bfloat16, + num_ubatches: int = 1, + padded_num_heads: int | None = None, + ) -> None: + if num_ubatches < 1: + raise ValueError( + "Direct DCP q-gather requires at least one ubatch slot, " + f"got {num_ubatches}" + ) + if max_num_tokens < 1 or heads_per_rank < 1 or head_dim < 1: + raise ValueError( + "Direct DCP q-gather dimensions must be positive, got " + f"T={max_num_tokens}, H={heads_per_rank}, D={head_dim}" + ) + gathered_num_heads = group.size() * heads_per_rank + if not _q_gather_layout_supported( + group.size(), heads_per_rank, head_dim, dtype, padded_num_heads + ): + raise ValueError("Direct DCP q-gather requires 16-byte-aligned query rows.") + super().__init__(group, device, num_ubatches) + if self.world_size <= 1: + raise ValueError("Direct DCP q-gather requires at least two ranks") + self.max_num_tokens = max_num_tokens + self.heads_per_rank = heads_per_rank + self.gathered_num_heads = gathered_num_heads + self.padded_num_heads = ( + self.gathered_num_heads if padded_num_heads is None else padded_num_heads + ) + if self.padded_num_heads < self.gathered_num_heads: + raise ValueError( + "Direct DCP q-gather padded heads must cover gathered heads: " + f"{self.padded_num_heads} < {self.gathered_num_heads}" + ) + self.head_dim = head_dim + + query_shape = ( + num_ubatches, + max_num_tokens, + self.padded_num_heads, + head_dim, + ) + signal_shape = (num_ubatches, 2, self.world_size) + self.final_query, _ = self._allocate(query_shape, dtype) + self.received_signal, _ = self._allocate(signal_shape, torch.int32) + query_multicast_ptrs = self._multicast_ptrs(self.final_query) + signal_multicast_ptrs = self._multicast_ptrs(self.received_signal) + self.multicast_ptrs = list( + zip(query_multicast_ptrs, signal_multicast_ptrs, strict=True) + ) + if not all( + query_ptr and signal_ptr for query_ptr, signal_ptr in self.multicast_ptrs + ): + raise RuntimeError( + "Direct DCP q-gather requires NVLS symmetric-memory multicast." + ) + self.completion = self.received_signal.new_zeros((num_ubatches, 1)) + torch.accelerator.synchronize() + + def gather(self, local_query: torch.Tensor) -> torch.Tensor: + ubatch = dbo_current_ubatch_id() + if not 0 <= ubatch < self.num_ubatches: + raise ValueError( + f"DCP q-gather ubatch {ubatch} exceeds {self.num_ubatches} slots" + ) + if local_query.ndim == 3 and local_query.shape[1] != self.heads_per_rank: + raise ValueError( + f"DCP q-gather expected {self.heads_per_rank} local query heads, " + f"got {local_query.shape[1]}" + ) + + num_tokens = local_query.shape[0] + output = torch.as_strided( + self.final_query[ubatch], + size=(num_tokens, self.gathered_num_heads, self.head_dim), + stride=( + self.gathered_num_heads * self.head_dim, + self.head_dim, + 1, + ), + ) + query_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] + torch.ops._C.direct_dcp_q_gather( + local_query, + output, + self.received_signal[ubatch], + self.completion[ubatch], + self.epoch[ubatch : ubatch + 1], + self.world_size, + self.rank, + self.max_num_tokens, + self.padded_num_heads, + query_multicast_ptr, + signal_multicast_ptr, + ) + return output + + +@functools.cache +def get_direct_dcp_q_gather_workspace( + group: GroupCoordinator, + device: torch.device, + max_num_tokens: int, + heads_per_rank: int, + head_dim: int, + dtype: torch.dtype, + num_ubatches: int, + padded_num_heads: int | None = None, +) -> DirectDCPQGatherWorkspace | None: + if not _direct_dcp_multicast_enabled( + group, dtype, envs.VLLM_USE_DIRECT_DCP_Q_GATHER + ): + return None + if not _q_gather_layout_supported( + group.world_size, heads_per_rank, head_dim, dtype, padded_num_heads + ): + return None + return DirectDCPQGatherWorkspace( + group.device_group, + device, + max_num_tokens, + heads_per_rank, + head_dim, + dtype, + num_ubatches, + padded_num_heads, + ) + + +_KV_GATHER_SUPPORTED_DTYPES = ( + torch.float16, + torch.bfloat16, + torch.float8_e4m3fn, +) + + +def _kv_gather_layout_supported(token_dim: int, dtype: torch.dtype) -> bool: + return token_dim * torch.empty((), dtype=dtype).element_size() % 16 == 0 + + +class DirectDCPKVGatherWorkspace(_DirectDCPWorkspace): + """Persistent symmetric buffers for direct DCP KV gather.""" + + def __init__( + self, + group: ProcessGroup, + device: torch.device, + max_gathered_tokens: int, + token_dim: int, + dtype: torch.dtype = torch.bfloat16, + num_ubatches: int = 1, + ) -> None: + if dtype not in _KV_GATHER_SUPPORTED_DTYPES: + raise ValueError(f"Direct DCP kv-gather does not support {dtype}") + if num_ubatches < 1: + raise ValueError( + "Direct DCP kv-gather requires at least one ubatch slot, " + f"got {num_ubatches}" + ) + if max_gathered_tokens < 1 or token_dim < 1: + raise ValueError( + "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.") + super().__init__(group, device, num_ubatches) + if self.world_size <= 1: + raise ValueError("Direct DCP kv-gather requires at least two ranks") + if max_gathered_tokens % self.world_size != 0: + raise ValueError( + "Direct DCP kv-gather capacity must divide evenly across " + f"ranks: {max_gathered_tokens} % {self.world_size} != 0" + ) + self.max_gathered_tokens = max_gathered_tokens + + 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) + 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.completion = self.received_signal.new_zeros((num_ubatches, 2)) + torch.accelerator.synchronize() + + def gather(self, gathered_kv: torch.Tensor, local_kv: torch.Tensor) -> None: + 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" + ) + kv_multicast_ptr, signal_multicast_ptr = self.multicast_ptrs[ubatch] + torch.ops._C.direct_dcp_kv_gather( + local_kv, + self.received_kv[ubatch], + self.received_signal[ubatch], + self.completion[ubatch], + self.epoch[ubatch : ubatch + 1], + gathered_kv, + self.world_size, + self.rank, + self.max_gathered_tokens, + kv_multicast_ptr, + signal_multicast_ptr, + ) + + +@functools.cache +def get_direct_dcp_kv_gather_workspace( + group: GroupCoordinator, + device: torch.device, + max_gathered_tokens: int, + token_dim: int, + dtype: torch.dtype, + num_ubatches: int, +) -> DirectDCPKVGatherWorkspace | None: + if not _direct_dcp_multicast_enabled( + group, + dtype, + envs.VLLM_USE_DIRECT_DCP_KV_GATHER, + _KV_GATHER_SUPPORTED_DTYPES, + ): + return None + if not _kv_gather_layout_supported(token_dim, dtype): + return None + return DirectDCPKVGatherWorkspace( + group.device_group, + device, + max_gathered_tokens, + token_dim, + dtype, + num_ubatches, + ) + + +class DCPCombine(Protocol): + def __call__( + self, + partial_output: torch.Tensor, + partial_lse: torch.Tensor, + *, + seq_lens: torch.Tensor, + query_start_loc: torch.Tensor, + ) -> torch.Tensor: ... + + +class MLADCPManager: + """Select and own layer-level collective implementations for MLA DCP.""" + + _kv_gather: Callable[[torch.Tensor, torch.Tensor], object] + + def __init__( + self, + vllm_config: VllmConfig, + device: torch.device, + num_heads: int, + query_head_dim: int, + output_head_dim: int, + query_dtype: torch.dtype, + output_dtype: torch.dtype, + padded_num_heads: int | None, + is_lse_base_on_e: bool, + 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) + 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.combine = self._init_combine( + num_heads, + output_head_dim, + output_dtype, + is_lse_base_on_e, + use_pcp, + ) + self.query_gather = ( + None + if use_pcp + else self._init_query_gather( + num_heads, + query_head_dim, + query_dtype, + ) + ) + + def _init_combine( + self, + num_heads: int, + head_dim: int, + dtype: torch.dtype, + is_lse_base_on_e: bool, + use_pcp: bool, + ) -> DCPCombine: + direct_workspace = None + if self.use_a2a: + direct_workspace = get_direct_dcp_a2a_workspace( + self.group, + self.device, + self.max_num_tokens, + num_heads, + head_dim, + dtype, + self.num_ubatches, + ) + if direct_workspace is not None: + logger.info_once("Using direct symmetric-memory DCP A2A for MLA.") + return functools.partial( + direct_workspace.lse_reduce, + is_lse_base_on_e=is_lse_base_on_e, + ) + + combine_fn = ( + dcp_a2a_lse_reduce + if self.use_a2a + else cp_lse_ag_out_ar + if use_pcp + else cp_lse_ag_out_rs + ) + return functools.partial( + combine_fn, + cp_group=self.group, + is_lse_base_on_e=is_lse_base_on_e, + ) + + def _init_query_gather( + self, + num_heads: int, + head_dim: int, + dtype: torch.dtype, + ) -> Callable[[torch.Tensor], torch.Tensor]: + direct_workspace = get_direct_dcp_q_gather_workspace( + self.group, + self.device, + self.max_num_tokens, + num_heads, + head_dim, + dtype, + self.num_ubatches, + self.padded_num_heads, + ) + if direct_workspace is not None: + logger.info_once("Using direct symmetric-memory DCP query gather for MLA.") + return direct_workspace.gather + return self._gather_query + + def _gather_query(self, query: torch.Tensor) -> torch.Tensor: + query = self.group.all_gather(query, dim=1) + 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, + workspace: torch.Tensor, + max_gathered_tokens: int, + ) -> None: + 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 + + direct_workspace = get_direct_dcp_kv_gather_workspace( + self.group, + workspace.device, + max_gathered_tokens, + workspace.shape[1], + workspace.dtype, + self.num_ubatches, + ) + if direct_workspace is not None: + logger.info_once( + "Using direct symmetric-memory DCP chunked-context KV gather for MLA." + ) + self._kv_gather = direct_workspace.gather + else: + self._kv_gather = functools.partial( + torch.distributed.all_gather_into_tensor, + group=self.group.device_group, + ) + + def kv_gather( + self, + gathered_kv: torch.Tensor, + local_kv: torch.Tensor, + ) -> object: + return self._kv_gather(gathered_kv, local_kv) From 635dd6aae656f3b47d174461f34d3de799f2dc94 Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Mon, 10 Aug 2026 13:34:39 -0700 Subject: [PATCH 08/30] [Build][gRPC] Publish protobuf schemas to Buf (#51276) Signed-off-by: Connor Carpenter Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .github/workflows/buf.yml | 56 +++++++++++++++++++++++++++++++++++++++ rust/proto/README.md | 12 +++++++++ rust/proto/buf.md | 8 ++++++ rust/proto/buf.yaml | 9 +++++++ 4 files changed, 85 insertions(+) create mode 100644 .github/workflows/buf.yml create mode 100644 rust/proto/README.md create mode 100644 rust/proto/buf.md create mode 100644 rust/proto/buf.yaml diff --git a/.github/workflows/buf.yml b/.github/workflows/buf.yml new file mode 100644 index 000000000000..e079240afda2 --- /dev/null +++ b/.github/workflows/buf.yml @@ -0,0 +1,56 @@ +name: Buf + +on: + pull_request: + paths: + - "rust/proto/**" + - ".github/workflows/buf.yml" + push: + tags: + - "v*" + schedule: + - cron: "0 8 * * *" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.event_name == 'pull_request' }} + +jobs: + proto: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 + with: + fetch-depth: 0 + - uses: bufbuild/buf-action@8c6a16e16f12ba20b6470afa9c2ba9b5ba8c97c3 # v1.5.0 + with: + version: 1.72.0 + input: rust/proto + token: ${{ github.repository == 'vllm-project/vllm' && github.event_name != 'pull_request' && secrets.BUF_TOKEN || '' }} + format: false + breaking: false + pr_comment: false + push: false + - name: Publish nightly schema + if: github.repository == 'vllm-project/vllm' && (github.event_name == 'schedule' || github.event_name == 'workflow_dispatch') + run: >- + buf push rust/proto + --create + --create-default-label main + --create-visibility public + --label nightly + --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" + - name: Publish release schema + if: github.repository == 'vllm-project/vllm' && github.event_name == 'push' + run: >- + buf push rust/proto + --create + --create-default-label main + --create-visibility public + --label main + --label "${GITHUB_REF_NAME}" + --source-control-url "${{ github.server_url }}/${{ github.repository }}/commit/${{ github.sha }}" diff --git a/rust/proto/README.md b/rust/proto/README.md new file mode 100644 index 000000000000..a9960f3bb8ab --- /dev/null +++ b/rust/proto/README.md @@ -0,0 +1,12 @@ +# gRPC protocol + +This directory is the canonical source for vLLM's gRPC schema. + +The schema is published to `buf.build/vllm-project/vllm`: + +- A daily workflow publishes the latest Git `main` schema to the `nightly` label. +- The workflow can be run manually to retry nightly publication. +- Tags matching `v*` update the Buf `main` label and publish the corresponding release label. +- Buf commits and generated SDK versions are immutable and can be pinned by consumers. + +Repository setup requires a `BUF_TOKEN` GitHub Actions secret with permission to create and push the public Buf module. Register the Prost and Tonic generated SDKs for the `main` and `nightly` labels once so subsequent pushes generate them automatically. diff --git a/rust/proto/buf.md b/rust/proto/buf.md new file mode 100644 index 000000000000..065124feee65 --- /dev/null +++ b/rust/proto/buf.md @@ -0,0 +1,8 @@ +# vLLM gRPC API + +Protocol definitions for the optional gRPC server exposed by the vLLM Rust frontend. + +- `Inference` provides unary and streaming generation. +- `Control` provides discovery and request control. + +The `nightly` label is updated daily from vLLM's `main` branch. The `main` label tracks the latest vLLM release, and version labels match tags such as `v0.27.0`. Pin a Buf commit or generated SDK version for reproducible builds. diff --git a/rust/proto/buf.yaml b/rust/proto/buf.yaml new file mode 100644 index 000000000000..64764f561ef3 --- /dev/null +++ b/rust/proto/buf.yaml @@ -0,0 +1,9 @@ +version: v2 +modules: + - path: . + name: buf.build/vllm-project/vllm +lint: + use: + - MINIMAL + except: + - PACKAGE_DIRECTORY_MATCH From 3c6b2e9c08447a7f839c6aa201d2683ab52d267e Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 10 Aug 2026 14:55:26 -0700 Subject: [PATCH 09/30] [CI] Add /ci cancel command (#51732) Signed-off-by: khluu Co-authored-by: OpenAI Codex --- .github/workflows/new_pr_bot.yml | 2 +- .github/workflows/run-ci-command.yml | 3 +- .github/workflows/scripts/run_ci_command.py | 58 +++++++++++- .../workflows/scripts/test_run_ci_command.py | 94 +++++++++++++++++++ 4 files changed, 151 insertions(+), 6 deletions(-) diff --git a/.github/workflows/new_pr_bot.yml b/.github/workflows/new_pr_bot.yml index a2f09eb8a45e..6e65ccd93a25 100644 --- a/.github/workflows/new_pr_bot.yml +++ b/.github/workflows/new_pr_bot.yml @@ -82,7 +82,7 @@ jobs: '', 'PRs do not trigger a full CI run by default. Reviewers with write access and configured trusted contributors can comment `/ci run` whenever CI signals are needed.', '', - 'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run` or `/ci retry`. New commits do not start CI automatically.', + 'Once the PR is approved or has the `ready` label, the PR author can also use `/ci run`, `/ci retry`, or `/ci cancel`. New commits do not start CI automatically.', '', 'If you have any questions, please reach out to us on Slack at https://slack.vllm.ai.', '', diff --git a/.github/workflows/run-ci-command.yml b/.github/workflows/run-ci-command.yml index 08edc4c2d022..6d33add7477b 100644 --- a/.github/workflows/run-ci-command.yml +++ b/.github/workflows/run-ci-command.yml @@ -20,7 +20,8 @@ jobs: (github.event.comment.body == '/ci run' || github.event.comment.body == '/ci run all' || github.event.comment.body == '/ci run nightly' || - github.event.comment.body == '/ci retry') + github.event.comment.body == '/ci retry' || + github.event.comment.body == '/ci cancel') runs-on: [self-hosted, linux, x64, vllm-runners] steps: - uses: actions/checkout@9c091bb21b7c1c1d1991bb908d89e4e9dddfe3e0 # v7.0.0 diff --git a/.github/workflows/scripts/run_ci_command.py b/.github/workflows/scripts/run_ci_command.py index bc9d7295e858..0dfa90cf3f48 100644 --- a/.github/workflows/scripts/run_ci_command.py +++ b/.github/workflows/scripts/run_ci_command.py @@ -16,6 +16,7 @@ COMMAND_RUN_CI_ALL = "/ci run all" COMMAND_RUN_CI_NIGHTLY = "/ci run nightly" COMMAND_RETRY_FAILED = "/ci retry" +COMMAND_CANCEL_CI = "/ci cancel" RUN_CI_COMMAND_ENV = { COMMAND_RUN_CI: {}, COMMAND_RUN_CI_ALL: {"RUN_ALL": "1"}, @@ -35,6 +36,7 @@ "waiting_failed", } RETRY_STATES = "failed,timed_out,expired" +CANCELABLE_BUILD_STATES = ("scheduled", "running", "failing") SETUP_STEP_KEYS = { "ensure-ci-base-amd", "pre-commit", @@ -335,7 +337,9 @@ def list_builds( self, commit: str | None, *, + branch: str | None = None, metadata: tuple[str, str] | None = None, + states: Sequence[str] = (), ) -> list[dict[str, Any]]: query = [ ("exclude_jobs", "true"), @@ -344,9 +348,12 @@ def list_builds( ] if commit: query.append(("commit", commit)) + if branch: + query.append(("branch", branch)) if metadata: key, value = metadata query.append((f"meta_data[{key}]", value)) + query.extend(("state[]", state) for state in states) response = self._request(query=query) if not isinstance(response, list): raise ApiError(None, "Buildkite API returned an invalid build list.") @@ -367,6 +374,10 @@ def retry_failed_jobs( path=f"/{number}/retry_failed_jobs", ) + def cancel_build(self, build_number: int) -> dict[str, Any]: + number = urllib.parse.quote(str(build_number), safe="") + return self._request(method="PUT", path=f"/{number}/cancel") + def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: number = urllib.parse.quote(str(build_number), safe="") query = [ @@ -397,7 +408,7 @@ def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: def parse_command(body: str) -> str | None: - if body in {*RUN_CI_COMMAND_ENV, COMMAND_RETRY_FAILED}: + if body in {*RUN_CI_COMMAND_ENV, COMMAND_RETRY_FAILED, COMMAND_CANCEL_CI}: return body return None @@ -432,7 +443,7 @@ def authorize( if actor.casefold() != pr["user"]["login"].casefold(): return ( False, - "Only reviewers with write access can run CI before it is " + "Only reviewers with write access can use CI commands before CI is " "delegated to the PR author.", ) if pr["draft"]: @@ -650,7 +661,9 @@ def notify_authorized( "- `/ci retry` retries failed jobs in the CI build for the current " "PR head. If the current head has no CI build, it starts a new CI " "build for the current head containing only jobs that failed in " - "the latest earlier CI build for this PR.\n\n" + "the latest earlier CI build for this PR.\n" + "- `/ci cancel` cancels scheduled or running CI builds for this PR " + "branch.\n\n" f"{CI_AUTHORIZED_COMMENT_MARKER}" ), ) @@ -837,6 +850,38 @@ def handle_retry_failed( ) +def handle_cancel_ci( + *, + buildkite: BuildkiteClient, + pr: Mapping[str, Any], +) -> str: + branch = pr["head"]["ref"] + builds = buildkite.list_builds( + None, + branch=branch, + states=CANCELABLE_BUILD_STATES, + ) + cancelable_builds = [ + build + for build in builds + if build.get("branch") == branch + and is_build_for_pr(build, pr["number"]) + and build.get("state") in CANCELABLE_BUILD_STATES + ] + if not cancelable_builds: + return f"No cancelable CI build is running for branch `{branch}`." + + for build in cancelable_builds: + buildkite.cancel_build(build["number"]) + + links = ", ".join( + f"[#{build['number']}]({build['web_url']})" for build in cancelable_builds + ) + count = len(cancelable_builds) + noun = "build" if count == 1 else "builds" + return f"Requested cancellation of {count} CI {noun} for `{branch}`: {links}." + + def run( event: Mapping[str, Any], github: GitHubClient, @@ -904,7 +949,7 @@ def run( github=github, pr=pr, ) - else: + elif command == COMMAND_RETRY_FAILED: message = handle_retry_failed( actor=actor, buildkite=buildkite, @@ -912,6 +957,11 @@ def run( github=github, pr=pr, ) + else: + message = handle_cancel_ci( + buildkite=buildkite, + pr=pr, + ) add_reaction_safely(github, comment_id, "rocket") github.add_comment(issue_number, f"✅ {message}") except Exception: diff --git a/.github/workflows/scripts/test_run_ci_command.py b/.github/workflows/scripts/test_run_ci_command.py index aba4d1dd86e1..f20b22525fd6 100644 --- a/.github/workflows/scripts/test_run_ci_command.py +++ b/.github/workflows/scripts/test_run_ci_command.py @@ -9,7 +9,9 @@ from unittest.mock import patch from run_ci_command import ( + CANCELABLE_BUILD_STATES, CI_AUTHORIZED_COMMENT_MARKER, + COMMAND_CANCEL_CI, COMMAND_RETRY_FAILED, COMMAND_RUN_CI, COMMAND_RUN_CI_ALL, @@ -136,12 +138,15 @@ def __init__( self.job_list_calls: list[int] = [] self.list_calls: list[tuple[str | None, tuple[str, str] | None]] = [] self.retry_calls: list[tuple[int, str]] = [] + self.cancel_calls: list[int] = [] def list_builds( self, commit: str | None, *, + branch: str | None = None, metadata: tuple[str, str] | None = None, + states: tuple[str, ...] = (), ) -> list[dict[str, Any]]: self.list_calls.append((commit, metadata)) return self.build_lists.pop(0) @@ -161,6 +166,10 @@ def retry_failed_jobs( self.retry_calls.append((build_number, states)) return {"retried_jobs_count": 3} + def cancel_build(self, build_number: int) -> dict[str, Any]: + self.cancel_calls.append(build_number) + return {"number": build_number, "state": "canceling"} + def list_failed_jobs(self, build_number: int) -> list[dict[str, Any]]: self.job_list_calls.append(build_number) return self.failed_job_lists.pop(0) @@ -292,8 +301,10 @@ def test_only_exact_ci_commands_are_accepted(self) -> None: parse_command(COMMAND_RETRY_FAILED), COMMAND_RETRY_FAILED, ) + self.assertEqual(parse_command(COMMAND_CANCEL_CI), COMMAND_CANCEL_CI) self.assertIsNone(parse_command("/ci run please")) self.assertIsNone(parse_command("/ci run all please")) + self.assertIsNone(parse_command("/ci cancel please")) self.assertIsNone(parse_command(" /ci run")) def test_write_access_authorizes_reviewers_and_authors(self) -> None: @@ -531,6 +542,7 @@ def test_ready_label_notifies_author_once(self) -> None: self.assertTrue(comment.startswith("✅ @author")) self.assertIn("`/ci run` starts a CI build", comment) self.assertIn("`/ci retry` retries failed jobs", comment) + self.assertIn("`/ci cancel` cancels scheduled or running", comment) self.assertIn("CI build for the current PR head", comment) self.assertIn("only jobs that failed in the latest earlier CI build", comment) self.assertNotIn(COMMAND_RUN_CI_ALL, comment) @@ -854,6 +866,82 @@ def test_buildkite_retry_uses_retry_failed_jobs_endpoint(self) -> None: self.assertTrue(call["url"].endswith("/123/retry_failed_jobs")) self.assertEqual(call["body"], {"states": RETRY_STATES}) + def test_ci_cancel_cancels_active_builds_for_pr_branch(self) -> None: + github = FakeGitHub() + buildkite = FakeBuildkite( + [ + [ + { + "branch": "feature", + "number": 123, + "pull_request": {"id": 42}, + "state": "running", + "web_url": "https://buildkite.example/builds/123", + }, + { + "branch": "feature", + "number": 124, + "meta_data": {"github-pr-number": "42"}, + "state": "failing", + "web_url": "https://buildkite.example/builds/124", + }, + { + "branch": "feature", + "number": 125, + "pull_request": {"id": 43}, + "state": "running", + "web_url": "https://buildkite.example/builds/125", + }, + { + "branch": "other-branch", + "number": 126, + "pull_request": {"id": 42}, + "state": "running", + "web_url": "https://buildkite.example/builds/126", + }, + { + "branch": "feature", + "number": 127, + "pull_request": {"id": 42}, + "state": "passed", + "web_url": "https://buildkite.example/builds/127", + }, + ] + ] + ) + + run(make_event(COMMAND_CANCEL_CI), github, buildkite) + + self.assertEqual(buildkite.cancel_calls, [123, 124]) + self.assertIn("Requested cancellation of 2 CI builds", github.comments[0]) + self.assertIn("#123", github.comments[0]) + self.assertIn("#124", github.comments[0]) + + def test_ci_cancel_is_a_noop_without_active_builds(self) -> None: + github = FakeGitHub() + buildkite = FakeBuildkite([[]]) + + run(make_event(COMMAND_CANCEL_CI), github, buildkite) + + self.assertEqual(buildkite.cancel_calls, []) + self.assertIn("No cancelable CI build is running", github.comments[0]) + + def test_buildkite_cancel_uses_cancel_build_endpoint(self) -> None: + transport = FakeTransport({"number": 123, "state": "canceling"}) + client = BuildkiteClient( + "secret", + "vllm", + "ci", + transport=transport, + ) + + client.cancel_build(123) + + call = transport.calls[0] + self.assertEqual(call["method"], "PUT") + self.assertTrue(call["url"].endswith("/123/cancel")) + self.assertIsNone(call["body"]) + def test_buildkite_list_builds_allows_query_on_builds_endpoint(self) -> None: transport = FakeTransport([]) client = BuildkiteClient( @@ -865,14 +953,20 @@ def test_buildkite_list_builds_allows_query_on_builds_endpoint(self) -> None: builds = client.list_builds( "current-commit", + branch="feature", metadata=("github-pr-number", "42"), + states=CANCELABLE_BUILD_STATES, ) self.assertEqual(builds, []) url = transport.calls[0]["url"] self.assertIn("?exclude_jobs=true", url) self.assertIn("commit=current-commit", url) + self.assertIn("branch=feature", url) self.assertIn("meta_data%5Bgithub-pr-number%5D=42", url) + self.assertIn("state%5B%5D=scheduled", url) + self.assertIn("state%5B%5D=running", url) + self.assertIn("state%5B%5D=failing", url) def test_buildkite_failed_jobs_follow_cursor_pagination(self) -> None: next_url = ( From fa722b9f0188179a66728f2f882e9d5bc92c9b0c Mon Sep 17 00:00:00 2001 From: Connor Carpenter Date: Mon, 10 Aug 2026 15:05:24 -0700 Subject: [PATCH 10/30] [Rust Frontend][gRPC] Add explicit data-parallel rank routing (#51178) Signed-off-by: Connor Carpenter Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/src/cmd/src/cli.rs | 8 ++++ rust/src/cmd/src/cli/tests.rs | 9 ++++ rust/src/engine-core-client/src/client.rs | 9 ---- .../engine-core-client/src/client/state.rs | 40 ++++++++++-------- rust/src/engine-core-client/src/error.rs | 9 +++- rust/src/engine-core-client/src/transport.rs | 26 ++++++++++-- rust/src/mock-engine/src/lib.rs | 15 ++++--- .../examples/external_engine_openai_qwen.rs | 1 + rust/src/server/src/config.rs | 42 +++++++++++++++++++ rust/src/server/src/grpc/control.rs | 2 +- rust/src/server/src/grpc/inference.rs | 27 +++++++++++- rust/src/server/src/grpc/tests.rs | 42 ++++++++++++++++--- rust/src/server/src/lib.rs | 1 + rust/src/server/src/routes/tests.rs | 18 ++++---- rust/src/server/src/routes/world_size.rs | 2 +- rust/src/server/src/state.rs | 15 +++++++ rust/src/text/src/error.rs | 5 ++- vllm/entrypoints/cli/serve.py | 1 + vllm/v1/utils.py | 3 ++ 19 files changed, 215 insertions(+), 60 deletions(-) diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index b62e75ebac19..833b93ee43e3 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -457,6 +457,7 @@ impl SharedRuntimeArgs { coordinator_address: Option, engine_start_index: u32, engine_count: usize, + data_parallel_size: usize, ) -> Config { let ready_timeout = self.ready_timeout(); let shutdown_timeout = self.shutdown_timeout(); @@ -474,6 +475,7 @@ impl SharedRuntimeArgs { engine_count, ready_timeout, }, + data_parallel_size, coordinator_mode: match coordinator_address { Some(address) => CoordinatorMode::External { address }, None => CoordinatorMode::None, @@ -530,6 +532,7 @@ impl SharedRuntimeArgs { local_input_address, local_output_address, }, + data_parallel_size: engine_count, coordinator_mode: CoordinatorMode::MaybeInProc, model: self.model, served_model_name: self.served_model_name, @@ -661,6 +664,9 @@ pub struct FrontendArgs { /// Total number of data-parallel engines expected for this frontend. #[arg(long, default_value_t = 1)] pub engine_count: usize, + /// Deployment-wide configured data-parallel size. Defaults to engine count. + #[arg(long)] + pub data_parallel_size: Option, /// Shared frontend arguments as one JSON object. #[arg(long = "args-json", value_parser = parse_runtime_args_json, value_name = "JSON")] @@ -670,6 +676,7 @@ pub struct FrontendArgs { impl FrontendArgs { /// Convert the CLI arguments into the OpenAI server's runtime config. pub fn into_config(self) -> Config { + let data_parallel_size = self.data_parallel_size.unwrap_or(self.engine_count); self.runtime.into_bootstrapped_config( self.listen_fd, self.input_address, @@ -677,6 +684,7 @@ impl FrontendArgs { self.coordinator_address, self.engine_start_index, self.engine_count, + data_parallel_size, ) } } diff --git a/rust/src/cmd/src/cli/tests.rs b/rust/src/cmd/src/cli/tests.rs index 28cea8a371f2..98a22ad3244c 100644 --- a/rust/src/cmd/src/cli/tests.rs +++ b/rust/src/cmd/src/cli/tests.rs @@ -785,6 +785,7 @@ fn frontend_args_accept_json() { ), engine_start_index: 0, engine_count: 1, + data_parallel_size: None, runtime: SharedRuntimeArgs { model: "Qwen/Qwen3-0.6B", engine_ready_timeout_secs: 600, @@ -846,6 +847,8 @@ fn frontend_args_json_applies_defaults() { "ipc:///tmp/input.sock", "--output-address", "ipc:///tmp/output.sock", + "--engine-count", + "4", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, ]) @@ -861,6 +864,7 @@ fn frontend_args_json_applies_defaults() { assert_eq!(args.runtime.renderer, RendererSelection::Auto); assert_eq!(args.runtime.max_logprobs, None); assert_eq!(args.runtime.shutdown_timeout, 0); + assert_eq!(args.into_config().data_parallel_size, 4); } #[test] @@ -1517,6 +1521,7 @@ fn serve_frontend_config_uses_dp_address_as_advertised_host() { "", ), }, + data_parallel_size: 4, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", served_model_name: [], @@ -1602,6 +1607,7 @@ fn serve_frontend_config_keeps_tcp_transport_for_non_local_only_topology() { local_input_address: None, local_output_address: None, }, + data_parallel_size: 4, coordinator_mode: MaybeInProc, model: "Qwen/Qwen3-0.6B", served_model_name: [], @@ -1685,6 +1691,8 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present "3", "--engine-count", "1", + "--data-parallel-size", + "4", "--args-json", r#"{"model_tag":"Qwen/Qwen3-0.6B"}"#, ]) @@ -1704,6 +1712,7 @@ fn frontend_config_uses_external_coordinator_when_coordinator_address_is_present engine_count: 1, ready_timeout: 600s, }, + data_parallel_size: 4, coordinator_mode: External { address: "tcp://127.0.0.1:7000", }, diff --git a/rust/src/engine-core-client/src/client.rs b/rust/src/engine-core-client/src/client.rs index a5709e52d67f..d1c20ad58220 100644 --- a/rust/src/engine-core-client/src/client.rs +++ b/rust/src/engine-core-client/src/client.rs @@ -451,15 +451,6 @@ impl EngineCoreClient { .world_size } - /// Return the data parallel size from the parallel config, if available. - pub fn data_parallel_size(&self) -> u64 { - self.engines - .first() - .expect("engine core client requires at least one engine") - .ready_response - .data_parallel_size - } - /// Get the model name associated with this client used for metrics /// labeling. pub fn model_name(&self) -> &str { diff --git a/rust/src/engine-core-client/src/client/state.rs b/rust/src/engine-core-client/src/client/state.rs index 28e701ba7320..59f77bf71cbd 100644 --- a/rust/src/engine-core-client/src/client/state.rs +++ b/rust/src/engine-core-client/src/client/state.rs @@ -161,15 +161,16 @@ impl RequestRegistry { fn choose_engine_for_request(&mut self, data_parallel_rank: Option) -> Result { if let Some(rank) = data_parallel_rank { - // Route to the engine at the specified rank index. - let engine_id = EngineId::from_engine_index(rank); - return self - .routing_per_engine - .contains_key(&engine_id) - .then_some(engine_id) + let engine_id = u16::try_from(rank).ok().map(EngineId::from_engine_index); + return engine_id + .filter(|engine_id| self.routing_per_engine.contains_key(engine_id)) .ok_or_else(|| Error::InvalidDataParallelRank { rank, - num_engines: self.routing_per_engine.len() as u32, + connected_ranks: self + .routing_per_engine + .keys() + .filter_map(EngineId::engine_index) + .collect(), }); } @@ -343,6 +344,9 @@ impl RequestRegistry { } fn apply_scheduler_counts(&mut self, engine_index: u32, next: EngineLoadSnapshot) -> bool { + let Ok(engine_index) = u16::try_from(engine_index) else { + return false; + }; let engine_id = EngineId::from_engine_index(engine_index); let Some(state) = self.routing_per_engine.get_mut(&engine_id) else { return false; @@ -840,26 +844,26 @@ mod tests { error, crate::error::Error::InvalidDataParallelRank { rank: 2, - num_engines: 2, - } + connected_ranks, + } if connected_ranks == vec![0, 1] )); } #[test] - fn register_with_rank_on_single_engine_only_accepts_zero() { - let engine_0 = EngineId::from_engine_index(0); - let mut registry = RequestRegistry::new(&[connected_engine(engine_0.clone())]); + fn register_with_rank_uses_global_engine_identity() { + let engine_3 = EngineId::from_engine_index(3); + let mut registry = RequestRegistry::new(&[connected_engine(engine_3.clone())]); - let (chosen, _) = registry.register("req-ok".to_string(), None, Some(0)).unwrap(); - assert_eq!(chosen, engine_0); + let (chosen, _) = registry.register("req-ok".to_string(), None, Some(3)).unwrap(); + assert_eq!(chosen, engine_3); - let error = registry.register("req-bad".to_string(), None, Some(1)).unwrap_err(); + let error = registry.register("req-bad".to_string(), None, Some(0)).unwrap_err(); assert!(matches!( error, crate::error::Error::InvalidDataParallelRank { - rank: 1, - num_engines: 1, - } + rank: 0, + connected_ranks, + } if connected_ranks == vec![3] )); } diff --git a/rust/src/engine-core-client/src/error.rs b/rust/src/engine-core-client/src/error.rs index 94809a67860f..98a578061efe 100644 --- a/rust/src/engine-core-client/src/error.rs +++ b/rust/src/engine-core-client/src/error.rs @@ -68,8 +68,13 @@ pub enum Error { ControlClosed { message: String }, #[error("request `{request_id}` is already in flight")] DuplicateRequestId { request_id: String }, - #[error("data parallel rank {rank} is out of range for {num_engines} engine(s)")] - InvalidDataParallelRank { rank: u32, num_engines: u32 }, + #[error( + "data parallel rank {rank} is not connected to this frontend; connected ranks: {connected_ranks:?}" + )] + InvalidDataParallelRank { + rank: u32, + connected_ranks: Vec, + }, #[error("engine-core output dispatcher closed: {message}")] DispatcherClosed { message: String }, #[error("engine-core client is closed: {message}")] diff --git a/rust/src/engine-core-client/src/transport.rs b/rust/src/engine-core-client/src/transport.rs index d3e3ecae2bf2..7aac8774cac7 100644 --- a/rust/src/engine-core-client/src/transport.rs +++ b/rust/src/engine-core-client/src/transport.rs @@ -66,8 +66,8 @@ impl EngineId { /// Construct an engine id from the Python-compatible engine index encoding /// (two-byte little-endian). - pub fn from_engine_index(value: u32) -> Self { - Self(Bytes::copy_from_slice(&(value as u16).to_le_bytes())) + pub fn from_engine_index(value: u16) -> Self { + Self(Bytes::copy_from_slice(&value.to_le_bytes())) } } @@ -336,6 +336,22 @@ pub async fn connect_bootstrapped( if engine_count == 0 { bail_unexpected_handshake_message!("expected engine_count >= 1"); } + let engine_start_index = + u16::try_from(engine_start_index).map_err(|_| Error::UnexpectedHandshakeMessage { + message: "engine_start_index exceeds the two-byte engine identity limit".to_string(), + })?; + let engine_end_index = + usize::from(engine_start_index).checked_add(engine_count).ok_or_else(|| { + Error::UnexpectedHandshakeMessage { + message: "engine_start_index + engine_count overflows".to_string(), + } + })?; + if engine_end_index > usize::from(u16::MAX) + 1 { + return Err(Error::UnexpectedHandshakeMessage { + message: "engine_start_index + engine_count exceeds the two-byte engine identity limit" + .to_string(), + }); + } let mut input_socket = RouterSocket::new(); let input_address = input_socket.bind(input_address).await?.to_string(); @@ -345,8 +361,10 @@ pub async fn connect_bootstrapped( let engines = wait_for_input_registrations( &mut input_socket, - (0..engine_count) - .map(|offset| EngineId::from_engine_index(engine_start_index + offset as u32)), + (0..engine_count).map(|offset| { + let offset = u16::try_from(offset).expect("validated engine offset fits u16"); + EngineId::from_engine_index(engine_start_index + offset) + }), ready_timeout, ) .await?; diff --git a/rust/src/mock-engine/src/lib.rs b/rust/src/mock-engine/src/lib.rs index d4e8581b2f97..63b1d432ba77 100644 --- a/rust/src/mock-engine/src/lib.rs +++ b/rust/src/mock-engine/src/lib.rs @@ -48,10 +48,13 @@ pub struct Opt { } /// Run one mock engine until shutdown or transport failure. -async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> Result<()> { +async fn run_engine(engine_index: usize, opt: Opt, shutdown: CancellationToken) -> Result<()> { + let engine_index = + u16::try_from(engine_index).context("mock engine index exceeds two-byte identity limit")?; + let engine_id = EngineId::from_engine_index(engine_index); let MockEngineSockets { data_sockets, .. } = connect_to_frontend( &opt.handshake_address, - EngineId::from_engine_index(engine_index), + engine_id, MockEngineConfig::default(), ) .await @@ -71,7 +74,7 @@ async fn run_engine(engine_index: u32, opt: Opt, shutdown: CancellationToken) -> )); // Engine loop: input_rx -> engine logic -> output_tx let mut engine_loop = tokio::spawn(engine::run_engine_loop( - engine_index, + u32::from(engine_index), opt, input_rx, output_tx, @@ -111,11 +114,7 @@ pub async fn run(opt: Opt, shutdown: CancellationToken) -> Result<()> { let mut engines = JoinSet::new(); for engine_index in 0..opt.engine_count { - engines.spawn(run_engine( - engine_index as u32, - opt.clone(), - shutdown.clone(), - )); + engines.spawn(run_engine(engine_index, opt.clone(), shutdown.clone())); } tokio::select! { diff --git a/rust/src/server/examples/external_engine_openai_qwen.rs b/rust/src/server/examples/external_engine_openai_qwen.rs index 387776263059..591252699a10 100644 --- a/rust/src/server/examples/external_engine_openai_qwen.rs +++ b/rust/src/server/examples/external_engine_openai_qwen.rs @@ -58,6 +58,7 @@ async fn main() -> Result<()> { local_input_address: None, local_output_address: None, }, + data_parallel_size: args.engine_count, coordinator_mode: CoordinatorMode::MaybeInProc, model: args.model, served_model_name: vec![], diff --git a/rust/src/server/src/config.rs b/rust/src/server/src/config.rs index 76848b3caa94..b67e87f45f3c 100644 --- a/rust/src/server/src/config.rs +++ b/rust/src/server/src/config.rs @@ -161,6 +161,8 @@ impl TlsConfig { pub struct Config { /// Frontend-to-engine transport setup. pub transport_mode: TransportMode, + /// Deployment-wide data-parallel size retained by the frontend. + pub data_parallel_size: usize, /// Requested frontend-side coordinator behavior. pub coordinator_mode: CoordinatorMode, /// Backend model identifier used for engine-core loading. @@ -237,6 +239,46 @@ impl Config { max_logprobs ); } + if self.data_parallel_size == 0 { + bail!("data parallel size must be at least 1"); + } + if self.data_parallel_size > usize::from(u16::MAX) + 1 { + bail!( + "data parallel size ({}) exceeds the two-byte engine identity limit", + self.data_parallel_size + ); + } + match &self.transport_mode { + TransportMode::HandshakeOwner { engine_count, .. } => { + if *engine_count != self.data_parallel_size { + bail!( + "managed frontend engine count ({engine_count}) must equal data parallel size ({})", + self.data_parallel_size + ); + } + } + TransportMode::Bootstrapped { + engine_start_index, + engine_count, + .. + } => { + if *engine_count == 0 { + bail!("engine count must be at least 1"); + } + let engine_start_index = usize::try_from(*engine_start_index) + .map_err(|_| anyhow::anyhow!("engine start index does not fit usize"))?; + let engine_end_index = + engine_start_index.checked_add(*engine_count).ok_or_else(|| { + anyhow::anyhow!("engine start index + engine count overflows") + })?; + if engine_end_index > self.data_parallel_size { + bail!( + "connected engine range [{engine_start_index}, {engine_end_index}) exceeds data parallel size {}", + self.data_parallel_size + ); + } + } + } Ok(()) } diff --git a/rust/src/server/src/grpc/control.rs b/rust/src/server/src/grpc/control.rs index 897209822252..201d8ace95ee 100644 --- a/rust/src/server/src/grpc/control.rs +++ b/rust/src/server/src/grpc/control.rs @@ -31,7 +31,7 @@ impl ControlServiceImpl { pb::ParallelismInfo { tensor_parallel_size: ready.tensor_parallel_size, pipeline_parallel_size: ready.pipeline_parallel_size, - data_parallel_size: ready.data_parallel_size.min(u64::from(u32::MAX)) as u32, + data_parallel_size: self.state.data_parallel_size().min(u32::MAX as usize) as u32, data_parallel_rank: ready.data_parallel_rank, decode_context_parallel_size: ready.decode_context_parallel_size, } diff --git a/rust/src/server/src/grpc/inference.rs b/rust/src/server/src/grpc/inference.rs index 3fc56fd0d92e..8eebdb721002 100644 --- a/rust/src/server/src/grpc/inference.rs +++ b/rust/src/server/src/grpc/inference.rs @@ -22,6 +22,8 @@ use crate::state::AppState; pub(crate) type InferenceGrpcService = InferenceServer; +const DATA_PARALLEL_RANK_METADATA_KEY: &str = "x-data-parallel-rank"; + /// gRPC inference service backed by the shared application state. pub struct InferenceServiceImpl { state: Arc, @@ -41,6 +43,7 @@ impl InferenceServiceImpl { async fn prepare_request( &self, mut proto_request: pb::GenerateRequest, + data_parallel_rank: Option, stream: bool, rpc: &'static str, ) -> Result { @@ -62,6 +65,7 @@ impl InferenceServiceImpl { rpc, %model, media_count, + ?data_parallel_rank, ); info!(parent: &request_span, "gRPC inference request received"); @@ -70,6 +74,7 @@ impl InferenceServiceImpl { let mut text_request = convert::to_text_request(proto_request, stream, self.state.served_model_names())?; text_request.arrival_time = Some(arrival_time); + text_request.data_parallel_rank = data_parallel_rank; let media = convert::media_parts_from_request(media)?; if !media.is_empty() { @@ -115,6 +120,20 @@ impl InferenceServiceImpl { } } +fn data_parallel_rank_from_metadata( + request: &Request, +) -> Result, Status> { + let Some(value) = request.metadata().get(DATA_PARALLEL_RANK_METADATA_KEY) else { + return Ok(None); + }; + let value = value.to_str().map_err(|_| { + Status::invalid_argument("x-data-parallel-rank metadata must be an unsigned 32-bit integer") + })?; + value.trim().parse::().map(Some).map_err(|_| { + Status::invalid_argument("x-data-parallel-rank metadata must be an unsigned 32-bit integer") + }) +} + #[tonic::async_trait] impl pb::inference_server::Inference for InferenceServiceImpl { type GenerateStreamStream = @@ -125,13 +144,14 @@ impl pb::inference_server::Inference for InferenceServiceImpl { &self, request: Request, ) -> Result, Status> { + let data_parallel_rank = data_parallel_rank_from_metadata(&request)?; let proto_req = request.into_inner(); let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); let PreparedGrpcRequest { text_request, request_span, started_at, - } = self.prepare_request(proto_req, false, "Generate").await?; + } = self.prepare_request(proto_req, data_parallel_rank, false, "Generate").await?; let stream = self .state @@ -186,13 +206,16 @@ impl pb::inference_server::Inference for InferenceServiceImpl { &self, request: Request, ) -> Result, Status> { + let data_parallel_rank = data_parallel_rank_from_metadata(&request)?; let proto_req = request.into_inner(); let response_opts = ResponseOpts::from_proto(proto_req.response.as_ref()); let PreparedGrpcRequest { text_request, request_span, started_at, - } = self.prepare_request(proto_req, true, "GenerateStream").await?; + } = self + .prepare_request(proto_req, data_parallel_rank, true, "GenerateStream") + .await?; let stream = self .state diff --git a/rust/src/server/src/grpc/tests.rs b/rust/src/server/src/grpc/tests.rs index fdc98333f178..274ac67fecf5 100644 --- a/rust/src/server/src/grpc/tests.rs +++ b/rust/src/server/src/grpc/tests.rs @@ -843,6 +843,37 @@ async fn unary_generate_missing_prompt_returns_invalid_argument() { server_task.abort(); } +#[tokio::test(flavor = "multi_thread", worker_threads = 2)] +#[serial] +async fn unary_generate_unconnected_data_parallel_rank_returns_invalid_argument() { + let (mut client, server_task, _engine_task) = grpc_test_server( + EngineId::from_engine_index(3), + default_stream_output_specs(), + ) + .await; + + let mut request = tonic::Request::new(pb::GenerateRequest { + request_id: "test-unconnected-dp-rank".to_string(), + model: "test-model".to_string(), + prompt: Some(pb::generate_request::Prompt::Text("hello".to_string())), + ..Default::default() + }); + request.metadata_mut().insert( + "x-data-parallel-rank", + "0".parse().expect("valid metadata value"), + ); + + let status = client + .generate(request) + .await + .expect_err("rank 0 should not select globally ranked engine 3"); + + assert_eq!(status.code(), tonic::Code::InvalidArgument); + assert!(status.message().contains("connected ranks: [3]")); + + server_task.abort(); +} + #[tokio::test(flavor = "multi_thread", worker_threads = 2)] #[serial] async fn unary_generate_min_tokens_above_max_tokens_returns_invalid_argument() { @@ -1523,7 +1554,9 @@ async fn control_aggregates_multi_engine_capacity() { ready_1.data_parallel_rank = 1; let engine_tasks = [ready_0, ready_1].map(|ready| { - let engine_id = EngineId::from_engine_index(ready.data_parallel_rank); + let engine_id = EngineId::from_engine_index( + ready.data_parallel_rank.try_into().expect("test rank fits engine identity"), + ); MockEngineTask::new(spawn_mock_engine_task_with_ready( handshake_address.clone(), engine_id, @@ -1551,10 +1584,8 @@ async fn control_aggregates_multi_engine_capacity() { Llm::new(client), Arc::new(FakeTextBackend) as Arc, ); - let service = ControlServiceImpl::new(Arc::new(AppState::new( - vec!["test-model".to_string()], - chat, - ))); + let state = AppState::new(vec!["test-model".to_string()], chat).with_data_parallel_size(4); + let service = ControlServiceImpl::new(Arc::new(state)); let server = pb::control_server::Control::get_server_info( &service, @@ -1565,6 +1596,7 @@ async fn control_aggregates_multi_engine_capacity() { .into_inner(); assert_eq!(server.max_model_len, 4_096); assert_eq!(server.total_kv_blocks, 30); + assert_eq!(server.parallelism.unwrap().data_parallel_size, 4); drop(engine_tasks); } diff --git a/rust/src/server/src/lib.rs b/rust/src/server/src/lib.rs index c36aaac2a19b..ca19d0ab9e9d 100644 --- a/rust/src/server/src/lib.rs +++ b/rust/src/server/src/lib.rs @@ -141,6 +141,7 @@ async fn build_state(config: &Config) -> Result> { .with_model_path(config.model.clone()) .with_api_server_options(config.api_server_options) .with_server_info(ServerInfoSnapshot::from_config(config)) + .with_data_parallel_size(config.data_parallel_size) .with_api_keys(config.api_keys.clone()) .with_cors(config.cors.clone()) .with_profiler(config.profiler.clone()), diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index 646d904ea924..92f1d44868b4 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -725,11 +725,11 @@ async fn test_app_with_dev_mode(dev_mode_enabled: bool) -> axum::Router { /// stays alive for the duration of the test. async fn test_dev_mode_app_with_ready( ready_response: vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse, + data_parallel_size: usize, ) -> (axum::Router, MockEngineTask) { let ipc = IpcNamespace::new().expect("create ipc namespace"); let handshake_address = ipc.handshake_endpoint(); let engine_id = b"engine-world-size".to_vec(); - let engine_task = MockEngineTask::new(spawn_mock_engine_task_with_ready( handshake_address.clone(), engine_id.clone(), @@ -750,10 +750,10 @@ async fn test_dev_mode_app_with_ready( let chat = ChatLlm::from_shared_backend(test_llm(client), Arc::new(FakeChatBackend::new())); let app = build_router_with_dev_mode( - Arc::new(AppState::new( - vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], - chat, - )), + Arc::new( + AppState::new(vec!["Qwen/Qwen1.5-0.5B-Chat".to_string()], chat) + .with_data_parallel_size(data_parallel_size), + ), true, ); (app, engine_task) @@ -6313,7 +6313,7 @@ async fn tokenize_allows_prompts_at_or_above_max_model_len() { max_model_len: 4, ..default_ready_response() }; - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await; + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 1).await; let (completion_status, completion_json) = post_json( &mut app, @@ -6569,10 +6569,10 @@ async fn world_size_endpoint_is_dev_mode_only() { async fn world_size_includes_data_parallelism_by_default() { let ready = vllm_engine_core_client::protocol::handshake::EngineCoreReadyResponse { world_size: 2, - data_parallel_size: 4, + data_parallel_size: 1, ..default_ready_response() }; - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await; + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await; let response = app .call( @@ -6598,7 +6598,7 @@ async fn world_size_excludes_data_parallelism_when_include_dp_false() { data_parallel_size: 4, ..default_ready_response() }; - let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready).await; + let (mut app, _engine_task) = test_dev_mode_app_with_ready(ready, 4).await; let response = app .call( diff --git a/rust/src/server/src/routes/world_size.rs b/rust/src/server/src/routes/world_size.rs index 2347158477b2..2be1bf76b7e7 100644 --- a/rust/src/server/src/routes/world_size.rs +++ b/rust/src/server/src/routes/world_size.rs @@ -47,7 +47,7 @@ pub async fn get_world_size( let ws = client.world_size(); let world_size = if params.include_dp { - let dp = client.data_parallel_size(); + let dp = state.data_parallel_size() as u64; ws * dp } else { ws diff --git a/rust/src/server/src/state.rs b/rust/src/server/src/state.rs index 722dede8dda0..f745cf7c0042 100644 --- a/rust/src/server/src/state.rs +++ b/rust/src/server/src/state.rs @@ -40,6 +40,8 @@ pub struct AppState { pub cors: CorsConfig, /// Runtime server information returned by `/server_info`, when available. server_info: Option, + /// Deployment-wide data-parallel size retained by the frontend. + data_parallel_size: usize, /// SHA-256 hashes of API keys accepted as bearer tokens for guarded routes. api_key_hashes: Vec, /// Number of in-flight inference requests currently owned by this frontend. @@ -69,12 +71,14 @@ impl AppState { !served_model_names.is_empty(), "served_model_names must not be empty" ); + let data_parallel_size = chat.engine_core_client().engine_count(); Self { served_model_names, chat, api_server_options: ApiServerOptions::default(), cors: CorsConfig::default(), server_info: None, + data_parallel_size, api_key_hashes: Vec::new(), server_load: AtomicU64::new(0), lora_manager: LoraManager::new(), @@ -114,6 +118,17 @@ impl AppState { self } + /// Set the deployment-wide data-parallel size reported by frontend APIs. + pub(crate) fn with_data_parallel_size(mut self, size: usize) -> Self { + self.data_parallel_size = size; + self + } + + /// Return the deployment-wide data-parallel size. + pub(crate) fn data_parallel_size(&self) -> usize { + self.data_parallel_size + } + /// Build a `/server_info` response payload. pub(crate) fn server_info_response( &self, diff --git a/rust/src/text/src/error.rs b/rust/src/text/src/error.rs index 551c2c9df800..58ed85242ce8 100644 --- a/rust/src/text/src/error.rs +++ b/rust/src/text/src/error.rs @@ -62,7 +62,10 @@ impl Error { | Self::InvalidRepetitionDetection { .. } // An empty tokenized prompt detected later, at request prepare // time, surfaces through the transparent Llm wrapper. - | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) => true, + | Self::Llm(LlmError::EmptyPromptTokenIds { .. }) + | Self::Llm(LlmError::EngineCoreClient( + EngineCoreError::InvalidDataParallelRank { .. }, + )) => true, _ => false, } } diff --git a/vllm/entrypoints/cli/serve.py b/vllm/entrypoints/cli/serve.py index f90a4f8c1d62..9f29cd545f32 100644 --- a/vllm/entrypoints/cli/serve.py +++ b/vllm/entrypoints/cli/serve.py @@ -352,6 +352,7 @@ def signal_handler(signum, frame): output_address=addresses.outputs[0], engine_start_index=expected_engine_start_index, engine_count=expected_engine_count, + data_parallel_size=parallel_config.data_parallel_size, stats_update_address=stats_update_address, ) else: diff --git a/vllm/v1/utils.py b/vllm/v1/utils.py index f5e0caa8188e..8d482f58fdf9 100644 --- a/vllm/v1/utils.py +++ b/vllm/v1/utils.py @@ -339,6 +339,7 @@ def __init__( output_address: str, engine_start_index: int, engine_count: int, + data_parallel_size: int, stats_update_address: str | None = None, ): import os @@ -360,6 +361,8 @@ def __init__( str(engine_start_index), "--engine-count", str(engine_count), + "--data-parallel-size", + str(data_parallel_size), ] if stats_update_address is not None: cmd.extend(["--coordinator-address", stats_update_address]) From 98a4144a4197c2b96e4df7bfd9c386d4fbdb1f4f Mon Sep 17 00:00:00 2001 From: Tristan Rice Date: Mon, 10 Aug 2026 15:14:53 -0700 Subject: [PATCH 11/30] [Bugfix] Preserve non-logitproc entry points in tests (#51097) Signed-off-by: Tristan Rice --- .../logits_processors/test_custom_offline.py | 28 +++++++++++++++++++ tests/v1/logits_processors/utils.py | 13 +++++---- 2 files changed, 36 insertions(+), 5 deletions(-) diff --git a/tests/v1/logits_processors/test_custom_offline.py b/tests/v1/logits_processors/test_custom_offline.py index 325ca48b597e..302a72057eaa 100644 --- a/tests/v1/logits_processors/test_custom_offline.py +++ b/tests/v1/logits_processors/test_custom_offline.py @@ -1,9 +1,11 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import importlib.metadata from typing import Any import pytest +import tests.v1.logits_processors.utils as logitproc_test_utils from tests.utils import create_new_process_for_each_test, set_random_seed from tests.v1.logits_processors.utils import ( DUMMY_LOGITPROC_ARG, @@ -42,6 +44,32 @@ ] +def test_fake_entrypoint_preserves_other_groups(monkeypatch): + other_group_entrypoints = object() + monkeypatch.setattr( + importlib.metadata, + "entry_points", + lambda **kwargs: other_group_entrypoints, + ) + monkeypatch.setattr( + logitproc_test_utils, "requires_spawn_multiprocessing", lambda: False + ) + + setup_fake_entrypoint(monkeypatch) + + logitproc_entrypoints = importlib.metadata.entry_points( + group=logitproc_test_utils.LOGITSPROCS_GROUP + ) + assert logitproc_entrypoints.names == [ + logitproc_test_utils.DUMMY_LOGITPROC_ENTRYPOINT + ] + assert ( + importlib.metadata.entry_points(group="another.entrypoint.group") + is other_group_entrypoints + ) + assert importlib.metadata.entry_points() is other_group_entrypoints + + def _run_test(kwargs: dict, logitproc_loaded: bool) -> None: """Compare `LLM` instance initialized with specified `kwargs` against reference `LLM` instance. diff --git a/tests/v1/logits_processors/utils.py b/tests/v1/logits_processors/utils.py index a2181d8e188c..fb45b73c9893 100644 --- a/tests/v1/logits_processors/utils.py +++ b/tests/v1/logits_processors/utils.py @@ -234,11 +234,6 @@ def register_fake_entrypoint(monkeypatch) -> str: return str(tmpdir) -def fake_entry_points(group: str) -> EntryPoints: - """Fake version of importlib.metadata.entry_points.""" - return EntryPoints(group) - - def setup_fake_entrypoint(monkeypatch) -> None: """Expose the dummy logitproc entrypoint for the current platform.""" if requires_spawn_multiprocessing(): @@ -248,5 +243,13 @@ def setup_fake_entrypoint(monkeypatch) -> None: import importlib.metadata + original_entry_points = importlib.metadata.entry_points + + def fake_entry_points(**kwargs): + group = kwargs.get("group") + if group == LOGITSPROCS_GROUP: + return EntryPoints(group) + return original_entry_points(**kwargs) + monkeypatch.setattr(importlib.metadata, "entry_points", fake_entry_points) monkeypatch.setenv("VLLM_WORKER_MULTIPROC_METHOD", "fork") From c3cac8c63d917759a38e94dfa933e71196de40bb Mon Sep 17 00:00:00 2001 From: Almog Tavor <70065337+almogtavor@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:36:42 +0300 Subject: [PATCH 12/30] [Bugfix][MiMo] Apply vision attention sinks in the window attention path (#49815) Signed-off-by: almogtavor --- tests/models/multimodal/test_mimo_v2_omni.py | 91 +++++++++++++++++++ vllm/model_executor/models/mimo_v2_omni.py | 37 ++++++-- .../attention/ops/triton_prefill_attention.py | 16 +++- 3 files changed, 132 insertions(+), 12 deletions(-) create mode 100644 tests/models/multimodal/test_mimo_v2_omni.py diff --git a/tests/models/multimodal/test_mimo_v2_omni.py b/tests/models/multimodal/test_mimo_v2_omni.py new file mode 100644 index 000000000000..d4e1d9ec18f5 --- /dev/null +++ b/tests/models/multimodal/test_mimo_v2_omni.py @@ -0,0 +1,91 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""MiMo-V2 vision window attention has to apply the per-head sink logits.""" + +import pytest +import torch + +from tests.utils import ensure_current_vllm_config +from vllm.distributed.parallel_state import ( + init_distributed_environment, + initialize_model_parallel, +) +from vllm.platforms import current_platform +from vllm.utils.network_utils import get_open_port + +EMBED_DIM = 256 +NUM_HEADS = 4 +HEAD_DIM = 64 +WINDOW = 8 +# One sequence shorter than the window, one longer. +SEQ_LENS = [5, 37] + + +@pytest.fixture(scope="module") +def vision_attn_env(): + init_distributed_environment( + world_size=1, + rank=0, + local_rank=0, + distributed_init_method=f"tcp://127.0.0.1:{get_open_port()}", + backend="nccl", + ) + default_dtype = torch.get_default_dtype() + torch.set_default_dtype(torch.bfloat16) + with ensure_current_vllm_config(): + initialize_model_parallel(tensor_model_parallel_size=1) + yield + torch.set_default_dtype(default_dtype) + + +def _reference(q, k, v, cu_seqlens, sinks, scale): + """Dense windowed softmax with the sink added to each sequence's key 0.""" + groups = q.shape[1] // k.shape[1] + out = torch.empty_like(q, dtype=torch.float32) + for start, end in zip(cu_seqlens[:-1].tolist(), cu_seqlens[1:].tolist()): + qs = q[start:end].float() + ks = k[start:end].float().repeat_interleave(groups, dim=1) + vs = v[start:end].float().repeat_interleave(groups, dim=1) + scores = torch.einsum("qhd,khd->hqk", qs, ks) * scale + scores[..., 0] += sinks.float().view(-1, 1) + pos = torch.arange(end - start, device=q.device) + outside = (pos.view(-1, 1) - pos.view(1, -1)).abs() > WINDOW + scores.masked_fill_(outside, -torch.inf) + out[start:end] = torch.einsum("hqk,khd->qhd", scores.softmax(-1), vs) + return out + + +@pytest.mark.skipif(not current_platform.is_cuda(), reason="requires flash-attn") +@pytest.mark.parametrize("num_kv_heads", [NUM_HEADS, NUM_HEADS // 2]) +def test_window_attention_applies_sinks(vision_attn_env, num_kv_heads): + from vllm.model_executor.models.mimo_v2_omni import MiMoVisionAttention + + torch.manual_seed(0) + attn = MiMoVisionAttention( + embed_dim=EMBED_DIM, + num_heads=NUM_HEADS, + num_kv_heads=num_kv_heads, + qk_channels=HEAD_DIM, + kv_channels=HEAD_DIM, + use_sink=True, + visual_token_window_size=WINDOW, + ).cuda() + attn.sinks.data.normal_() + + total = sum(SEQ_LENS) + opts = dict(device="cuda", dtype=torch.bfloat16) + q = torch.randn(total, NUM_HEADS, HEAD_DIM, **opts) + k = torch.randn(total, num_kv_heads, HEAD_DIM, **opts) + v = torch.randn(total, num_kv_heads, HEAD_DIM, **opts) + cu_seqlens = torch.tensor( + [0, *torch.tensor(SEQ_LENS).cumsum(0).tolist()], + device="cuda", + dtype=torch.int32, + ) + + out = attn._forward_window_attn(q, k, v, cu_seqlens, max(SEQ_LENS)) + ref = _reference(q, k, v, cu_seqlens, attn.sinks, attn.scale) + + # bf16 attention lands at ~2e-3 here; dropping the sinks lands at ~1e-1. + error = ((out.float() - ref).norm() / ref.norm()).item() + assert error < 1e-2, f"sink-corrected output is off by {error:.2e}" diff --git a/vllm/model_executor/models/mimo_v2_omni.py b/vllm/model_executor/models/mimo_v2_omni.py index 747cb0e88b2b..589ca4f61eab 100644 --- a/vllm/model_executor/models/mimo_v2_omni.py +++ b/vllm/model_executor/models/mimo_v2_omni.py @@ -193,7 +193,7 @@ def __init__( # Rotary embeddings applied separately to Q and K self.apply_rotary_emb = ApplyRotaryEmb(enforce_enable=True) - # Sink attention weights (loaded but not used in vLLM flash_attn) + # Per-head sink logits, applied in the window attention path. # The checkpoint stores these only for non-full-attention blocks self.use_sink = use_sink if use_sink: @@ -214,21 +214,38 @@ def _forward_window_attn( cu_seqlens: torch.Tensor, max_seqlen: torch.Tensor, ) -> torch.Tensor: - """Window attention via flash_attn_varlen_func with window_size.""" - from vllm.vllm_flash_attn import flash_attn_varlen_func + """Window attention with the per-head sink applied to key 0. + + The reference adds ``sinks[h]`` to the logit of each sequence's first + key, which the Triton prefill kernel supports directly, so the softmax + normalizes over the biased scores in one pass. + """ + from vllm.v1.attention.ops.triton_prefill_attention import ( + context_attention_fwd, + ) w = self.visual_token_window_size - output = flash_attn_varlen_func( + output = torch.empty_like(q) + head_start = self.tp_rank * self.num_heads_per_partition + sinks = ( + self.sinks[head_start : head_start + self.num_heads_per_partition] + if self.sinks is not None + else None + ) + context_attention_fwd( q, k, v, - cu_seqlens_q=cu_seqlens, - cu_seqlens_k=cu_seqlens, - max_seqlen_q=max_seqlen, - max_seqlen_k=max_seqlen, + output, + b_start_loc=cu_seqlens[:-1], + b_seq_len=cu_seqlens[1:] - cu_seqlens[:-1], + max_input_len=max_seqlen, + is_causal=False, softmax_scale=self.scale, - causal=False, - window_size=[w, w], + sliding_window_q=w, + sliding_window_k=w, + sinks=sinks, + sinks_bias_key0=True, ) return output diff --git a/vllm/v1/attention/ops/triton_prefill_attention.py b/vllm/v1/attention/ops/triton_prefill_attention.py index 9628a9e9614a..0fb31ce7e1dc 100644 --- a/vllm/v1/attention/ops/triton_prefill_attention.py +++ b/vllm/v1/attention/ops/triton_prefill_attention.py @@ -58,6 +58,7 @@ def _fwd_kernel( IS_CAUSAL: tl.constexpr, SLIDING_WINDOW_Q: tl.constexpr, SLIDING_WINDOW_K: tl.constexpr, + SINKS_BIAS_KEY0: tl.constexpr, USE_SINKS: tl.constexpr, Lk: tl.constexpr, ): @@ -98,8 +99,15 @@ def _fwd_kernel( # initialize pointer to m and l if USE_SINKS: sink = tl.load(Sinks + cur_head) * 1.4426950408889634 - m_i = tl.full([BLOCK_M], sink, dtype=tl.float32) - l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) + if SINKS_BIAS_KEY0: + # Sinks bias the logit of key 0, so the softmax starts empty and + # normalizes over the biased scores as usual. + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + else: + # Sinks are a null logit that only inflates the denominator. + m_i = tl.full([BLOCK_M], sink, dtype=tl.float32) + l_i = tl.full([BLOCK_M], 1.0, dtype=tl.float32) else: m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") l_i = tl.zeros([BLOCK_M], dtype=tl.float32) @@ -160,6 +168,8 @@ def _fwd_kernel( qk = tl.dot(q, k) qk = tl.where(mask, qk * sm_scale, -1.0e8) + if USE_SINKS and SINKS_BIAS_KEY0: + qk = tl.where(mask & (pos_k == 0), qk + sink, qk) m_ij = tl.maximum(m_i, tl.max(qk, 1)) qk -= m_ij[:, None] p = tl.math.exp2(qk) @@ -217,6 +227,7 @@ def context_attention_fwd( sliding_window_q: int | None = None, sliding_window_k: int | None = None, sinks: torch.Tensor | None = None, + sinks_bias_key0: bool = False, ): """ q, k, v: [b * s, head, head_dim] @@ -268,6 +279,7 @@ def context_attention_fwd( SLIDING_WINDOW_Q=sliding_window_q, SLIDING_WINDOW_K=sliding_window_k, USE_SINKS=sinks is not None, + SINKS_BIAS_KEY0=sinks_bias_key0, num_warps=num_warps, num_stages=1, Lk=Lk, From 3e174bb73c70b677dd339c59c39d7460abc5d02d Mon Sep 17 00:00:00 2001 From: Giancarlo Delfin <32987265+TheEpicDolphin@users.noreply.github.com> Date: Mon, 10 Aug 2026 15:46:13 -0700 Subject: [PATCH 13/30] [Model Runner V2][MTP] Share topk index buffer between draft steps (#47352) Signed-off-by: Giancarlo Delfin Co-authored-by: Woosuk Kwon --- .../spec_decode/autoregressive/speculator.py | 32 +++++++++++--- .../worker/gpu/spec_decode/mtp/speculator.py | 42 ++++++++++++++++++- 2 files changed, 68 insertions(+), 6 deletions(-) diff --git a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py index 7dbd81125325..5ceb6c7558d6 100644 --- a/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/autoregressive/speculator.py @@ -54,6 +54,18 @@ def load_model(self, target_model: nn.Module) -> None: device=self.device, ) + # Lifecycle hooks for model-specific optimizations. Subclasses override + # the ones they need. These fire in both `capture` and `propose` so that + # any state they toggle (e.g. attention flags baked into a CUDA graph) is + # identical at capture time and replay time. + def on_prefill_begin(self, num_reqs: int) -> None: ... + + def on_prefill_end(self, num_reqs: int) -> None: ... + + def on_multi_step_decode_begin(self, num_reqs: int) -> None: ... + + def on_multi_step_decode_end(self, num_reqs: int) -> None: ... + @property def advance_draft_positions(self) -> bool: """ @@ -104,6 +116,8 @@ def capture(self) -> None: assert self.prefill_cudagraph_manager is not None if self.prefill_cudagraph_manager.use_breakable_cg: self.prefill_cudagraph_manager.init_breakable_cg_runner(self.model) + + self.on_prefill_begin(self.max_num_reqs) self.prefill_cudagraph_manager.capture( self._prefill, self.model_state, @@ -113,10 +127,12 @@ def capture(self) -> None: self.kv_cache_config, progress_bar_desc="Capturing prefill CUDA graphs", ) + self.on_prefill_end(self.max_num_reqs) if self.num_speculative_steps == 1: return + self.on_multi_step_decode_begin(self.max_num_reqs) # Capture the decode draft generation routine (model forward + # sample + update_draft_inputs) for a single # step. @@ -130,6 +146,7 @@ def capture(self) -> None: self.kv_cache_config, progress_bar_desc="Capturing decode CUDA graphs", ) + self.on_multi_step_decode_end(self.max_num_reqs) @torch.inference_mode() def propose( @@ -159,7 +176,8 @@ def propose( mm_inputs: tuple[list[torch.Tensor], torch.Tensor] | None = None, is_profile: bool = False, ) -> torch.Tensor: - num_tokens = input_batch.num_tokens_after_padding + num_tokens = input_batch.num_tokens + num_tokens_padded = input_batch.num_tokens_after_padding num_reqs = input_batch.num_reqs max_query_len = input_batch.num_scheduled_tokens.max() max_seq_len = input_batch.seq_lens_cpu_upper_bound[:num_reqs].max().item() @@ -180,7 +198,7 @@ def propose( ) else: hidden_states = last_hidden_states - self.hidden_states[:num_tokens].copy_(hidden_states) + self.hidden_states[:num_tokens_padded].copy_(hidden_states) self._copy_request_inputs( num_reqs, @@ -208,21 +226,22 @@ def propose( num_reqs, # Use the actual number of tokens without padding added by # the target model during FULL cudagraph. - input_batch.num_tokens, + num_tokens, max_query_len, ) prefill_batch_desc, num_tokens_across_dp = dispatch_cg_and_sync_dp( self.prefill_cudagraph_manager, num_reqs, - num_tokens, + num_tokens_padded, uniform_token_count, dp_size=self.dp_size, dp_rank=self.dp_rank, need_eager=is_profile, ) - self._prepare_eplb_forward(input_batch.num_tokens) + self._prepare_eplb_forward(num_tokens) + self.on_prefill_begin(num_reqs) if prefill_batch_desc.cg_mode == CUDAGraphMode.FULL: # Replay the full graph for draft prefill. assert self.prefill_cudagraph_manager is not None @@ -240,6 +259,7 @@ def propose( cudagraph_runtime_mode=prefill_batch_desc.cg_mode, mm_inputs=mm_inputs, ) + self.on_prefill_end(num_reqs) if self.num_speculative_steps == 1: # Early exit. @@ -268,6 +288,7 @@ def propose( need_eager=is_profile, ) + self.on_multi_step_decode_begin(num_reqs) # Generate the remaining num_speculative_steps - 1 draft tokens. self._multi_step_decode( num_reqs, @@ -276,6 +297,7 @@ def propose( num_tokens_across_dp, input_batch.seq_lens_cpu_upper_bound, ) + self.on_multi_step_decode_end(num_reqs) return self.draft_tokens[:num_reqs] diff --git a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py index 4b9354f23e70..4889e9a750f8 100644 --- a/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/mtp/speculator.py @@ -10,9 +10,49 @@ class MTPSpeculator(AutoRegressiveSpeculator): + share_mtp_topk_indices: bool = False + def load_draft_model( self, target_model: nn.Module, target_attn_layer_names: set[str], ) -> nn.Module: - return load_eagle_model(target_model, self.vllm_config) + draft_model = load_eagle_model(target_model, self.vllm_config) + spec_config = self.vllm_config.speculative_config + draft_hf_config = ( + spec_config.draft_model_config.hf_config + if spec_config is not None + else None + ) + # Detect index_share_for_mtp_iteration. When True, the proposer + # toggles skip_topk so step 0 computes MTP's own indices and + # steps 1+ reuse them. + self.share_mtp_topk_indices = ( + getattr(draft_hf_config, "index_share_for_mtp_iteration", False) + and hasattr(draft_model.model, "set_skip_topk") + and hasattr(draft_model.model, "compact_topk_indices") + ) + return draft_model + + def on_prefill_begin(self, num_reqs: int) -> None: + # Step 0 computes its own top-k. Unconditional, so a step that died + # midway cannot leave reuse mode on. + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) + + def on_prefill_end(self, num_reqs: int) -> None: + # Step 0 (prefill) wrote topk indices for every query token in the + # multi-token batch. Compact them down to each request's last token so + # steps 1+ can reuse them from the shared buffer. + if self.share_mtp_topk_indices and self.num_speculative_steps > 1: + self.model.model.compact_topk_indices(self.last_token_indices[:num_reqs]) + + def on_multi_step_decode_begin(self, num_reqs: int) -> None: + # Switch to reuse mode so draft steps 1+ skip the indexer op and read + # the indices that step 0 wrote into the shared buffer. + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(True) + + def on_multi_step_decode_end(self, num_reqs: int) -> None: + if self.share_mtp_topk_indices: + self.model.model.set_skip_topk(False) From 8bcc916a98a90822882455ab30aba02c409da2a6 Mon Sep 17 00:00:00 2001 From: Flora Feng <4florafeng@gmail.com> Date: Mon, 10 Aug 2026 19:00:08 -0400 Subject: [PATCH 14/30] [Bugfix] Fix DeepSeek V4/3.2 tokenizer vocab size overcount crashing guided decoding (#51727) Signed-off-by: sfeng33 <4florafeng@gmail.com> Co-authored-by: Wentao Ye <44945378+yewentao256@users.noreply.github.com> --- vllm/tokenizers/deepseek_v32.py | 6 ------ vllm/tokenizers/deepseek_v4.py | 5 ----- 2 files changed, 11 deletions(-) diff --git a/vllm/tokenizers/deepseek_v32.py b/vllm/tokenizers/deepseek_v32.py index b388f0579300..0e9aaa2cae5a 100644 --- a/vllm/tokenizers/deepseek_v32.py +++ b/vllm/tokenizers/deepseek_v32.py @@ -19,8 +19,6 @@ def get_deepseek_v32_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: dsv32_tokenizer = copy.copy(tokenizer) added_vocab = tokenizer.get_added_vocab() - added_vocab_size = len(added_vocab) - tokenizer_vocab_size = tokenizer.vocab_size class _DeepseekV32Tokenizer(tokenizer.__class__): # type: ignore def apply_chat_template( @@ -66,10 +64,6 @@ def apply_chat_template( def num_special_tokens_to_add(self) -> int: return len(self.encode("")) - def __len__(self) -> int: - # is an added token in DeepseekV32 tokenizer - return tokenizer_vocab_size + added_vocab_size - def get_added_vocab(self) -> dict[str, int]: return added_vocab.copy() diff --git a/vllm/tokenizers/deepseek_v4.py b/vllm/tokenizers/deepseek_v4.py index ce5b902c31e6..2a0db215b80f 100644 --- a/vllm/tokenizers/deepseek_v4.py +++ b/vllm/tokenizers/deepseek_v4.py @@ -19,8 +19,6 @@ def get_deepseek_v4_tokenizer(tokenizer: HfTokenizer) -> HfTokenizer: dsv4_tokenizer = copy.copy(tokenizer) added_vocab = tokenizer.get_added_vocab() - added_vocab_size = len(added_vocab) - tokenizer_vocab_size = tokenizer.vocab_size class _DeepseekV4Tokenizer(tokenizer.__class__): # type: ignore def apply_chat_template( @@ -78,9 +76,6 @@ def apply_chat_template( def num_special_tokens_to_add(self) -> int: return len(self.encode("")) - def __len__(self) -> int: - return tokenizer_vocab_size + added_vocab_size - def get_added_vocab(self) -> dict[str, int]: return added_vocab.copy() From 8977ea8895b19023f729db474d0f609e5bbac144 Mon Sep 17 00:00:00 2001 From: Myeongwon Kim <65876994+samuelkim7@users.noreply.github.com> Date: Tue, 11 Aug 2026 01:17:38 +0200 Subject: [PATCH 15/30] [Perf] Skip detokenization in offline beam search (#50333) Signed-off-by: samuelkim7 --- vllm/entrypoints/generate/beam_search/offline.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/vllm/entrypoints/generate/beam_search/offline.py b/vllm/entrypoints/generate/beam_search/offline.py index e1bceb5d23d4..7814d4d44e84 100644 --- a/vllm/entrypoints/generate/beam_search/offline.py +++ b/vllm/entrypoints/generate/beam_search/offline.py @@ -119,6 +119,7 @@ def beam_search( logprobs=2 * beam_width, max_tokens=1, temperature=temperature, + detokenize=False, skip_clone=True, # Internal beam search, safe to skip clone ) instances: list[BeamSearchInstance] = [] @@ -443,6 +444,7 @@ def _build_beam_sampling_params( logprobs=base_params.logprobs, max_tokens=1, temperature=base_params.temperature, + detokenize=False, allowed_token_ids=( allowed_ids if len(allowed_ids) <= _MAX_NUM_ALLOWED_TOKEN_IDS From 8a9f9f762a47e24ed8b48a8472f6757346cb30ae Mon Sep 17 00:00:00 2001 From: "Kevin H. Luu" Date: Mon, 10 Aug 2026 16:38:22 -0700 Subject: [PATCH 16/30] [CI] Parallelize release image publishing (#51735) Signed-off-by: khluu Co-authored-by: OpenAI Codex --- .buildkite/release-pipeline.yaml | 12 +- .buildkite/scripts/publish-release-images.sh | 307 ++++++++++--------- 2 files changed, 178 insertions(+), 141 deletions(-) diff --git a/.buildkite/release-pipeline.yaml b/.buildkite/release-pipeline.yaml index de6709547412..6b97bc57daa1 100644 --- a/.buildkite/release-pipeline.yaml +++ b/.buildkite/release-pipeline.yaml @@ -959,14 +959,22 @@ steps: - step: build-cpu-release-image-arm64 allow_failure: true - - label: "Publish release images to DockerHub" + - label: "Publish {{matrix}} release images to DockerHub" depends_on: - block-publish-release-images key: publish-release-images-dockerhub agents: queue: small_cpu_queue_release commands: - - "bash .buildkite/scripts/publish-release-images.sh" + - "bash .buildkite/scripts/publish-release-images.sh {{matrix}}" + matrix: + - "cuda-13-0" + - "cuda-12-9" + - "cuda-13-0-ubuntu-24-04" + - "cuda-12-9-ubuntu-24-04" + - "rocm" + - "xpu" + - "cpu" plugins: - docker-login#v3.0.0: username: vllmbot diff --git a/.buildkite/scripts/publish-release-images.sh b/.buildkite/scripts/publish-release-images.sh index 91b5c3ace1be..dd54b1717e10 100755 --- a/.buildkite/scripts/publish-release-images.sh +++ b/.buildkite/scripts/publish-release-images.sh @@ -2,12 +2,26 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project # -# Publish release Docker images from ECR to DockerHub. +# Publish a release Docker image family from ECR to DockerHub. # Pulls per-arch images, tags with latest and versioned tags, pushes them, # then creates and pushes multi-arch manifests. set -euo pipefail +TARGET="${1:-all}" +case "${TARGET}" in + cuda-13-0 | cuda-12-9 | cuda-13-0-ubuntu-24-04 | \ + cuda-12-9-ubuntu-24-04 | rocm | xpu | cpu | all) ;; + *) + echo "Usage: $0 {cuda-13-0|cuda-12-9|cuda-13-0-ubuntu-24-04|cuda-12-9-ubuntu-24-04|rocm|xpu|cpu|all}" + exit 2 + ;; +esac + +target_enabled() { + [ "${TARGET}" = "all" ] || [ "${TARGET}" = "$1" ] +} + RELEASE_VERSION=$(buildkite-agent meta-data get release-version --default "" | sed 's/^v//') if [ -z "${RELEASE_VERSION}" ]; then echo "ERROR: release-version metadata not set" @@ -15,12 +29,10 @@ if [ -z "${RELEASE_VERSION}" ]; then fi COMMIT="$BUILDKITE_COMMIT" -ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) echo "========================================" -echo "Publishing release images v${RELEASE_VERSION}" +echo "Publishing ${TARGET} release images v${RELEASE_VERSION}" echo " Commit: ${COMMIT}" -echo " ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" echo "========================================" # Login to ECR to pull staging images @@ -29,122 +41,137 @@ aws ecr-public get-login-password --region us-east-1 | \ # ---- CUDA (default: 13.0) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 -docker push vllm/vllm-openai:latest-x86_64 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker push vllm/vllm-openai:latest-aarch64 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 - -docker manifest rm vllm/vllm-openai:latest || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true -docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 -docker manifest push vllm/vllm-openai:latest -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} +if target_enabled cuda-13-0; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:latest-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai:latest-x86_64 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:latest-aarch64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker push vllm/vllm-openai:latest-aarch64 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + + docker manifest rm vllm/vllm-openai:latest || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai:latest vllm/vllm-openai:latest-x86_64 vllm/vllm-openai:latest-aarch64 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION} vllm/vllm-openai:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64 + docker manifest push vllm/vllm-openai:latest + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION} +fi # ---- CUDA 12.9 ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 -docker push vllm/vllm-openai:latest-x86_64-cu129 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 -docker push vllm/vllm-openai:latest-aarch64-cu129 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 - -docker manifest rm vllm/vllm-openai:latest-cu129 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true -docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 -docker manifest push vllm/vllm-openai:latest-cu129 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 +if target_enabled cuda-12-9; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:latest-x86_64-cu129 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + docker push vllm/vllm-openai:latest-x86_64-cu129 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:latest-aarch64-cu129 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker push vllm/vllm-openai:latest-aarch64-cu129 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + + docker manifest rm vllm/vllm-openai:latest-cu129 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129 || true + docker manifest create vllm/vllm-openai:latest-cu129 vllm/vllm-openai:latest-x86_64-cu129 vllm/vllm-openai:latest-aarch64-cu129 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129 + docker manifest push vllm/vllm-openai:latest-cu129 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129 +fi # ---- Ubuntu 24.04 (CUDA 13.0) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 -docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 -docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 - -docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true -docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 -docker manifest push vllm/vllm-openai:latest-ubuntu2404 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 +if target_enabled cuda-13-0-ubuntu-24-04; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + docker push vllm/vllm-openai:latest-x86_64-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker push vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + + docker manifest rm vllm/vllm-openai:latest-ubuntu2404 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 || true + docker manifest create vllm/vllm-openai:latest-ubuntu2404 vllm/vllm-openai:latest-x86_64-ubuntu2404 vllm/vllm-openai:latest-aarch64-ubuntu2404 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-ubuntu2404 + docker manifest push vllm/vllm-openai:latest-ubuntu2404 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-ubuntu2404 +fi # ---- Ubuntu 24.04 (CUDA 12.9) ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 -docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 - -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 -docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 - -docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true -docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true -docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 -docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 -docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 -docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 +if target_enabled cuda-12-9-ubuntu-24-04; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + docker push vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 + + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-aarch64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker push vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker push vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + + docker manifest rm vllm/vllm-openai:latest-cu129-ubuntu2404 || true + docker manifest rm vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 || true + docker manifest create vllm/vllm-openai:latest-cu129-ubuntu2404 vllm/vllm-openai:latest-x86_64-cu129-ubuntu2404 vllm/vllm-openai:latest-aarch64-cu129-ubuntu2404 + docker manifest create vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-x86_64-cu129-ubuntu2404 vllm/vllm-openai:v${RELEASE_VERSION}-aarch64-cu129-ubuntu2404 + docker manifest push vllm/vllm-openai:latest-cu129-ubuntu2404 + docker manifest push vllm/vllm-openai:v${RELEASE_VERSION}-cu129-ubuntu2404 +fi # ---- ROCm ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base +if target_enabled rocm; then + ROCM_BASE_CACHE_KEY=$(.buildkite/scripts/cache-rocm-base-wheels.sh key) + echo "ROCm base cache key: ${ROCM_BASE_CACHE_KEY}" + + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker push vllm/vllm-openai-rocm:latest -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:latest + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-rocm vllm/vllm-openai-rocm:v${RELEASE_VERSION} + docker push vllm/vllm-openai-rocm:latest + docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION} -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base -docker push vllm/vllm-openai-rocm:latest-base -docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:latest-base + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${ROCM_BASE_CACHE_KEY}-rocm-base vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base + docker push vllm/vllm-openai-rocm:latest-base + docker push vllm/vllm-openai-rocm:v${RELEASE_VERSION}-base +fi # ---- XPU ---- -docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu +if target_enabled xpu; then + docker pull public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 -docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 -docker push vllm/vllm-openai-xpu:latest-x86_64 -docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:latest-x86_64 + docker tag public.ecr.aws/q9t5s3a7/vllm-release-repo:${COMMIT}-x86_64-xpu vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai-xpu:latest-x86_64 + docker push vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 -docker manifest rm vllm/vllm-openai-xpu:latest || true -docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true -docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend -docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend -docker manifest push vllm/vllm-openai-xpu:latest -docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} + docker manifest rm vllm/vllm-openai-xpu:latest || true + docker manifest rm vllm/vllm-openai-xpu:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai-xpu:latest vllm/vllm-openai-xpu:latest-x86_64 --amend + docker manifest create vllm/vllm-openai-xpu:v${RELEASE_VERSION} vllm/vllm-openai-xpu:v${RELEASE_VERSION}-x86_64 --amend + docker manifest push vllm/vllm-openai-xpu:latest + docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} +fi # ---- CPU ---- # CPU images are behind separate block steps and may not have been built. @@ -153,44 +180,46 @@ docker manifest push vllm/vllm-openai-xpu:v${RELEASE_VERSION} # arch would leave `:latest-x86_64` pointing at the new release while the # `:latest` multi-arch manifest still resolves to the previous release. -CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} -CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} - -CPU_X86_AVAILABLE=false -CPU_ARM_AVAILABLE=false -docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true -docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true - -if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then - docker pull "${CPU_X86_TAG}" - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 - docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 - docker push vllm/vllm-openai-cpu:latest-x86_64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 - - docker pull "${CPU_ARM_TAG}" - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 - docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - docker push vllm/vllm-openai-cpu:latest-arm64 - docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - - docker manifest rm vllm/vllm-openai-cpu:latest || true - docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true - docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 - docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 - docker manifest push vllm/vllm-openai-cpu:latest - docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} -elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then - echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" -else - # Partial state: one arch built, the other did not. Fail loudly rather than - # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch) - # disagree on which release they point at. - echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})." - echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags." - echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended." - exit 1 +if target_enabled cpu; then + CPU_X86_TAG=public.ecr.aws/q9t5s3a7/vllm-cpu-release-repo:v${RELEASE_VERSION} + CPU_ARM_TAG=public.ecr.aws/q9t5s3a7/vllm-arm64-cpu-release-repo:v${RELEASE_VERSION} + + CPU_X86_AVAILABLE=false + CPU_ARM_AVAILABLE=false + docker manifest inspect "${CPU_X86_TAG}" >/dev/null 2>&1 && CPU_X86_AVAILABLE=true + docker manifest inspect "${CPU_ARM_TAG}" >/dev/null 2>&1 && CPU_ARM_AVAILABLE=true + + if [ "$CPU_X86_AVAILABLE" = "true" ] && [ "$CPU_ARM_AVAILABLE" = "true" ]; then + docker pull "${CPU_X86_TAG}" + docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:latest-x86_64 + docker tag "${CPU_X86_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + docker push vllm/vllm-openai-cpu:latest-x86_64 + docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 + + docker pull "${CPU_ARM_TAG}" + docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:latest-arm64 + docker tag "${CPU_ARM_TAG}" vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker push vllm/vllm-openai-cpu:latest-arm64 + docker push vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + + docker manifest rm vllm/vllm-openai-cpu:latest || true + docker manifest rm vllm/vllm-openai-cpu:v${RELEASE_VERSION} || true + docker manifest create vllm/vllm-openai-cpu:latest vllm/vllm-openai-cpu:latest-x86_64 vllm/vllm-openai-cpu:latest-arm64 + docker manifest create vllm/vllm-openai-cpu:v${RELEASE_VERSION} vllm/vllm-openai-cpu:v${RELEASE_VERSION}-x86_64 vllm/vllm-openai-cpu:v${RELEASE_VERSION}-arm64 + docker manifest push vllm/vllm-openai-cpu:latest + docker manifest push vllm/vllm-openai-cpu:v${RELEASE_VERSION} + elif [ "$CPU_X86_AVAILABLE" = "false" ] && [ "$CPU_ARM_AVAILABLE" = "false" ]; then + echo "WARNING: Neither CPU image found in ECR, skipping CPU publish (ensure block-cpu-release-image-build and block-arm64-cpu-release-image-build were unblocked and the builds finished pushing)" + else + # Partial state: one arch built, the other did not. Fail loudly rather than + # ship a Docker Hub state where `:latest-${arch}` and `:latest` (multi-arch) + # disagree on which release they point at. + echo "ERROR: Partial CPU build detected (x86_64=${CPU_X86_AVAILABLE}, arm64=${CPU_ARM_AVAILABLE})." + echo " Refusing to publish to avoid split-tag drift between per-arch and multi-arch tags." + echo " Re-run the missing CPU build and retry, or manually publish if a single-arch release is intended." + exit 1 + fi fi echo "" -echo "Successfully published release images for v${RELEASE_VERSION}" +echo "Successfully published ${TARGET} release images for v${RELEASE_VERSION}" From d8c70f22434afcbd6644aa43d3f23aecb6e5a09f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 11 Aug 2026 10:23:24 +1000 Subject: [PATCH 17/30] [Rust Frontend] Upgrade MiniJinja to 2.22 & remove method lookup workaround (#51235) Signed-off-by: Bugen Zhao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/Cargo.lock | 9 ++- rust/Cargo.toml | 4 +- rust/src/chat/Cargo.toml | 1 - rust/src/chat/src/renderer/hf/mod.rs | 16 ++---- rust/src/chat/src/renderer/hf/value.rs | 80 -------------------------- 5 files changed, 12 insertions(+), 98 deletions(-) delete mode 100644 rust/src/chat/src/renderer/hf/value.rs diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 3c75e01dc59e..092a79843124 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2402,9 +2402,9 @@ dependencies = [ [[package]] name = "minijinja" -version = "2.18.0" +version = "2.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "328251e58ad8e415be6198888fc207502727dc77945806421ab34f35bf012e7d" +checksum = "ef84a52be188a1d4124bd717903fdde96ca4705f2b56adfe2d91fcc57fcb6987" dependencies = [ "indexmap 2.13.0", "memo-map", @@ -2414,9 +2414,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.18.0" +version = "2.22.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "8c6302e47d2b51f9fc978268ff7f5a014de5caa2ad48440309fd10ee711480d7" +checksum = "fd6e279dc925840c2d9ebc1e9f85410611d01292561297463ba3e516c3ad94dd" dependencies = [ "minijinja", "serde", @@ -5520,7 +5520,6 @@ dependencies = [ "expect-test", "futures", "half", - "indexmap 2.13.0", "itertools 0.14.0", "llm-multimodal", "minijinja", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 09f55cf07cdf..405d6775bcf4 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -61,8 +61,8 @@ itertools = "0.14.0" libc = "0.2.177" llm-multimodal = { git = "https://github.com/smg-project/llm-multimodal", rev = "15adba5e025d8636ba4a334fb379b1371f6196a1", default-features = false, features = ["native-tls"] } mimalloc = "0.1.52" -minijinja = { version = "2.0", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } -minijinja-contrib = { version = "2.0", features = ["pycompat"] } +minijinja = { version = "2.22", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } +minijinja-contrib = { version = "2.22", features = ["pycompat"] } native-tls-vendored = { package = "native-tls", version = "0.2.18", features = ["vendored"] } ndarray = { version = "0.17", features = ["serde"] } openai-harmony = { package = "oss-harmony", git = "https://github.com/oss-harmony/harmony", tag = "v0.0.11", default-features = false } diff --git a/rust/src/chat/Cargo.toml b/rust/src/chat/Cargo.toml index 7bc66dd5987c..f7526b75d9ba 100644 --- a/rust/src/chat/Cargo.toml +++ b/rust/src/chat/Cargo.toml @@ -10,7 +10,6 @@ asynk-strim-attr.workspace = true easy-ext.workspace = true futures.workspace = true half.workspace = true -indexmap.workspace = true itertools.workspace = true llm-multimodal.workspace = true minijinja.workspace = true diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index cfbc3f924a34..901447dc27f4 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -16,7 +16,6 @@ use self::format::{ ChatTemplateContentFormat, ChatTemplateContentFormatOption as ContentFormatOption, }; use self::template::{CompiledChatTemplate, TemplateContext}; -use self::value::{TemplateValue, to_template_value}; use super::{ChatRenderer, RenderedPrompt, effective_template_kwargs}; use crate::error::Result; use crate::request::{ChatContent, ChatContentPart, ChatMessage, ChatRequest}; @@ -28,7 +27,6 @@ mod error; mod format; mod template; mod tojson; -mod value; pub use template::{load_chat_template, resolve_chat_template}; @@ -277,7 +275,7 @@ struct TemplateToolCall { #[derive(Debug, Serialize)] struct TemplateToolFunction { name: String, - arguments: TemplateValue, + arguments: JsonValue, } #[derive(Debug, Serialize)] @@ -291,7 +289,7 @@ pub(super) struct TemplateTool { struct TemplateToolDefinition { name: String, description: Option, - parameters: TemplateValue, + parameters: JsonValue, strict: Option, } @@ -384,8 +382,6 @@ fn to_template_tool_calls( error.as_report() )) })?; - let arguments = to_template_value(arguments); - tool_calls.push(TemplateToolCall { id: tool_call.id.clone(), r#type: "function", @@ -564,7 +560,7 @@ fn to_template_tools(tools: &[ChatTool]) -> Vec { function: TemplateToolDefinition { name: tool.name.clone(), description: tool.description.clone(), - parameters: to_template_value(tool.parameters.clone()), + parameters: tool.parameters.clone(), strict: tool.strict, }, }) @@ -1008,7 +1004,7 @@ mod tests { .apply_chat_template(&request) .unwrap(); - assert_eq!(rendered.prompt, Prompt::Text("|true".to_string())); + assert_eq!(rendered.prompt, Prompt::Text("|True".to_string())); } #[test] @@ -1093,7 +1089,7 @@ mod tests { let rendered = renderer.render(&request).unwrap().prompt; - assert_eq!(rendered, Prompt::Text("true|x".to_string())); + assert_eq!(rendered, Prompt::Text("True|x".to_string())); } #[test] @@ -1146,7 +1142,7 @@ mod tests { let rendered = renderer.render(&request).unwrap(); - assert_eq!(rendered.prompt, Prompt::Text("none|true".to_string())); + assert_eq!(rendered.prompt, Prompt::Text("none|True".to_string())); assert_eq!( rendered.effective_template_kwargs.get("reasoning_effort"), Some(&Value::String("none".to_string())) diff --git a/rust/src/chat/src/renderer/hf/value.rs b/rust/src/chat/src/renderer/hf/value.rs deleted file mode 100644 index 574e4a45928e..000000000000 --- a/rust/src/chat/src/renderer/hf/value.rs +++ /dev/null @@ -1,80 +0,0 @@ -// SPDX-License-Identifier: Apache-2.0 -// SPDX-FileCopyrightText: Copyright contributors to the vLLM project - -use std::sync::Arc; - -use indexmap::IndexMap; -use minijinja::value::{Enumerator, Object, ObjectExt, ObjectRepr}; -use minijinja::{Error as TemplateError, ErrorKind as TemplateErrorKind, State}; -use serde::Serialize; -use serde_json::Value as JsonValue; - -/// A wrapper around `minijinja::Value` that can be constructed with `to_template_value` and used -/// as a value in the chat template. -#[derive(Debug, Serialize)] -#[serde(transparent)] -pub(super) struct TemplateValue(minijinja::Value); - -pub(super) fn to_template_value(value: JsonValue) -> TemplateValue { - TemplateValue(match value { - JsonValue::Array(values) => values - .into_iter() - .map(to_template_value) - .map(|value| value.0) - .collect::(), - JsonValue::Object(values) => minijinja::Value::from_object(TemplateMap( - values - .into_iter() - .map(|(key, value)| (key, to_template_value(value).0)) - .collect(), - )), - // For primitive values, directly convert them to `minijinja::Value` using `from_serialize`. - value => minijinja::Value::from_serialize(value), - }) -} - -/// A custom map type that always returns `UnknownMethod` for method calls, so that pycompat can -/// always handle dict methods through the unknown-method callback. -/// -/// Use `IndexMap` to preserve the original key order when iterating. -/// -/// MiniJinja's default map can resolve a same-named field before Python dict methods. HF templates -/// commonly call `dict.items()`, which would fail if the map had an `items` field. -/// See issue: https://github.com/mitsuhiko/minijinja/issues/903 -#[derive(Debug)] -struct TemplateMap(IndexMap); - -impl Object for TemplateMap { - fn repr(self: &Arc) -> ObjectRepr { - ObjectRepr::Map - } - - fn get_value(self: &Arc, key: &minijinja::Value) -> Option { - self.0.get(key.as_str()?).cloned() - } - - fn get_value_by_str(self: &Arc, key: &str) -> Option { - self.0.get(key).cloned() - } - - fn enumerate(self: &Arc) -> Enumerator { - self.mapped_rev_enumerator(|this| { - Box::new(this.0.keys().map(|key| minijinja::Value::from(key.as_str()))) - }) - } - - fn enumerator_len(self: &Arc) -> Option { - Some(self.0.len()) - } - - fn call_method( - self: &Arc, - _state: &State<'_, '_>, - _method: &str, - _args: &[minijinja::Value], - ) -> std::result::Result { - // Always return `UnknownMethod` for method calls, - // so that pycompat can handle dict methods through the unknown-method callback. - Err(TemplateError::from(TemplateErrorKind::UnknownMethod)) - } -} From b2506d62aec7e6bccc5959b829221a7ae217abf3 Mon Sep 17 00:00:00 2001 From: Qiuyang Yue Date: Mon, 10 Aug 2026 17:39:04 -0700 Subject: [PATCH 18/30] [MM][CG][BugFix] Fix Ernie-4.5-VL encoder CG postprocess for multi-path outputs (#51461) Signed-off-by: qyYue1389 --- vllm/model_executor/models/ernie45_vl.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/vllm/model_executor/models/ernie45_vl.py b/vllm/model_executor/models/ernie45_vl.py index 74de1a0e9f45..f27e047b476c 100644 --- a/vllm/model_executor/models/ernie45_vl.py +++ b/vllm/model_executor/models/ernie45_vl.py @@ -1721,7 +1721,7 @@ def encoder_eager_forward( def postprocess_encoder_output( self, - output: torch.Tensor, + outputs: dict[str, torch.Tensor], indices: list[int], per_item_out_tokens: list[int], dest, @@ -1731,6 +1731,8 @@ def postprocess_encoder_output( # The graph output is the raw ViT features (pre-merge), padded to the # token budget. Run the resampler eagerly on the valid portion using # the actual batch grid_thw, then scatter the post-merge embeddings. + # Ernie only uses the single "default" encoder path. + output = outputs["default"] grid_thw = batch_mm_kwargs["image_grid_thw"].to(output.device) num_valid = int((grid_thw[:, 0] * grid_thw[:, 1] * grid_thw[:, 2]).sum()) image_embeds = self.resampler_model(output[:num_valid], grid_thw) From f1e921b6d6404ad1debd14ec0275b007c8c2ee11 Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 11 Aug 2026 11:30:38 +1000 Subject: [PATCH 19/30] [Rust Frontend] Support dynamic tools from developer messages (#51144) Signed-off-by: Bugen Zhao Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- rust/src/chat/src/error.rs | 11 +- rust/src/chat/src/lib.rs | 2 +- rust/src/chat/src/output/default/mod.rs | 6 +- .../chat/src/output/default/structural_tag.rs | 80 ++++-- rust/src/chat/src/output/harmony/mod.rs | 2 +- rust/src/chat/src/output/harmony/tests.rs | 25 +- .../src/renderer/deepseek_v32/encoding.rs | 14 +- .../fixtures/test_input_developer_tools.json | 32 +++ .../fixtures/test_output_developer_tools.txt | 37 +++ .../chat/src/renderer/deepseek_v32/tests.rs | 49 ++-- .../chat/src/renderer/deepseek_v4/encoding.rs | 6 +- .../fixtures/test_input_developer_tools.json | 31 +++ .../fixtures/test_output_developer_tools.txt | 28 +++ .../chat/src/renderer/deepseek_v4/tests.rs | 8 + rust/src/chat/src/renderer/harmony/mod.rs | 2 +- rust/src/chat/src/renderer/hf/mod.rs | 51 +++- rust/src/chat/src/renderer/inkling/mod.rs | 13 +- .../src/chat/src/renderer/kimi_k3/encoding.rs | 6 +- .../dynamic_only_tool_declare_input.json | 40 +++ .../dynamic_only_tool_declare_output.txt | 10 + rust/src/chat/src/renderer/kimi_k3/tests.rs | 15 +- rust/src/chat/src/renderer/test_utils.rs | 26 +- rust/src/chat/src/request.rs | 238 ++++++++++++++++-- rust/src/chat/tests/chat.rs | 10 +- rust/src/chat/tests/roundtrip.rs | 15 +- .../routes/openai/chat_completions/convert.rs | 63 +++-- .../routes/openai/chat_completions/types.rs | 111 +------- rust/src/server/src/routes/tokenize/types.rs | 13 +- 28 files changed, 686 insertions(+), 258 deletions(-) create mode 100644 rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input_developer_tools.json create mode 100644 rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_developer_tools.txt create mode 100644 rust/src/chat/src/renderer/deepseek_v4/fixtures/test_input_developer_tools.json create mode 100644 rust/src/chat/src/renderer/deepseek_v4/fixtures/test_output_developer_tools.txt create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_input.json create mode 100644 rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_output.txt diff --git a/rust/src/chat/src/error.rs b/rust/src/chat/src/error.rs index b032e9bf24cf..80e3a7f074e8 100644 --- a/rust/src/chat/src/error.rs +++ b/rust/src/chat/src/error.rs @@ -71,6 +71,12 @@ pub enum Error { StreamClosedBeforeTerminalOutput { request_id: String }, #[error("tool call stream state is inconsistent: {message}")] ToolCallStreamInvariant { message: String }, + #[error("duplicate tool name `{name}`")] + DuplicateToolName { name: String }, + #[error("tool_choice requires at least one available tool")] + ToolChoiceRequiresTools, + #[error("tool_choice function `{name}` was not found in the available tools")] + ToolChoiceFunctionNotFound { name: String }, #[error("failed to build structural tag: {message}")] StructuralTag { message: String }, #[error(transparent)] @@ -85,7 +91,10 @@ impl Error { /// Whether this error represents invalid user request parameters. pub fn is_request_validation_error(&self) -> bool { match self { - Self::PromptTooLong { .. } => true, + Self::PromptTooLong { .. } + | Self::DuplicateToolName { .. } + | Self::ToolChoiceRequiresTools + | Self::ToolChoiceFunctionNotFound { .. } => true, Self::Text(error) => error.is_request_validation_error(), Self::UnsupportedMultimodalRenderer | Self::UnsupportedMultimodalContent(_) diff --git a/rust/src/chat/src/lib.rs b/rust/src/chat/src/lib.rs index ca5c7a125a11..cf22722dd7f1 100644 --- a/rust/src/chat/src/lib.rs +++ b/rust/src/chat/src/lib.rs @@ -39,7 +39,7 @@ pub use renderer::{ }; pub use request::{ ChatContent, ChatContentPart, ChatMessage, ChatOptions, ChatRequest, ChatRole, ChatTool, - ChatToolChoice, GenerationPromptMode, ReasoningEffort, SamplingParams, + ChatToolChoice, GenerationPromptMode, ReasoningEffort, ResolvedToolContext, SamplingParams, }; pub use stream::{ChatEventStream, ChatEventStreamTrait, CollectedAssistantMessage}; pub use vllm_engine_core_client::protocol::multimodal::MmFeatures; diff --git a/rust/src/chat/src/output/default/mod.rs b/rust/src/chat/src/output/default/mod.rs index 4317e48fce2d..7c983fefd133 100644 --- a/rust/src/chat/src/output/default/mod.rs +++ b/rust/src/chat/src/output/default/mod.rs @@ -52,7 +52,7 @@ impl DefaultChatOutputProcessor { ) -> ChatResult { let parser = if tool_call_parser == reasoning_parser && let Some(parser) = Self::resolve_optional_unified_parser( - &request.tools, + request.tools(), model_id, tokenizer.clone(), tool_call_parser, @@ -62,7 +62,7 @@ impl DefaultChatOutputProcessor { let tool_parsing_enabled = request.tool_parsing_enabled(); let tool_parser = if tool_parsing_enabled { Some(Self::resolve_tool_parser( - &request.tools, + request.tools(), model_id, tool_call_parser, )?) @@ -82,7 +82,7 @@ impl DefaultChatOutputProcessor { Ok(Self { parser, - parallel_tool_calls: request.parallel_tool_calls, + parallel_tool_calls: request.parallel_tool_calls(), }) } diff --git a/rust/src/chat/src/output/default/structural_tag.rs b/rust/src/chat/src/output/default/structural_tag.rs index c0b31bd74564..2ef6f0868593 100644 --- a/rust/src/chat/src/output/default/structural_tag.rs +++ b/rust/src/chat/src/output/default/structural_tag.rs @@ -31,7 +31,7 @@ pub(super) fn apply_structural_tag_constraint( }; let tools = request - .tools + .tools() .iter() .map(|tool| { ToolParam::Function(FunctionToolParam::new(FunctionDefinition { @@ -67,13 +67,13 @@ pub(super) fn apply_structural_tag_constraint( /// /// Returns `None` if no structural tag constraints should be applied. fn structural_tag_tool_choice(request: &ChatRequest) -> Option { - if request.tools.is_empty() { + if request.tools().is_empty() { return None; } - match &request.tool_choice { + match request.tool_choice() { // For `Auto`, only apply the structural tag if there's at least one strict tool. - ChatToolChoice::Auto if request.tools.iter().any(|tool| tool.strict == Some(true)) => { + ChatToolChoice::Auto if request.tools().iter().any(|tool| tool.strict == Some(true)) => { Some(StructuralTagToolChoice::auto()) } ChatToolChoice::Auto | ChatToolChoice::None => None, @@ -92,6 +92,7 @@ mod tests { use vllm_parser::tool::{Qwen3CoderToolParser, Tool, ToolParser}; use super::*; + use crate::request::{ChatMessage, ResolvedToolContext}; fn chat_tool(name: &str, strict: Option) -> Tool { Tool { @@ -114,8 +115,8 @@ mod tests { fn request(tool_choice: ChatToolChoice, tools: Vec) -> ChatRequest { ChatRequest { - tool_choice, - tools, + tool_context: ResolvedToolContext::new(&[], tools, Some(tool_choice), true) + .expect("tool context should resolve"), ..ChatRequest::for_test() } } @@ -145,7 +146,7 @@ mod tests { #[test] fn auto_strict_tool_choice_builds_structural_tag() { let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", Some(true))]); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); @@ -155,10 +156,61 @@ mod tests { assert!(tag.to_string().contains("search")); } + #[test] + fn required_dynamic_tool_builds_structural_tag_from_effective_tools() { + let messages = vec![ChatMessage::developer( + "", + Some(vec![chat_tool("lookup", None)]), + )]; + let tool_context = + ResolvedToolContext::new(&messages, Vec::new(), Some(ChatToolChoice::Required), true) + .expect("dynamic tool context should resolve"); + let mut request = ChatRequest { + messages, + tool_context, + ..ChatRequest::for_test() + }; + let parser = qwen3_coder_parser(request.tools()); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) + .expect("structural tag should build"); + + let tag = structural_tag_value(&request); + assert!(tag.to_string().contains("lookup")); + } + + #[test] + fn required_initial_and_dynamic_tools_build_one_structural_tag() { + let messages = vec![ChatMessage::developer( + "", + Some(vec![chat_tool("lookup", None)]), + )]; + let tool_context = ResolvedToolContext::new( + &messages, + vec![chat_tool("search", None)], + Some(ChatToolChoice::Required), + true, + ) + .expect("tool context should resolve"); + let mut request = ChatRequest { + messages, + tool_context, + ..ChatRequest::for_test() + }; + let parser = qwen3_coder_parser(request.tools()); + + apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) + .expect("structural tag should build"); + + let tag = structural_tag_value(&request).to_string(); + assert!(tag.contains("search")); + assert!(tag.contains("lookup")); + } + #[test] fn auto_non_strict_tool_choice_skips_structural_tag() { let mut request = request(ChatToolChoice::Auto, vec![chat_tool("search", None)]); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag decision should succeed"); @@ -173,7 +225,7 @@ mod tests { backend: StructuredOutputBackend::Xgrammar, ..StructuredOutputsParams::json(json!({"type": "object"})) }); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); @@ -188,7 +240,7 @@ mod tests { #[test] fn required_tool_choice_builds_structural_tag_without_strict_tools() { let mut request = request(ChatToolChoice::Required, vec![chat_tool("search", None)]); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); @@ -205,7 +257,7 @@ mod tests { backend: StructuredOutputBackend::Xgrammar, ..StructuredOutputsParams::json_object() }); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); @@ -225,7 +277,7 @@ mod tests { }, vec![chat_tool("search", None), chat_tool("lookup", None)], ); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag should build"); @@ -238,7 +290,7 @@ mod tests { #[test] fn none_tool_choice_skips_structural_tag() { let mut request = request(ChatToolChoice::None, vec![chat_tool("search", Some(true))]); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag decision should succeed"); @@ -253,7 +305,7 @@ mod tests { backend: StructuredOutputBackend::Xgrammar, ..StructuredOutputsParams::json_object() }); - let parser = qwen3_coder_parser(&request.tools); + let parser = qwen3_coder_parser(request.tools()); apply_structural_tag_constraint(&mut request, parser.structural_tag_builder()) .expect("structural tag decision should succeed"); diff --git a/rust/src/chat/src/output/harmony/mod.rs b/rust/src/chat/src/output/harmony/mod.rs index a84687882b10..14206751bb49 100644 --- a/rust/src/chat/src/output/harmony/mod.rs +++ b/rust/src/chat/src/output/harmony/mod.rs @@ -75,7 +75,7 @@ impl HarmonyChatOutputProcessor { Ok(Self { encoding: harmony_encoding()?, tool_calls_enabled: request.tool_parsing_enabled(), - parallel_tool_calls: request.parallel_tool_calls, + parallel_tool_calls: request.parallel_tool_calls(), }) } } diff --git a/rust/src/chat/src/output/harmony/tests.rs b/rust/src/chat/src/output/harmony/tests.rs index 9b9544bdc6bd..e798f092694c 100644 --- a/rust/src/chat/src/output/harmony/tests.rs +++ b/rust/src/chat/src/output/harmony/tests.rs @@ -10,7 +10,7 @@ use vllm_text::output::{DecodedLogprobs, DecodedPositionLogprobs, DecodedTextEve use super::*; use crate::output::ChatOutputProcessor; -use crate::request::{ChatRequest, ChatTool, ChatToolChoice}; +use crate::request::{ChatRequest, ChatTool, ChatToolChoice, ResolvedToolContext}; use crate::{AssistantMessageExt, ChatEvent, FinishReason}; fn assistant_prefix() -> Vec { @@ -72,18 +72,19 @@ async fn collect_events( } fn request_with_tools() -> ChatRequest { + let tools = vec![ChatTool { + name: "get_weather".to_string(), + description: Some("Get weather".to_string()), + parameters: serde_json::json!({ + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"] + }), + strict: None, + }]; ChatRequest { - tool_choice: ChatToolChoice::Auto, - tools: vec![ChatTool { - name: "get_weather".to_string(), - description: Some("Get weather".to_string()), - parameters: serde_json::json!({ - "type": "object", - "properties": {"city": {"type": "string"}}, - "required": ["city"] - }), - strict: None, - }], + tool_context: ResolvedToolContext::new(&[], tools, Some(ChatToolChoice::Auto), true) + .expect("tool context should resolve"), ..ChatRequest::for_test() } } diff --git a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs index 86925c87c4c3..3ccf4d11db7f 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/encoding.rs @@ -2,6 +2,9 @@ // SPDX-FileCopyrightText: Copyright contributors to the vLLM project //! DeepSeek V3.2 prompt renderer. +//! +//! Official Python reference: +//! use std::collections::{HashMap, HashSet}; use std::fmt::Write as _; @@ -48,15 +51,20 @@ pub(super) fn render_request(request: &ChatRequest) -> Result { request.messages.last().map(ChatMessage::role), Some(ChatRole::User | ChatRole::Developer) ); - let render_offset = isize::from(request.tool_parsing_enabled()); + let initial_tools = if request.tool_parsing_enabled() { + request.initial_tools() + } else { + &[] + }; + let render_offset = isize::from(!initial_tools.is_empty()); let last_user_render_index = find_last_user_render_index(request.messages.as_slice(), render_offset); let last_user_actual_index = find_last_user_actual_index(request.messages.as_slice()); let continue_final_message = request.chat_options.continue_final_message(); let mut prompt = String::from(BOS_TOKEN); - if request.tool_parsing_enabled() { - render_system_message(&mut prompt, None, &request.tools)?; + if !initial_tools.is_empty() { + render_system_message(&mut prompt, None, initial_tools)?; } for (message_index, message) in request.messages.iter().enumerate() { diff --git a/rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input_developer_tools.json b/rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input_developer_tools.json new file mode 100644 index 000000000000..088c5fa8ef40 --- /dev/null +++ b/rust/src/chat/src/renderer/deepseek_v32/fixtures/test_input_developer_tools.json @@ -0,0 +1,32 @@ +{ + "messages": [ + { + "role": "user", + "content": "question" + }, + { + "role": "developer", + "content": "policy", + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object" + } + } + } + ] + }, + { + "role": "assistant", + "reasoning_content": "thought", + "content": "answer" + } + ], + "add_generation_prompt": false, + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_developer_tools.txt b/rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_developer_tools.txt new file mode 100644 index 000000000000..b3e732dbd5ea --- /dev/null +++ b/rust/src/chat/src/renderer/deepseek_v32/fixtures/test_output_developer_tools.txt @@ -0,0 +1,37 @@ +<｜begin▁of▁sentence｜><｜User｜>question<｜Assistant｜><｜User｜> + +## Tools + +You have access to a set of tools you can use to answer the user's question. +You can invoke functions by writing a "<｜DSML｜function_calls>" block like the following as part of your reply to the user: +<｜DSML｜function_calls> +<｜DSML｜invoke name="$FUNCTION_NAME"> +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<｜DSML｜invoke name="$FUNCTION_NAME2"> +... + + + +String and scalar parameters should be specified as is without any escaping or quotes, while lists and objects should use JSON format. The "string" attribute should be set to "true" for string type parameters and "false" for other types (numbers, booleans, arrays, objects). + +If the thinking_mode is enabled, then after function results you should strongly consider outputting a thinking block. Here is an example: + +<｜DSML｜function_calls> +... + + + +... + + +...thinking about results + +Here are the functions available in JSONSchema format: + +{"name": "lookup", "parameters": {"type": "object"}} + + + +# The user's message is: policy<｜Assistant｜>thoughtanswer<｜end▁of▁sentence｜> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/deepseek_v32/tests.rs b/rust/src/chat/src/renderer/deepseek_v32/tests.rs index 7960927219b3..3f93e3653efd 100644 --- a/rust/src/chat/src/renderer/deepseek_v32/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v32/tests.rs @@ -13,6 +13,7 @@ use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::renderer::test_utils::{FixtureRequestOptions, fixture_chat_request}; use crate::request::{ ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, GenerationPromptMode, + ResolvedToolContext, }; use crate::{ChatRenderer, ChatRole}; @@ -85,6 +86,14 @@ fn renders_vllm_parity_prompt_for_request_level_tools_fixture() { ); } +#[test] +fn renders_developer_tools_like_hf_python() { + assert_fixture( + "test_input_developer_tools.json", + expect_file!["fixtures/test_output_developer_tools.txt"], + ); +} + #[test] fn renders_official_search_fixture_without_date() { assert_fixture( @@ -103,27 +112,29 @@ fn renders_official_search_fixture_with_date() { #[test] fn request_level_tools_are_lowered_as_synthetic_leading_system_message() { + let messages = vec![ + ChatMessage::system("System prompt."), + ChatMessage::text(ChatRole::User, "Hello"), + ]; + let tools = vec![ChatTool { + name: "lookup".to_string(), + description: Some("Look things up".to_string()), + parameters: json!({ + "type": "object", + "properties": { + "query": { + "type": "string" + } + }, + "required": ["query"] + }), + strict: None, + }]; let mut request = ChatRequest { request_id: "deepseek-v32-tools".to_string(), - messages: vec![ - ChatMessage::system("System prompt."), - ChatMessage::text(ChatRole::User, "Hello"), - ], - tools: vec![ChatTool { - name: "lookup".to_string(), - description: Some("Look things up".to_string()), - parameters: json!({ - "type": "object", - "properties": { - "query": { - "type": "string" - } - }, - "required": ["query"] - }), - strict: None, - }], - tool_choice: ChatToolChoice::Auto, + tool_context: ResolvedToolContext::new(&messages, tools, Some(ChatToolChoice::Auto), true) + .expect("tool context should resolve"), + messages, ..ChatRequest::for_test() }; request diff --git a/rust/src/chat/src/renderer/deepseek_v4/encoding.rs b/rust/src/chat/src/renderer/deepseek_v4/encoding.rs index db86dd5b8253..c0bea8260a1c 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/encoding.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/encoding.rs @@ -3,8 +3,8 @@ //! DeepSeek V4 prompt renderer. //! -//! Original Python implementation: -//! +//! Official Python reference: +//! use std::collections::HashMap; use std::fmt::Write as _; @@ -156,7 +156,7 @@ fn resolve_thinking_options(request: &ChatRequest) -> Result<(ThinkingMode, &'st /// Return request-level tools only when native tool parsing is enabled. fn request_tools(request: &ChatRequest) -> &[ChatTool] { if request.tool_parsing_enabled() { - request.tools.as_slice() + request.initial_tools() } else { &[] } diff --git a/rust/src/chat/src/renderer/deepseek_v4/fixtures/test_input_developer_tools.json b/rust/src/chat/src/renderer/deepseek_v4/fixtures/test_input_developer_tools.json new file mode 100644 index 000000000000..a602f3efe607 --- /dev/null +++ b/rust/src/chat/src/renderer/deepseek_v4/fixtures/test_input_developer_tools.json @@ -0,0 +1,31 @@ +{ + "messages": [ + { + "role": "user", + "content": "question" + }, + { + "role": "developer", + "content": "policy", + "tools": [ + { + "type": "function", + "function": { + "name": "lookup", + "parameters": { + "type": "object" + } + } + } + ] + }, + { + "role": "assistant", + "content": "answer" + } + ], + "add_generation_prompt": false, + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/deepseek_v4/fixtures/test_output_developer_tools.txt b/rust/src/chat/src/renderer/deepseek_v4/fixtures/test_output_developer_tools.txt new file mode 100644 index 000000000000..85cbe91dd0fd --- /dev/null +++ b/rust/src/chat/src/renderer/deepseek_v4/fixtures/test_output_developer_tools.txt @@ -0,0 +1,28 @@ +<｜begin▁of▁sentence｜><｜User｜>question<｜User｜>policy + +## Tools + +You have access to a set of tools to help answer the user's question. You can invoke tools by writing a "<｜DSML｜tool_calls>" block like the following: + +<｜DSML｜tool_calls> +<｜DSML｜invoke name="$TOOL_NAME"> +<｜DSML｜parameter name="$PARAMETER_NAME" string="true|false">$PARAMETER_VALUE +... + +<｜DSML｜invoke name="$TOOL_NAME2"> +... + + + +String parameters should be specified as is and set `string="true"`. For all other types (numbers, booleans, arrays, objects), pass the value in JSON format and set `string="false"`. + +If thinking_mode is enabled (triggered by ), you MUST output your complete reasoning inside ... BEFORE any tool calls or final response. + +Otherwise, output directly after with tool calls or final response. + +### Available Tool Schemas + +{"name": "lookup", "parameters": {"type": "object"}} + +You MUST strictly follow the above defined tool name and parameter schemas to invoke tool calls. +<｜Assistant｜>answer<｜end▁of▁sentence｜> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/deepseek_v4/tests.rs b/rust/src/chat/src/renderer/deepseek_v4/tests.rs index 5c92e0f0ae0a..ee054a3ac96b 100644 --- a/rust/src/chat/src/renderer/deepseek_v4/tests.rs +++ b/rust/src/chat/src/renderer/deepseek_v4/tests.rs @@ -63,6 +63,14 @@ fn renders_v4_fixture_2_multi_turn_drop_thinking() { ); } +#[test] +fn renders_developer_tools_like_hf_python() { + assert_fixture( + "test_input_developer_tools.json", + expect_file!["fixtures/test_output_developer_tools.txt"], + ); +} + #[test] fn reasoning_effort_max_adds_prefix_when_thinking_is_enabled() { let mut request = ChatRequest { diff --git a/rust/src/chat/src/renderer/harmony/mod.rs b/rust/src/chat/src/renderer/harmony/mod.rs index 49375abe8a51..b07d8458d739 100644 --- a/rust/src/chat/src/renderer/harmony/mod.rs +++ b/rust/src/chat/src/renderer/harmony/mod.rs @@ -216,7 +216,7 @@ fn preamble_tool_descriptions( ) -> Vec { let mut tools = Vec::new(); if request.tool_parsing_enabled() { - tools.extend(to_tool_descriptions(&request.tools)); + tools.extend(to_tool_descriptions(request.initial_tools())); } if let Some(leading_developer_tools) = leading_developer_tools { tools.extend(to_tool_descriptions(leading_developer_tools)); diff --git a/rust/src/chat/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index 901447dc27f4..df0d94331a8b 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -179,7 +179,7 @@ impl HfChatRenderer { None }; - let tools = request.tool_parsing_enabled().then(|| to_template_tools(&request.tools)); + let tools = request.tool_parsing_enabled().then(|| to_template_tools(request.tools())); trace!( message_count = messages.len(), content_format = ?effective_template.content_format(), @@ -579,7 +579,7 @@ mod tests { use super::{ChatTemplateContentFormatOption, HfChatRenderer, MultimodalRenderInfo}; use crate::request::{ ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, - GenerationPromptMode, ReasoningEffort, + GenerationPromptMode, ReasoningEffort, ResolvedToolContext, }; use crate::{AssistantContentBlock, ChatRenderer, Error, Result}; @@ -587,13 +587,48 @@ mod tests { const QWEN3_5_0_8B_TEMPLATE: &str = include_str!("../../../tests/templates/qwen35.jinja"); fn sample_request(messages: Vec) -> ChatRequest { + let tool_context = ResolvedToolContext::new(&messages, Vec::new(), None, true).unwrap(); ChatRequest { messages, + tool_context, request_id: "render-test".to_string(), ..ChatRequest::for_test() } } + #[test] + fn dynamic_tools_are_exposed_as_effective_template_tools() { + let messages = vec![ + ChatMessage::developer( + "", + Some(vec![ChatTool { + name: "lookup".to_string(), + description: None, + parameters: serde_json::json!({"type": "object"}), + strict: None, + }]), + ), + ChatMessage::user("hello"), + ]; + let tool_context = ResolvedToolContext::new(&messages, Vec::new(), None, true).unwrap(); + let request = ChatRequest { + messages, + tool_context, + request_id: "render-test".to_string(), + ..ChatRequest::for_test() + }; + + let rendered = render( + Some("{{ messages|length }}:{{ messages[0].role }}:{{ messages[0].tools[0].function.name }}:{{ tools[0].function.name }}"), + &request, + ) + .unwrap(); + + assert_eq!(rendered, "2:developer:lookup:lookup"); + assert_eq!(request.tool_choice(), &ChatToolChoice::Auto); + assert_eq!(request.tools().len(), 1); + } + fn render(template: Option<&str>, request: &ChatRequest) -> Result { HfChatRenderer::new( template.map(str::to_owned), @@ -897,7 +932,7 @@ mod tests { } #[test] - fn chat_template_exposes_developer_tools() { + fn chat_template_preserves_developer_tools() { let request = sample_request(vec![ChatMessage::developer( "policy", Some(vec![ChatTool { @@ -1218,7 +1253,7 @@ mod tests { #[test] fn chat_template_exposes_tools_to_templates_when_auto_enabled() { let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); - request.tools = vec![ChatTool { + let tools = vec![ChatTool { name: "get_weather".to_string(), description: Some("Get weather".to_string()), parameters: serde_json::json!({ @@ -1228,7 +1263,13 @@ mod tests { }), strict: None, }]; - request.tool_choice = ChatToolChoice::Auto; + request.tool_context = crate::request::ResolvedToolContext::new( + &request.messages, + tools, + Some(ChatToolChoice::Auto), + true, + ) + .expect("tool context should resolve"); let rendered = render( Some("{{ tools[0].function.name }}|{{ tools[0].function.parameters.required[0] }}"), diff --git a/rust/src/chat/src/renderer/inkling/mod.rs b/rust/src/chat/src/renderer/inkling/mod.rs index 8bf1a204fbf2..a33e3414f2db 100644 --- a/rust/src/chat/src/renderer/inkling/mod.rs +++ b/rust/src/chat/src/renderer/inkling/mod.rs @@ -383,18 +383,7 @@ fn rendered_tools(request: &ChatRequest) -> Vec<&ChatTool> { return Vec::new(); } - let mut tools = Vec::with_capacity(request.tools.len()); - tools.extend(request.tools.iter()); - for message in &request.messages { - if let ChatMessage::Developer { - tools: Some(local_tools), - .. - } = message - { - tools.extend(local_tools.iter()); - } - } - tools + request.tools().iter().collect() } fn tool_call_json(tool_call: &AssistantToolCall) -> Result { diff --git a/rust/src/chat/src/renderer/kimi_k3/encoding.rs b/rust/src/chat/src/renderer/kimi_k3/encoding.rs index b2958b3d8492..3ff9242ff4bf 100644 --- a/rust/src/chat/src/renderer/kimi_k3/encoding.rs +++ b/rust/src/chat/src/renderer/kimi_k3/encoding.rs @@ -191,7 +191,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - } flush_tool_run(&mut out, &mut pending_tool_run, &tool_call_id_index)?; - match &request.tool_choice { + match request.tool_choice() { ChatToolChoice::Required => { write_internal_system( &mut out, @@ -202,7 +202,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - } // Emit only when tools are present: Rust defaults tool_choice to None // for tool-free requests, which must not inject a tool-choice message. - ChatToolChoice::None if !request.tools.is_empty() => { + ChatToolChoice::None if !request.tools().is_empty() => { write_internal_system( &mut out, "tool-choice", @@ -226,7 +226,7 @@ pub(super) fn render_request(request: &ChatRequest, tokenizer: &dyn Tokenizer) - fn request_tools(request: &ChatRequest) -> &[ChatTool] { // Declare tools whenever the request carries them. K3 tool-declare is // independent of tool_choice; tool_choice only injects control messages. - request.tools.as_slice() + request.initial_tools() } fn thinking_enabled(request: &ChatRequest) -> Result { diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_input.json b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_input.json new file mode 100644 index 000000000000..da1138e593be --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_input.json @@ -0,0 +1,40 @@ +{ + "messages": [ + { + "role": "user", + "content": "Load the calculator." + }, + { + "role": "developer", + "tools": [ + { + "type": "function", + "function": { + "name": "calc", + "description": "Evaluate an expression", + "parameters": { + "type": "object", + "properties": { + "expression": { + "type": "string" + } + }, + "required": [ + "expression" + ] + } + } + } + ] + }, + { + "role": "user", + "content": "What is 2 + 2?" + } + ], + "add_generation_prompt": true, + "tool_choice": "none", + "template_kwargs": { + "thinking": true + } +} diff --git a/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_output.txt b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_output.txt new file mode 100644 index 000000000000..376de0777990 --- /dev/null +++ b/rust/src/chat/src/renderer/kimi_k3/fixtures/dynamic_only_tool_declare_output.txt @@ -0,0 +1,10 @@ +<|open|>message role="system" type="thinking-effort"<|sep|>`thinking_effort` guides on how much to think in your thinking channel (not including the response channel), supported values include `low`, `medium`, `high`, and `max`. +Now the system is invoked with `thinking_effort=max`.<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>Load the calculator.<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-declare"<|sep|>## New Tools Available +The system dynamically extends the toolset via lazy-loading. +You have access to all existing and extended tools. +Here are the specs for the extended tools. + +```json +[{"function":{"description":"Evaluate an expression","name":"calc","parameters":{"properties":{"expression":{"type":"string"}},"required":["expression"],"type":"object"}},"type":"function"}] +```<|close|>message<|sep|><|end_of_msg|><|open|>message role="user"<|sep|>What is 2 + 2?<|close|>message<|sep|><|end_of_msg|><|open|>message role="system" type="tool-choice"<|sep|>The system is invoked with `tool_choice=none`. +You MUST NOT call any tools in the next message.<|close|>message<|sep|><|end_of_msg|><|open|>message role="assistant"<|sep|><|open|>think<|sep|> \ No newline at end of file diff --git a/rust/src/chat/src/renderer/kimi_k3/tests.rs b/rust/src/chat/src/renderer/kimi_k3/tests.rs index 8b2799cdb3da..94a3ad4cfef2 100644 --- a/rust/src/chat/src/renderer/kimi_k3/tests.rs +++ b/rust/src/chat/src/renderer/kimi_k3/tests.rs @@ -92,10 +92,15 @@ fn golden_dynamic_system_tool_declare() { assert_golden("dynamic_system_tool_declare"); } +#[test] +fn golden_dynamic_only_tool_declare() { + assert_golden("dynamic_only_tool_declare"); +} + #[test] fn tool_declare_omits_absent_fields_and_preserves_parameter_nulls() { let mut request = crate::request::ChatRequest::for_test(); - request.tools = vec![ChatTool { + let tools = vec![ChatTool { name: "lookup".to_string(), description: None, parameters: json!({ @@ -109,6 +114,9 @@ fn tool_declare_omits_absent_fields_and_preserves_parameter_nulls() { }), strict: None, }]; + request.tool_context = + crate::request::ResolvedToolContext::new(&request.messages, tools, None, true) + .expect("tool context should resolve"); let rendered = render_request(&request); @@ -122,12 +130,15 @@ fn tool_declare_omits_absent_fields_and_preserves_parameter_nulls() { #[test] fn tool_declare_preserves_present_optional_fields() { let mut request = crate::request::ChatRequest::for_test(); - request.tools = vec![ChatTool { + let tools = vec![ChatTool { name: "lookup".to_string(), description: Some("Look up a record".to_string()), parameters: json!({"type": "object"}), strict: Some(false), }]; + request.tool_context = + crate::request::ResolvedToolContext::new(&request.messages, tools, None, true) + .expect("tool context should resolve"); let rendered = render_request(&request); diff --git a/rust/src/chat/src/renderer/test_utils.rs b/rust/src/chat/src/renderer/test_utils.rs index 0a4edc813798..1090a4039f5e 100644 --- a/rust/src/chat/src/renderer/test_utils.rs +++ b/rust/src/chat/src/renderer/test_utils.rs @@ -11,7 +11,7 @@ use serde_json::Value; use crate::event::{AssistantContentBlock, AssistantToolCall}; use crate::request::{ ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatTool, ChatToolChoice, - GenerationPromptMode, ReasoningEffort, + GenerationPromptMode, ReasoningEffort, ResolvedToolContext, }; /// Options for constructing a [`ChatRequest`] from a fixture file. @@ -158,22 +158,20 @@ struct FixtureToolCallFunction { impl FixtureRequest { fn into_chat_request(self, options: FixtureRequestOptions) -> ChatRequest { let tools = to_chat_tools(&self.tools); - let tool_choice = self.tool_choice.unwrap_or(if tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }); + + let messages = self + .messages + .into_iter() + .enumerate() + .map(|(index, message)| fixture_message_to_chat_message(index, message)) + .collect::>(); + let tool_context = ResolvedToolContext::new(&messages, tools, self.tool_choice, true) + .expect("fixture tool context should resolve"); let mut request = ChatRequest { request_id: "renderer-fixture".to_string(), - messages: self - .messages - .into_iter() - .enumerate() - .map(|(index, message)| fixture_message_to_chat_message(index, message)) - .collect(), - tools, - tool_choice, + messages, + tool_context, ..ChatRequest::for_test() }; diff --git a/rust/src/chat/src/request.rs b/rust/src/chat/src/request.rs index f4cc252d81e2..44d833eca252 100644 --- a/rust/src/chat/src/request.rs +++ b/rust/src/chat/src/request.rs @@ -1,7 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project -use std::collections::HashMap; +use std::collections::{HashMap, HashSet}; use llm_multimodal::ImageDetail; use serde::{Deserialize, Serialize}; @@ -248,6 +248,16 @@ impl ChatMessage { } } + /// Return message-scoped tool declarations carried by this message. + pub fn declared_tools(&self) -> Option<&[ChatTool]> { + match self { + Self::Developer { + tools: Some(tools), .. + } if !tools.is_empty() => Some(tools), + _ => None, + } + } + /// Construct one chat message with user role. pub fn user(content: impl Into) -> Self { Self::User { @@ -450,6 +460,93 @@ pub enum ChatToolChoice { }, } +/// Resolved tool state shared by rendering, output parsing, and constraints. +#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] +pub struct ResolvedToolContext { + /// Request-level tools rendered before the ordered chat history. + pub initial_tools: Vec, + /// Request-level and message-scoped tools available for this generation. + pub effective_tools: Vec, + /// Resolved tool-choice behavior for this request. + pub tool_choice: ChatToolChoice, + /// Whether more than one parsed tool call may be surfaced northbound. + pub parallel_tool_calls: bool, +} + +impl ResolvedToolContext { + /// Resolve one complete tool context from request and message declarations. + pub fn new( + messages: &[ChatMessage], + initial_tools: Vec, + requested_tool_choice: Option, + parallel_tool_calls: bool, + ) -> Result { + let dynamic_tool_count = messages + .iter() + .filter_map(ChatMessage::declared_tools) + .map(<[ChatTool]>::len) + .sum::(); + let mut effective_tools = Vec::with_capacity(initial_tools.len() + dynamic_tool_count); + let mut names = HashSet::with_capacity(effective_tools.capacity()); + + for tool in initial_tools + .iter() + .chain(messages.iter().filter_map(ChatMessage::declared_tools).flatten()) + { + if !names.insert(tool.name.clone()) { + return Err(Error::DuplicateToolName { + name: tool.name.clone(), + }); + } + effective_tools.push(tool.clone()); + } + + let tool_choice = requested_tool_choice.unwrap_or({ + if effective_tools.is_empty() { + ChatToolChoice::None + } else { + ChatToolChoice::Auto + } + }); + + match &tool_choice { + ChatToolChoice::Auto | ChatToolChoice::Required if effective_tools.is_empty() => { + return Err(Error::ToolChoiceRequiresTools); + } + ChatToolChoice::Function { name } if !names.contains(name) => { + return Err(Error::ToolChoiceFunctionNotFound { name: name.clone() }); + } + ChatToolChoice::None + | ChatToolChoice::Auto + | ChatToolChoice::Required + | ChatToolChoice::Function { .. } => {} + } + + Ok(Self { + initial_tools, + effective_tools, + tool_choice, + parallel_tool_calls, + }) + } + + /// Return whether tool output parsing is active for this request. + pub fn parsing_enabled(&self) -> bool { + !matches!(self.tool_choice, ChatToolChoice::None) && !self.effective_tools.is_empty() + } +} + +impl Default for ResolvedToolContext { + fn default() -> Self { + Self { + initial_tools: Vec::new(), + effective_tools: Vec::new(), + tool_choice: ChatToolChoice::None, + parallel_tool_calls: true, + } + } +} + /// One chat request ready to be rendered into a prompt and lowered into a /// generate request. #[derive(Debug, Clone, PartialEq, Serialize, Deserialize)] @@ -462,14 +559,8 @@ pub struct ChatRequest { pub sampling_params: SamplingParams, /// Chat-specific rendering options. pub chat_options: ChatOptions, - /// Function tools made available to the model for this request. - pub tools: Vec, - /// Tool-choice behavior for this request. - pub tool_choice: ChatToolChoice, - /// Whether the model may return more than one tool call per response. - /// - /// When `false`, only the first parsed tool call is surfaced northbound. - pub parallel_tool_calls: bool, + /// Resolved request-level and message-scoped tool behavior. + pub tool_context: ResolvedToolContext, /// Text decode options for incremental detokenization. pub decode_options: TextDecodeOptions, /// Whether to emit intermediate northbound content deltas before the @@ -507,9 +598,7 @@ impl ChatRequest { messages: vec![ChatMessage::text(ChatRole::User, "test")], sampling_params: SamplingParams::default(), chat_options: ChatOptions::default(), - tools: Vec::new(), - tool_choice: ChatToolChoice::None, - parallel_tool_calls: true, + tool_context: ResolvedToolContext::default(), decode_options: TextDecodeOptions::default(), intermediate: true, priority: 0, @@ -547,10 +636,33 @@ impl ChatRequest { self.messages.iter().any(ChatMessage::has_multimodal) } + /// Return every tool available for this generation. + pub fn tools(&self) -> &[ChatTool] { + &self.tool_context.effective_tools + } + + /// Return request-level tools before message-scoped declarations. + /// + /// Renderers should use [`Self::tools`] unless their prompt format places + /// initial and message-scoped declarations differently. + pub fn initial_tools(&self) -> &[ChatTool] { + &self.tool_context.initial_tools + } + + /// Return the resolved tool-choice behavior. + pub fn tool_choice(&self) -> &ChatToolChoice { + &self.tool_context.tool_choice + } + + /// Return whether more than one parsed tool call may be surfaced. + pub fn parallel_tool_calls(&self) -> bool { + self.tool_context.parallel_tool_calls + } + /// Return true if this request should enable tool parsing based on the tool /// choice and tool list. pub(crate) fn tool_parsing_enabled(&self) -> bool { - !matches!(self.tool_choice, ChatToolChoice::None) && !self.tools.is_empty() + self.tool_context.parsing_enabled() } /// Return the request-level thinking toggle when explicitly requested. @@ -605,7 +717,10 @@ impl ChatRole { mod tests { use serde_json::{json, to_value}; - use super::{ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool}; + use super::{ + ChatContent, ChatContentPart, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, + ResolvedToolContext, + }; use crate::Error; use crate::event::AssistantContentBlock; @@ -702,6 +817,101 @@ mod tests { assert_eq!(decoded, message); } + fn chat_tool(name: &str) -> ChatTool { + ChatTool { + name: name.to_string(), + description: None, + parameters: json!({"type": "object"}), + strict: None, + } + } + + #[test] + fn resolved_tool_context_defaults_dynamic_only_tools_to_auto() { + let messages = vec![ + ChatMessage::user("hello"), + ChatMessage::developer("", Some(vec![chat_tool("lookup")])), + ]; + + let context = ResolvedToolContext::new(&messages, Vec::new(), None, false).unwrap(); + + assert!(context.initial_tools.is_empty()); + assert_eq!(context.effective_tools, vec![chat_tool("lookup")]); + assert_eq!(context.tool_choice, ChatToolChoice::Auto); + assert!(!context.parallel_tool_calls); + assert!(context.parsing_enabled()); + } + + #[test] + fn resolved_tool_context_preserves_initial_then_message_order() { + let messages = vec![ + ChatMessage::developer("", Some(vec![chat_tool("middle")])), + ChatMessage::developer("policy", Some(vec![chat_tool("last")])), + ]; + + let context = ResolvedToolContext::new( + &messages, + vec![chat_tool("first")], + Some(ChatToolChoice::Required), + true, + ) + .unwrap(); + + assert_eq!( + context + .effective_tools + .iter() + .map(|tool| tool.name.as_str()) + .collect::>(), + vec!["first", "middle", "last"] + ); + assert_eq!(context.tool_choice, ChatToolChoice::Required); + } + + #[test] + fn resolved_tool_context_rejects_duplicate_names() { + let messages = vec![ChatMessage::developer("", Some(vec![chat_tool("lookup")]))]; + + let error = + ResolvedToolContext::new(&messages, vec![chat_tool("lookup")], None, true).unwrap_err(); + + assert!(matches!(error, Error::DuplicateToolName { name } if name == "lookup")); + } + + #[test] + fn resolved_tool_context_validates_named_choice_against_effective_tools() { + let messages = vec![ChatMessage::developer("", Some(vec![chat_tool("lookup")]))]; + let context = ResolvedToolContext::new( + &messages, + Vec::new(), + Some(ChatToolChoice::Function { + name: "lookup".to_string(), + }), + true, + ) + .unwrap(); + assert_eq!( + context.tool_choice, + ChatToolChoice::Function { + name: "lookup".to_string() + } + ); + + let error = ResolvedToolContext::new( + &messages, + Vec::new(), + Some(ChatToolChoice::Function { + name: "missing".to_string(), + }), + true, + ) + .unwrap_err(); + assert!(matches!( + error, + Error::ToolChoiceFunctionNotFound { name } if name == "missing" + )); + } + #[test] fn enable_thinking_is_none_when_no_kwargs_are_present() { let request = ChatRequest::for_test(); diff --git a/rust/src/chat/tests/chat.rs b/rust/src/chat/tests/chat.rs index e14f91a7b1a4..2bbc39512bb0 100644 --- a/rust/src/chat/tests/chat.rs +++ b/rust/src/chat/tests/chat.rs @@ -302,7 +302,7 @@ fn sample_request(request_id: &str) -> ChatRequest { fn sample_tool_request(request_id: &str) -> ChatRequest { let mut request = sample_request(request_id); - request.tools = vec![ChatTool { + let tools = vec![ChatTool { name: "get_weather".to_string(), description: Some("Get weather".to_string()), parameters: serde_json::json!({ @@ -312,7 +312,13 @@ fn sample_tool_request(request_id: &str) -> ChatRequest { }), strict: None, }]; - request.tool_choice = ChatToolChoice::Auto; + request.tool_context = vllm_chat::ResolvedToolContext::new( + &request.messages, + tools, + Some(ChatToolChoice::Auto), + true, + ) + .expect("tool context should resolve"); request } diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index 9cd670d29d08..f8c3519bbfdd 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -16,9 +16,9 @@ use serde_json_fmt::JsonFormat as JsonFmt; use serial_test::file_serial; use vllm_chat::{ AssistantContentBlock, AssistantMessage, AssistantMessageExt as _, AssistantToolCall, - ChatEvent, ChatMessage, ChatRequest, ChatRole, ChatTool, ChatToolChoice, FinishReason, - GenerationPromptMode, LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection, - RendererSelection, load_model_backends, + ChatEvent, ChatMessage, ChatRequest, ChatRole, ChatTool, FinishReason, GenerationPromptMode, + LoadModelBackendsOptions, NewChatOutputProcessorOptions, ParserSelection, RendererSelection, + load_model_backends, }; use vllm_text::{DecodedTextEvent, Finished, Prompt}; use vllm_tokenizer::Tokenizer; @@ -825,15 +825,12 @@ fn roundtrip_request( tools: Vec, thinking: Option, ) -> ChatRequest { + let tool_context = vllm_chat::ResolvedToolContext::new(&messages, tools, None, true) + .expect("tool context should resolve"); let mut request = ChatRequest { request_id: request_id.into(), messages, - tool_choice: if tools.is_empty() { - ChatToolChoice::None - } else { - ChatToolChoice::Auto - }, - tools, + tool_context, ..ChatRequest::for_test() }; diff --git a/rust/src/server/src/routes/openai/chat_completions/convert.rs b/rust/src/server/src/routes/openai/chat_completions/convert.rs index 521643fb6955..2c563f27bcbe 100644 --- a/rust/src/server/src/routes/openai/chat_completions/convert.rs +++ b/rust/src/server/src/routes/openai/chat_completions/convert.rs @@ -5,16 +5,16 @@ use itertools::Itertools as _; use vllm_chat::{ AssistantContentBlock, AssistantToolCall, ChatContent, ChatContentPart, ChatMessage as VllmChatMessage, ChatOptions, ChatRequest, ChatTool, ChatToolChoice, - GenerationPromptMode, SamplingParams, + GenerationPromptMode, ResolvedToolContext, SamplingParams, }; use super::types::ChatCompletionRequest; use super::validate; -use crate::error::{ApiError, bail_invalid_request}; +use crate::error::{ApiError, bail_invalid_request, chat_submit_error}; use crate::lora::LoraModelResolution; use crate::routes::openai::utils::structured_outputs::convert_from_response_format; use crate::routes::openai::utils::types::{ - ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue, + ChatMessage, ContentPart, MessageContent, Tool, ToolChoice, ToolChoiceValue, ToolReference, }; use crate::utils::{ ResolvedRequestContext, convert_logit_bias, merge_ec_transfer_params, merge_kv_transfer_params, @@ -130,6 +130,14 @@ pub(super) fn prepare_chat_request( request.vllm_xargs.as_ref(), ); + let tool_context = ResolvedToolContext::new( + &messages, + convert_tools(request.tools)?, + request.tool_choice.as_ref().map(convert_tool_choice).transpose()?, + request.parallel_tool_calls.unwrap_or(true), + ) + .map_err(|error| chat_submit_error("failed to resolve request tools", error))?; + let chat_request = ChatRequest { request_id: request_id.clone(), messages, @@ -168,9 +176,7 @@ pub(super) fn prepare_chat_request( response_format, template_kwargs, }, - tools: convert_tools(request.tools)?, - tool_choice: convert_tool_choice(request.tool_choice.as_ref())?, - parallel_tool_calls: request.parallel_tool_calls.unwrap_or(true), + tool_context, decode_options: vllm_text::output::TextDecodeOptions { skip_special_tokens: request.skip_special_tokens, include_stop_str_in_output: request.include_stop_str_in_output, @@ -399,17 +405,21 @@ fn convert_message_tools(tools: Option>) -> Result) -> Result { +fn convert_tool_choice(tool_choice: &ToolChoice) -> Result { match tool_choice { - None | Some(ToolChoice::Value(ToolChoiceValue::Auto)) => Ok(ChatToolChoice::Auto), - Some(ToolChoice::Value(ToolChoiceValue::None)) => Ok(ChatToolChoice::None), - Some(ToolChoice::Value(ToolChoiceValue::Required)) => Ok(ChatToolChoice::Required), - Some(ToolChoice::Function { + ToolChoice::Value(ToolChoiceValue::Auto) => Ok(ChatToolChoice::Auto), + ToolChoice::Value(ToolChoiceValue::None) => Ok(ChatToolChoice::None), + ToolChoice::Value(ToolChoiceValue::Required) => Ok(ChatToolChoice::Required), + ToolChoice::Function { tool_type, function, - }) if tool_type == "function" => Ok(ChatToolChoice::Function { + } if tool_type == "function" => Ok(ChatToolChoice::Function { name: function.name.clone(), }), + ToolChoice::AllowedTools { tools, .. } => bail_invalid_request!( + "allowed_tools tool_choice is not supported yet: {}.", + tools.iter().map(ToolReference::identifier).join(", ") + ), _ => bail_invalid_request!("tool_choice={:?} is not supported yet.", tool_choice), } } @@ -478,7 +488,7 @@ mod tests { ) .expect("request is valid"); - assert!(!prepared.chat_request.parallel_tool_calls); + assert!(!prepared.chat_request.parallel_tool_calls()); } #[test] @@ -490,7 +500,7 @@ mod tests { ) .expect("request is valid"); - assert!(prepared.chat_request.parallel_tool_calls); + assert!(prepared.chat_request.parallel_tool_calls()); } #[test] @@ -596,8 +606,8 @@ mod tests { min_tokens: 0, } ); - assert!(prepared.chat_request.tools.is_empty()); - assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto); + assert!(prepared.chat_request.initial_tools().is_empty()); + assert_eq!(prepared.chat_request.tool_choice(), &ChatToolChoice::None); } #[test] @@ -671,8 +681,8 @@ mod tests { min_tokens: 0, } ); - assert!(prepared.chat_request.tools.is_empty()); - assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto); + assert!(prepared.chat_request.initial_tools().is_empty()); + assert_eq!(prepared.chat_request.tool_choice(), &ChatToolChoice::None); } #[test] @@ -1009,8 +1019,8 @@ mod tests { }, ])] ); - assert!(prepared.chat_request.tools.is_empty()); - assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Auto); + assert!(prepared.chat_request.initial_tools().is_empty()); + assert_eq!(prepared.chat_request.tool_choice(), &ChatToolChoice::None); } #[test] @@ -1105,7 +1115,7 @@ mod tests { ] ); assert_eq!( - prepared.chat_request.tools, + prepared.chat_request.initial_tools(), vec![VllmChatTool { name: "get_weather".to_string(), description: Some("Get weather".to_string()), @@ -1116,7 +1126,7 @@ mod tests { strict: None, }] ); - assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::None); + assert_eq!(prepared.chat_request.tool_choice(), &ChatToolChoice::None); } #[test] @@ -1145,7 +1155,10 @@ mod tests { ) .expect("request is valid"); - assert_eq!(prepared.chat_request.tool_choice, ChatToolChoice::Required); + assert_eq!( + prepared.chat_request.tool_choice(), + &ChatToolChoice::Required + ); assert!(!prepared.options.is_named_tool_choice); } @@ -1181,8 +1194,8 @@ mod tests { .expect("request is valid"); assert_eq!( - prepared.chat_request.tool_choice, - ChatToolChoice::Function { + prepared.chat_request.tool_choice(), + &ChatToolChoice::Function { name: "get_weather".to_string(), } ); diff --git a/rust/src/server/src/routes/openai/chat_completions/types.rs b/rust/src/server/src/routes/openai/chat_completions/types.rs index c0735553c5ad..27911bc3f30e 100644 --- a/rust/src/server/src/routes/openai/chat_completions/types.rs +++ b/rust/src/server/src/routes/openai/chat_completions/types.rs @@ -14,8 +14,8 @@ use vllm_engine_core_client::protocol::sampling::RepetitionDetectionParams; use crate::routes::openai::utils::structured_outputs::ResponseFormat; use crate::routes::openai::utils::types::{ ChatLogProbs, ChatMessage, Normalizable, StreamOptions, StringOrArray, Tool, ToolCall, - ToolCallDelta, ToolChoice, ToolChoiceValue, ToolReference, UNKNOWN_MODEL_ID, Usage, - default_true, validate_messages, validate_stop, validate_top_p_value, + ToolCallDelta, ToolChoice, UNKNOWN_MODEL_ID, Usage, default_true, validate_messages, + validate_stop, validate_top_p_value, }; /// vLLM-compatible request type for the Chat Completions API. @@ -325,19 +325,6 @@ impl Normalizable for ChatCompletionRequest { self.max_completion_tokens = self.max_tokens; self.max_tokens = None; } - - // Apply tool_choice defaults - // If tools is None, leave tool_choice as None (don't set it) - if self.tool_choice.is_none() - && let Some(tools) = &self.tools - { - let choice_value = if tools.is_empty() { - ToolChoiceValue::None - } else { - ToolChoiceValue::Auto - }; - self.tool_choice = Some(ToolChoice::Value(choice_value)); - } } } @@ -492,99 +479,5 @@ fn validate_chat_cross_parameters( return Err(e); } - // 5. Validate tool_choice requires tools (except for "none") - if let Some(ref tool_choice) = req.tool_choice { - let has_tools = req.tools.as_ref().is_some_and(|t| !t.is_empty()); - - // Check if tool_choice is anything other than "none" - let is_some_choice = !matches!(tool_choice, ToolChoice::Value(ToolChoiceValue::None)); - - if is_some_choice && !has_tools { - let mut e = validator::ValidationError::new("tool_choice_requires_tools"); - e.message = Some("Invalid value for 'tool_choice': 'tool_choice' is only allowed when 'tools' are specified.".into()); - return Err(e); - } - - // Additional validation when tools are present - if let Some(tools) = req.tools.as_ref().filter(|t| !t.is_empty()) { - match tool_choice { - ToolChoice::Function { function, .. } => { - // Validate that the specified function name exists in tools - let function_exists = tools.iter().any(|tool| { - tool.tool_type == "function" && tool.function.name == function.name - }); - - if !function_exists { - let mut e = - validator::ValidationError::new("tool_choice_function_not_found"); - e.message = Some( - format!( - "Invalid value for 'tool_choice': function '{}' not found in 'tools'.", - function.name - ) - .into(), - ); - return Err(e); - } - } - ToolChoice::AllowedTools { - mode, - tools: allowed_tools, - .. - } => { - // Validate mode is "auto" or "required" - if mode != "auto" && mode != "required" { - let mut e = validator::ValidationError::new("tool_choice_invalid_mode"); - e.message = Some(format!( - "Invalid value for 'tool_choice.mode': must be 'auto' or 'required', got '{mode}'." - ).into()); - return Err(e); - } - - // Validate that all ToolReferences are Function type (Chat API only supports - // function tools) - for tool_ref in allowed_tools { - match tool_ref { - ToolReference::Function { name } => { - // Validate that the function exists in tools array - let tool_exists = tools.iter().any(|tool| { - tool.tool_type == "function" && tool.function.name == *name - }); - - if !tool_exists { - let mut e = validator::ValidationError::new( - "tool_choice_tool_not_found", - ); - e.message = Some( - format!( - "Invalid value for 'tool_choice.tools': tool '{name}' not found in 'tools'." - ) - .into(), - ); - return Err(e); - } - } - _ => { - // Chat Completion API only supports function tools in tool_choice - let mut e = validator::ValidationError::new( - "tool_choice_invalid_tool_type", - ); - e.message = Some( - format!( - "Invalid value for 'tool_choice.tools': Chat Completion API only supports function tools, got '{}'.", - tool_ref.identifier() - ) - .into(), - ); - return Err(e); - } - } - } - } - ToolChoice::Value(_) => {} - } - } - } - Ok(()) } diff --git a/rust/src/server/src/routes/tokenize/types.rs b/rust/src/server/src/routes/tokenize/types.rs index 65b3e7a2b58e..5a2d49b095fb 100644 --- a/rust/src/server/src/routes/tokenize/types.rs +++ b/rust/src/server/src/routes/tokenize/types.rs @@ -7,7 +7,7 @@ use itertools::Itertools as _; use serde::{Deserialize, Serialize}; use serde_json::Value; use validator::{Validate, ValidationErrors}; -use vllm_chat::{ChatOptions, ChatRequest, ChatToolChoice, SamplingParams}; +use vllm_chat::{ChatOptions, ChatRequest, ResolvedToolContext, SamplingParams}; use vllm_text::output::TextDecodeOptions; use vllm_text::{Prompt, TextRequest}; @@ -97,6 +97,11 @@ impl TokenizeChatRequest { &messages, )?; + let tool_context = + ResolvedToolContext::new(&messages, convert_tools(self.tools)?, None, true).map_err( + |error| crate::error::chat_submit_error("failed to resolve request tools", error), + )?; + Ok(ChatRequest { request_id, messages, @@ -108,9 +113,7 @@ impl TokenizeChatRequest { response_format: None, template_kwargs: self.chat_template_kwargs.unwrap_or_default(), }, - tools: convert_tools(self.tools)?, - tool_choice: ChatToolChoice::Auto, - parallel_tool_calls: true, + tool_context, decode_options: TextDecodeOptions::default(), intermediate: false, priority: 0, @@ -199,7 +202,7 @@ mod tests { req.into_chat_request("tokenize-test".to_string()).expect("request is valid"); assert_eq!( - chat_request.tools, + chat_request.initial_tools(), vec![ChatTool { name: "get_weather".to_string(), description: Some("Get weather".to_string()), From 48bada6ea49ad7f3ecbe03128aa76562089c8b00 Mon Sep 17 00:00:00 2001 From: Tarun Kumar Date: Tue, 11 Aug 2026 07:01:19 +0530 Subject: [PATCH 20/30] Fix chat completion 500 on non-object JSON bodies (#51654) Signed-off-by: Tarun Kumar --- .../test_non_object_body_validation.py | 38 +++++++++++++++++++ .../openai/chat_completion/protocol.py | 27 +++++++++---- 2 files changed, 58 insertions(+), 7 deletions(-) create mode 100644 tests/entrypoints/openai/chat_completion/test_non_object_body_validation.py diff --git a/tests/entrypoints/openai/chat_completion/test_non_object_body_validation.py b/tests/entrypoints/openai/chat_completion/test_non_object_body_validation.py new file mode 100644 index 000000000000..7dba0018610b --- /dev/null +++ b/tests/entrypoints/openai/chat_completion/test_non_object_body_validation.py @@ -0,0 +1,38 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Non-object JSON bodies must fail validation cleanly (4xx), not AttributeError (500). + +mode=before validators that call data.get(...) without an isinstance(data, dict) +guard raise AttributeError for string/list/scalar bodies and surface as HTTP 500. +""" + +import pytest +from pydantic import ValidationError + +from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionRequest + + +@pytest.mark.parametrize( + "payload", + [ + "this is not valid json{{{", + ["not", "an", "object"], + 42, + None, + ], +) +def test_chat_completion_request_rejects_non_object_body(payload): + with pytest.raises(ValidationError): + ChatCompletionRequest.model_validate(payload) + + +def test_chat_completion_request_cache_salt_still_validated_on_dict(): + with pytest.raises(ValidationError, match="cache_salt"): + ChatCompletionRequest.model_validate( + { + "model": "qwen", + "messages": [{"role": "user", "content": "hello"}], + "cache_salt": "", + } + ) diff --git a/vllm/entrypoints/openai/chat_completion/protocol.py b/vllm/entrypoints/openai/chat_completion/protocol.py index d25b3886ad54..eda0f5a3bae0 100644 --- a/vllm/entrypoints/openai/chat_completion/protocol.py +++ b/vllm/entrypoints/openai/chat_completion/protocol.py @@ -745,6 +745,8 @@ def to_sampling_params( @model_validator(mode="before") @classmethod def validate_response_format(cls, data): + if not isinstance(data, dict): + return data response_format = data.get("response_format") if response_format is None: return data @@ -776,6 +778,8 @@ def validate_response_format(cls, data): @model_validator(mode="before") @classmethod def validate_stream_options(cls, data): + if not isinstance(data, dict): + return data if data.get("stream_options") and not data.get("stream"): raise VLLMValidationError( "Stream options can only be defined when `stream=True`.", @@ -787,6 +791,8 @@ def validate_stream_options(cls, data): @model_validator(mode="before") @classmethod def check_logprobs(cls, data): + if not isinstance(data, dict): + return data if data.get("logprob_token_ids") and data.get("use_beam_search"): raise VLLMValidationError( "`logprob_token_ids` is not supported with beam search.", @@ -845,6 +851,8 @@ def check_logprobs(cls, data): def check_structured_outputs_count(cls, data): if isinstance(data, ValueError): raise data + if not isinstance(data, dict): + return data if data.get("structured_outputs", None) is None: return data @@ -970,6 +978,8 @@ def check_tool_usage(cls, data): @model_validator(mode="before") @classmethod def check_generation_prompt(cls, data): + if not isinstance(data, dict): + return data if data.get("continue_final_message") and data.get("add_generation_prompt"): raise VLLMValidationError( "Cannot set both `continue_final_message` and " @@ -1105,6 +1115,8 @@ class BatchChatCompletionRequest(OpenAIBaseModel): def check_batch_mode(cls, data: Any) -> Any: if isinstance(data, BatchChatCompletionRequest): data = data.model_dump(exclude_unset=True) + if not isinstance(data, dict): + return data if data.get("use_beam_search"): raise VLLMValidationError( "Batch chat completions do not support beam search. " @@ -1117,13 +1129,14 @@ def check_batch_mode(cls, data: Any) -> Any: parameter="logprob_token_ids", ) response_format = data.get("response_format") - rf_type = ( - response_format.get("type") - if isinstance(response_format, dict) - else getattr(response_format, "type", None) - ) - if rf_type == "structural_tag": - validate_structural_tag_response_format(response_format) + if response_format is not None: + rf_type = ( + response_format.get("type") + if isinstance(response_format, dict) + else getattr(response_format, "type", None) + ) + if rf_type == "structural_tag": + validate_structural_tag_response_format(response_format) if (structured_outputs := data.get("structured_outputs")) is not None: validate_structured_outputs_structural_tag(structured_outputs) n = data.get("n", 1) From 1a1727330a14cfa9f4bdcebfaba403fd02ae7011 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 10 Aug 2026 20:56:39 -0500 Subject: [PATCH 21/30] [CI] Solidify speculative decoding E2E coverage (#50713) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas Co-authored-by: OpenAI Codex --- .buildkite/test_areas/spec_decode.yaml | 104 +++++++++++++- .../acceptance_rates/dflash/test_dflash.py | 59 ++++---- .../acceptance_rates/dspark/test_dspark.py | 53 ++++--- .../acceptance_rates/mtp_other/test_medusa.py | 26 ++-- .../acceptance_rates/mtp_other/test_mtp.py | 2 + .../mtp_other/test_synthetic.py | 60 ++++---- .../e2e/spec_decode/acceptance_rates/utils.py | 102 +++++++------ .../e2e/spec_decode/draft_model/test_async.py | 54 ++++--- .../draft_model/test_draft_model.py | 134 +++++++++++------- .../e2e/spec_decode/draft_model/test_lora.py | 79 +++++------ .../eagle/test_eagle_correctness.py | 19 ++- tests/v1/e2e/spec_decode/eagle/utils.py | 117 ++++++++------- tests/v1/e2e/spec_decode/mtp/test_mtp.py | 94 ++++++------ .../ngram_suffix/test_ngram_suffix.py | 106 +++++++++----- .../speculators/test_speculators.py | 29 ++-- .../e2e/spec_decode/test_mtp_parallel_load.py | 70 ++++++--- tests/v1/e2e/spec_decode/utils.py | 110 +++++++++++--- vllm/config/speculative.py | 2 +- 18 files changed, 749 insertions(+), 471 deletions(-) diff --git a/.buildkite/test_areas/spec_decode.yaml b/.buildkite/test_areas/spec_decode.yaml index 464c4c021dbd..852e5a463491 100644 --- a/.buildkite/test_areas/spec_decode.yaml +++ b/.buildkite/test_areas/spec_decode.yaml @@ -29,8 +29,8 @@ steps: depends_on: - image-build-amd -- label: Spec Decode Eagle - key: spec-decode-eagle +- label: "Spec Decode Eagle 1: DeepSeek + Qwen" + key: spec-decode-eagle-1-deepseek-qwen timeout_in_minutes: 25 device: h200_35gb source_file_dependencies: @@ -38,12 +38,38 @@ steps: - vllm/v1/worker/gpu/spec_decode/ - tests/v1/e2e/spec_decode/ commands: - - pytest -v -s v1/e2e/spec_decode/eagle/ + - pytest -v -s v1/e2e/spec_decode/eagle/ -k "deepseek_eagle or qwen3_eagle3" mirror: amd: dind: false device: mi300_1 - timeout_in_minutes: 55 + timeout_in_minutes: 25 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/model_executor/model_loader/ + - vllm/v1/sample/ + - vllm/model_executor/layers/ + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py + +- label: "Spec Decode Eagle 2: Llama 3 + Qwen VL + Other" + key: spec-decode-eagle-2-llama3-qwen-vl-other + timeout_in_minutes: 25 + device: h200_35gb + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - tests/v1/e2e/spec_decode/ + commands: + - pytest -v -s v1/e2e/spec_decode/eagle/ -k "not deepseek_eagle and not qwen3_eagle3" + mirror: + amd: + dind: false + device: mi300_1 + timeout_in_minutes: 25 depends_on: - image-build-amd source_file_dependencies: @@ -192,8 +218,8 @@ steps: - export VLLM_ALLOW_INSECURE_SERIALIZATION=1 - pytest -v -s v1/spec_decode/test_speculators_correctness.py -m slow_test -- label: Spec Decode DeepSeek MTP Parallel Load (B200) - key: spec-decode-deepseek-mtp-parallel-load-b200 +- label: Spec Decode DeepSeek MTP Parallel Load (2xB200-2xMI300) + key: spec-decode-deepseek-mtp-parallel-load-2xb200-2xmi300 timeout_in_minutes: 30 device: b200-k8s optional: true @@ -207,6 +233,26 @@ steps: - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py commands: - pytest -v -s v1/e2e/spec_decode/test_mtp_parallel_load.py + mirror: + amd: + dind: false + device: mi300_2 + timeout_in_minutes: 45 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/distributed/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/model_executor/models/deepseek_mtp.py + - vllm/model_executor/models/deepseek_v2.py + - vllm/v1/attention/ + - vllm/v1/engine/ + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/worker/gpu_model_runner.py + - tests/v1/e2e/spec_decode/test_mtp_parallel_load.py + - vllm/platforms/rocm.py - label: Spec Decode AL DFlash Nightly key: spec-decode-dflash-nightly @@ -221,6 +267,27 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dflash/ + mirror: + amd: + dind: false + device: mi300_1 + timeout_in_minutes: 120 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/attention/ + - vllm/v1/sample/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/laguna.py + - vllm/model_executor/models/laguna_dflash.py + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py - label: Spec Decode AL DSpark Nightly key: spec-decode-dspark-nightly @@ -235,6 +302,31 @@ steps: - tests/v1/e2e/spec_decode/ commands: - pytest -v -s v1/e2e/spec_decode/acceptance_rates/dspark/ + mirror: + amd: + dind: false + device: mi300_1 + timeout_in_minutes: 90 + depends_on: + - image-build-amd + source_file_dependencies: + - vllm/v1/spec_decode/ + - vllm/v1/worker/gpu/spec_decode/ + - vllm/v1/attention/ + - vllm/v1/sample/ + - vllm/model_executor/model_loader/ + - vllm/model_executor/layers/ + - vllm/model_executor/models/qwen2.py + - vllm/model_executor/models/qwen3.py + - vllm/model_executor/models/qwen3_dflash.py + - vllm/model_executor/models/qwen3_dspark.py + - vllm/model_executor/models/gemma4.py + - vllm/model_executor/models/gemma4_mm.py + - vllm/model_executor/models/gemma4_unified.py + - vllm/model_executor/models/gemma4_mtp.py + - vllm/model_executor/models/gemma4_dspark.py + - tests/v1/e2e/spec_decode/ + - vllm/platforms/rocm.py - label: Spec Decode AL MTP + Other Acceptance Nightly key: spec-decode-mtp-other-acceptance-nightly diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py index a4c9c67a0523..2b1800f78234 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dflash/test_dflash.py @@ -4,12 +4,10 @@ from dataclasses import dataclass import pytest -import torch from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline from tests.utils import single_gpu_only -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory +from vllm.config import CompilationConfig from ...utils import compute_acceptance_len from ..utils import run_acceptance_length_eval @@ -56,9 +54,11 @@ class DFlashCorrectnessConfig: def test_dflash_reference_acceptance_lengths( monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, + vllm_runner, ): run_acceptance_length_eval( monkeypatch, + vllm_runner, spec_config={ "model": QWEN3_DFLASH.model, "trust_remote_code": True, @@ -110,14 +110,16 @@ def test_dflash_correctness( monkeypatch: pytest.MonkeyPatch, config: DFlashCorrectnessConfig, use_mrv2: bool, + vllm_runner, ): """Guard GSM8K accuracy and batched acceptance length for DFlash models.""" monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") if config.disable_flashinfer_sampler: monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") - spec_llm = LLM( - model=config.model, + with vllm_runner( + config.model, + block_size=None, trust_remote_code=True, speculative_config={ "method": "dflash", @@ -130,23 +132,30 @@ def test_dflash_correctness( gpu_memory_utilization=0.85, enforce_eager=config.enforce_eager, disable_log_stats=False, - ) - - results = evaluate_gsm8k_offline( - spec_llm, - num_questions=config.num_questions, - use_chat_completions=config.use_chat_completions, - ) - accuracy = results["accuracy"] - acceptance_len = compute_acceptance_len(spec_llm.get_metrics()) - print( - f"{config.model}: GSM8K accuracy={accuracy:.3f}, " - f"acceptance_len={acceptance_len:.2f}" - ) - - assert accuracy >= config.expected_accuracy - assert acceptance_len >= config.expected_acceptance_len - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + ) as spec_runner: + spec_llm = spec_runner.llm + results = evaluate_gsm8k_offline( + spec_llm, + num_questions=config.num_questions, + use_chat_completions=config.use_chat_completions, + ) + accuracy = results["accuracy"] + acceptance_len = compute_acceptance_len(spec_llm.get_metrics()) + context = ( + f"DFlash target={config.model}, draft={config.draft_model}, MRV2={use_mrv2}" + ) + print( + f"{context}: GSM8K accuracy={accuracy:.3f}, " + f"acceptance_len={acceptance_len:.2f}" + ) + + assert accuracy >= config.expected_accuracy, ( + f"{context}: GSM8K accuracy {accuracy:.3f} is below " + f"{config.expected_accuracy:.3f}; acceptance_len={acceptance_len:.3f}" + ) + assert acceptance_len >= config.expected_acceptance_len, ( + f"{context}: acceptance_len {acceptance_len:.3f} is below " + f"{config.expected_acceptance_len:.3f}; accuracy={accuracy:.3f}" + ) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py index 781374749dbf..e90100472400 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py @@ -4,11 +4,10 @@ from dataclasses import dataclass import pytest -import torch from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline -from vllm import LLM -from vllm.distributed import cleanup_dist_env_and_memory +from vllm.config import CompilationConfig +from vllm.platforms import current_platform from ...utils import compute_acceptance_len, compute_acceptance_rate @@ -44,6 +43,7 @@ class DSparkCorrectnessConfig: reference_accuracy=0.801, reference_acceptance_rate=0.428, reference_acceptance_len=3.994, + attention_backend="FLASH_ATTN", ) # References from five 200-question runs at temperature 1.0: accuracy @@ -58,6 +58,9 @@ class DSparkCorrectnessConfig: num_questions=200, use_chat_completions=True, gpu_memory_utilization=0.85, + enforce_eager=True, + disable_flashinfer_sampler=True, + language_model_only=True, ) # Reference from 200 GSM8K questions at temperature 1.0. @@ -88,34 +91,40 @@ class DSparkCorrectnessConfig: def test_dspark_correctness_and_acceptance_rate( monkeypatch: pytest.MonkeyPatch, config: DSparkCorrectnessConfig, + vllm_runner, ): """Guard GSM8K accuracy and acceptance metrics for DSpark models.""" + if config.disable_flashinfer_sampler: + monkeypatch.setenv("VLLM_USE_FLASHINFER_SAMPLER", "0") + speculative_config = { "method": "dspark", "model": config.draft_model, "num_speculative_tokens": config.num_speculative_tokens, "draft_sample_method": "probabilistic", } - if config.attention_backend is not None: + if config.attention_backend is not None and current_platform.is_cuda(): speculative_config["attention_backend"] = config.attention_backend - llm_kwargs = { - "model": config.model, + runner_config = { + "block_size": None, "trust_remote_code": True, "speculative_config": speculative_config, "max_model_len": config.max_model_len, "gpu_memory_utilization": config.gpu_memory_utilization, "enforce_eager": config.enforce_eager, + "enable_chunked_prefill": None, "enable_prefix_caching": False, "disable_log_stats": False, + "compilation_config": CompilationConfig(), } if config.max_num_seqs is not None: - llm_kwargs["max_num_seqs"] = config.max_num_seqs + runner_config["max_num_seqs"] = config.max_num_seqs if config.language_model_only: - llm_kwargs["language_model_only"] = True + runner_config["language_model_only"] = True - spec_llm = LLM(**llm_kwargs) - try: + with vllm_runner(config.model, **runner_config) as spec_runner: + spec_llm = spec_runner.llm results = evaluate_gsm8k_offline( spec_llm, num_questions=config.num_questions, @@ -128,18 +137,20 @@ def test_dspark_correctness_and_acceptance_rate( metrics = spec_llm.get_metrics() acceptance_rate = compute_acceptance_rate(metrics) acceptance_len = compute_acceptance_len(metrics) - print( - f"{config.model}: GSM8K accuracy={accuracy:.3f}, " - f"acceptance_rate={acceptance_rate:.2f}, " - f"acceptance_len={acceptance_len:.2f}" + context = f"DSpark target={config.model}, draft={config.draft_model}" + metrics_summary = ( + f"gsm8k_accuracy={accuracy:.3f}, " + f"acceptance_rate={acceptance_rate:.3f}, " + f"acceptance_len={acceptance_len:.3f}" ) + print(f"{context}: {metrics_summary}") - assert accuracy >= config.reference_accuracy * REGRESSION_TOLERANCE + assert accuracy >= config.reference_accuracy * REGRESSION_TOLERANCE, ( + f"{context}: {metrics_summary}" + ) assert ( acceptance_rate >= config.reference_acceptance_rate * REGRESSION_TOLERANCE - ) - assert acceptance_len >= config.reference_acceptance_len * REGRESSION_TOLERANCE - finally: - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + ), f"{context}: {metrics_summary}" + assert ( + acceptance_len >= config.reference_acceptance_len * REGRESSION_TOLERANCE + ), f"{context}: {metrics_summary}" diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py index cd9c115a9dc3..1e10d34527ac 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_medusa.py @@ -1,17 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch - from tests.evals.gsm8k.gsm8k_eval import _build_gsm8k_prompts -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig from ...utils import compute_acceptance_rate def test_medusa_acceptance_rate( sampling_config: SamplingParams, + vllm_runner, ): """Verify a trained Medusa checkpoint achieves nonzero acceptance rate. @@ -23,8 +22,10 @@ def test_medusa_acceptance_rate( medusa_model = "FasterDecoding/medusa-vicuna-7b-v1.3" prompts = _build_gsm8k_prompts(num_questions=10, num_shots=1)[0] - spec_llm = LLM( - model=target_model, + with vllm_runner( + target_model, + block_size=None, + trust_remote_code=False, speculative_config={ "method": "medusa", "model": medusa_model, @@ -32,14 +33,13 @@ def test_medusa_acceptance_rate( }, max_model_len=1024, enforce_eager=True, + enable_chunked_prefill=None, disable_log_stats=False, - ) - spec_llm.generate(prompts, sampling_config) - metrics = spec_llm.get_metrics() - acceptance_rate = compute_acceptance_rate(metrics) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + compilation_config=CompilationConfig(), + ) as spec_runner: + spec_runner.llm.generate(prompts, sampling_config) + metrics = spec_runner.llm.get_metrics() + acceptance_rate = compute_acceptance_rate(metrics) min_acceptance_rate = 0.198 print(f"Medusa acceptance rate: {acceptance_rate:.4f} (min {min_acceptance_rate})") diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py index 831501361ac7..ea5f139fdf41 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_mtp.py @@ -10,9 +10,11 @@ def test_gemma4_mtp_acceptance_lengths( monkeypatch: pytest.MonkeyPatch, use_mrv2: bool, + vllm_runner, ): run_acceptance_length_eval( monkeypatch, + vllm_runner, spec_config={ "model": "google/gemma-4-E4B-it", "trust_remote_code": True, diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py index aef0805513ad..8714e4f56a1e 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/mtp_other/test_synthetic.py @@ -1,25 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -import torch - from tests.utils import single_gpu_only -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig from ...utils import compute_acceptance_len, get_test_prompts @single_gpu_only -def test_synthetic_acceptance_rate(): +def test_synthetic_acceptance_rate(vllm_runner): """Verify that synthetic rejection sampling produces an acceptance length close to the requested mean acceptance length.""" num_spec_tokens = 3 expected_acceptance_len = 1.875 tolerance = 0.15 - spec_llm = LLM( - model="meta-llama/Llama-3.2-1B-Instruct", + with vllm_runner( + "meta-llama/Llama-3.2-1B-Instruct", + block_size=None, trust_remote_code=True, speculative_config={ "method": "eagle3", @@ -31,28 +30,27 @@ def test_synthetic_acceptance_rate(): }, max_model_len=2048, enforce_eager=True, + enable_chunked_prefill=None, disable_log_stats=False, - ) - - test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50) - spec_llm.chat( - test_prompts, - SamplingParams(temperature=0, max_tokens=64, ignore_eos=True), - ) - - metrics = spec_llm.get_metrics() - acceptance_len = compute_acceptance_len(metrics) - - print( - f"Synthetic acceptance length: {acceptance_len:.3f}" - f" (expected={expected_acceptance_len:.3f}," - f" tolerance=±{tolerance})" - ) - assert abs(acceptance_len - expected_acceptance_len) <= tolerance, ( - f"Synthetic acceptance length {acceptance_len:.3f} is not within" - f" ±{tolerance} of expected {expected_acceptance_len:.3f}" - ) - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + compilation_config=CompilationConfig(), + ) as spec_runner: + spec_llm = spec_runner.llm + + test_prompts = get_test_prompts(mm_enabled=False, num_prompts=50) + spec_llm.chat( + test_prompts, + SamplingParams(temperature=0, max_tokens=64, ignore_eos=True), + ) + + metrics = spec_llm.get_metrics() + acceptance_len = compute_acceptance_len(metrics) + + print( + f"Synthetic acceptance length: {acceptance_len:.3f}" + f" (expected={expected_acceptance_len:.3f}," + f" tolerance=±{tolerance})" + ) + assert abs(acceptance_len - expected_acceptance_len) <= tolerance, ( + f"Synthetic acceptance length {acceptance_len:.3f} is not within" + f" ±{tolerance} of expected {expected_acceptance_len:.3f}" + ) diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py index 631dc95b8186..97aa82fb5456 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/utils.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/utils.py @@ -4,11 +4,9 @@ from typing import Any import pytest -import torch -from tqdm import tqdm -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig from ..utils import compute_acceptance_len @@ -39,6 +37,7 @@ def load_and_process_dataset(data_name: str): def run_acceptance_length_eval( monkeypatch: pytest.MonkeyPatch, + vllm_runner, spec_config: dict[str, Any], expected_acceptance_lengths: dict[str, float], chat_template_kwargs: dict[str, Any], @@ -52,54 +51,53 @@ def run_acceptance_length_eval( stays within tolerance of the reference figure for each dataset. """ monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if use_mrv2 else "0") - spec_llm = LLM(**spec_config) - - max_prompts_per_dataset = 200 # mt-bench has 80, humaneval has 164, truncates gsm8k - - tokenizer = spec_llm.get_tokenizer() - for dataset_name, expected_len in expected_acceptance_lengths.items(): - dataset = load_and_process_dataset(dataset_name) - prev_metrics = None - acceptance_lengths = [] - for i in tqdm( - range(min(max_prompts_per_dataset, len(dataset))), - desc=f"Processing {dataset_name}", - ): - user_content = dataset[i]["turns"][0] - prompt_text = tokenizer.apply_chat_template( - [{"role": "user", "content": user_content}], - tokenize=False, - add_generation_prompt=True, - **chat_template_kwargs, + runner_config = spec_config.copy() + model = runner_config.pop("model") + runner_config.setdefault("block_size", None) + runner_config.setdefault("enable_chunked_prefill", None) + runner_config.setdefault("compilation_config", CompilationConfig()) + with vllm_runner(model, **runner_config) as spec_runner: + spec_llm = spec_runner.llm + # MT-Bench has 80 prompts and HumanEval has 164; truncate GSM8K. + max_prompts_per_dataset = 200 + + tokenizer = spec_llm.get_tokenizer() + for dataset_name, expected_len in expected_acceptance_lengths.items(): + dataset = load_and_process_dataset(dataset_name) + prev_metrics = None + acceptance_lengths = [] + for i in range(min(max_prompts_per_dataset, len(dataset))): + user_content = dataset[i]["turns"][0] + prompt_text = tokenizer.apply_chat_template( + [{"role": "user", "content": user_content}], + tokenize=False, + add_generation_prompt=True, + **chat_template_kwargs, + ) + + # Greedy (temp=0) so acceptance length is deterministic and comparable + # across runs. + spec_llm.generate( + [prompt_text], + SamplingParams(temperature=0, max_tokens=2048), + use_tqdm=False, + ) + current_metrics = spec_llm.get_metrics() + acceptance_len = compute_acceptance_len(current_metrics, prev_metrics) + prev_metrics = current_metrics + acceptance_lengths.append(acceptance_len) + + mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths) + # Fairly tight tolerance of 95% against the reference figures, + # watching for regressions. Can be relaxed if test is flaky but be sure to + # check for genuine issues such as #40727. + expected_len = expected_len * 0.95 + print( + f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" + f" (expected at least {expected_len:.2f})" ) - # Greedy (temp=0) so acceptance length is deterministic and comparable - # across runs. - spec_llm.generate( - [prompt_text], - SamplingParams(temperature=0, max_tokens=2048), - use_tqdm=False, + assert mean_acceptance_length >= expected_len, ( + f"acceptance_len for {dataset_name} is below expected threshold: " + f"{mean_acceptance_length:.2f} < {expected_len:.2f}" ) - current_metrics = spec_llm.get_metrics() - acceptance_len = compute_acceptance_len(current_metrics, prev_metrics) - prev_metrics = current_metrics - acceptance_lengths.append(acceptance_len) - - mean_acceptance_length = sum(acceptance_lengths) / len(acceptance_lengths) - # Fairly tight tolerance of 95% against the reference figures, - # watching for regressions. Can be relaxed if test is flaky but be sure to - # check for genuine issues such as #40727. - expected_len = expected_len * 0.95 - print( - f"acceptance_len for {dataset_name}: {mean_acceptance_length:.2f}" - f" (expected at least {expected_len:.2f})" - ) - - assert mean_acceptance_length >= expected_len, ( - f"acceptance_len for {dataset_name} is below expected threshold: " - f"{mean_acceptance_length:.2f} < {expected_len:.2f}" - ) - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() diff --git a/tests/v1/e2e/spec_decode/draft_model/test_async.py b/tests/v1/e2e/spec_decode/draft_model/test_async.py index b19f90e2cdc6..d7f932b0c256 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_async.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_async.py @@ -102,14 +102,15 @@ def test_no_sync_with_spec_decode( spec_model: str, method: str, num_spec_tokens: int, + vllm_runner, ): """ Test that no implicit GPU-CPU sync occurs during speculative decoding generation. """ # Import vLLM AFTER sync_tracker fixture has applied the patch - from vllm import LLM, SamplingParams - from vllm.distributed import cleanup_dist_env_and_memory + from vllm import SamplingParams + from vllm.config import CompilationConfig # Qwen3.5 is a VLM; without this, profile_run runs the ViT warmup # and peaks well above the 18GB MIG slice used by one of the CI lanes. @@ -119,8 +120,10 @@ def test_no_sync_with_spec_decode( if "Qwen3.5" in model: extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0} - llm = LLM( - model=model, + with vllm_runner( + model, + block_size=None, + trust_remote_code=False, max_model_len=256, speculative_config={ "method": method, @@ -129,24 +132,31 @@ def test_no_sync_with_spec_decode( }, enforce_eager=True, async_scheduling=True, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), **extra_kwargs, - ) - - # Assert async scheduling is actually active before running inference. - assert llm.llm_engine.vllm_config.scheduler_config.async_scheduling, ( - f"Expected async_scheduling=True for spec decode, got False. method={method}" - ) - - outputs = llm.generate( - ["Hello, my name is"], - SamplingParams(temperature=0, max_tokens=10), - ) - - assert len(outputs) == 1 - assert len(outputs[0].outputs[0].text) > 0 - - del llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + ) as runner: + llm = runner.llm + + # Assert async scheduling is actually active before running inference. + assert llm.llm_engine.vllm_config.scheduler_config.async_scheduling, ( + f"Expected async_scheduling=True for spec decode, got False. " + f"method={method}, target={model}, draft={spec_model}" + ) + + outputs = llm.generate( + ["Hello, my name is"], + SamplingParams(temperature=0, max_tokens=10), + ) + + assert len(outputs) == 1, ( + f"{method} target={model}: expected one request output, got {len(outputs)}" + ) + assert outputs[0].outputs, ( + f"{method} target={model}: request output has no completion candidates" + ) + assert outputs[0].outputs[0].text, ( + f"{method} target={model}: generated completion text is empty" + ) sync_tracker.assert_no_sync() diff --git a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py index 2e9821bc4119..cf7b44fb226a 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_draft_model.py @@ -4,17 +4,16 @@ from dataclasses import dataclass import pytest -import torch from tests.utils import multi_gpu_only, single_gpu_only -from vllm import LLM, SamplingParams -from vllm.config import VllmConfig, replace -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig, VllmConfig, replace +from vllm.config.kernel import MoEBackend from vllm.engine.arg_utils import EngineArgs +from vllm.platforms import current_platform from ..utils import ( Messages, - _skip_if_insufficient_gpus_for_tp, compute_acceptance_len, compute_acceptance_rate, evaluate_llm_for_gsm8k, @@ -42,7 +41,8 @@ class ArgsTest: max_model_len: int = 2048 gpu_memory_utilization: float = 0.5 dataset: str = "test_prompts" - num_prompts: int = 100 + # Doubling 100 to 200 reduces sampling standard error by about 29%. + num_prompts: int = 200 def get_messages(dataset: str, n: int) -> list[Messages]: @@ -89,13 +89,13 @@ def some_high_acceptance_metrics() -> dict: @pytest.mark.parametrize("args", cases) @pytest.mark.parametrize("enforce_eager", [True, False]) @single_gpu_only -def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool): +def test_draft_model_correctness(args: ArgsTest, enforce_eager: bool, vllm_runner): args.enforce_eager = enforce_eager - assert_draft_model_correctness(args) + assert_draft_model_correctness(args, vllm_runner) @single_gpu_only -def test_draft_model_realistic_example(): +def test_draft_model_realistic_example(vllm_runner): args = ArgsTest( target_model="Qwen/Qwen3-1.7B", draft_model="Qwen/Qwen3-0.6B", @@ -106,11 +106,11 @@ def test_draft_model_realistic_example(): expected_acceptance_len=2.6, # ref: 2.86 expected_acceptance_rate=0.5, # ref: 0.62 ) - assert_draft_model_correctness(args) + assert_draft_model_correctness(args, vllm_runner) @single_gpu_only -def test_draft_model_parallel_drafting(): +def test_draft_model_parallel_drafting(vllm_runner): args = ArgsTest( target_model="Qwen/Qwen3-1.7B", draft_model="amd/PARD-Qwen3-0.6B", @@ -122,7 +122,7 @@ def test_draft_model_parallel_drafting(): expected_acceptance_len=2.3, # ref: 2.52 expected_acceptance_rate=0.4, # ref: 0.51 ) - assert_draft_model_correctness(args) + assert_draft_model_correctness(args, vllm_runner) @pytest.mark.parametrize( @@ -136,7 +136,9 @@ def test_draft_model_parallel_drafting(): ) @pytest.mark.parametrize("enforce_eager", [True, False]) @single_gpu_only -def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): +def test_draft_model_quantization( + models: tuple[str, str], enforce_eager: bool, vllm_runner +): tgt_model, draft_model = models sd_case = ArgsTest( target_model=tgt_model, @@ -144,13 +146,12 @@ def test_draft_model_quantization(models: tuple[str, str], enforce_eager: bool): **some_high_acceptance_metrics(), enforce_eager=enforce_eager, ) - assert_draft_model_correctness(sd_case) + assert_draft_model_correctness(sd_case, vllm_runner) @multi_gpu_only(num_gpus=2) -def test_draft_model_tensor_parallelism(): +def test_draft_model_tensor_parallelism(vllm_runner): """Ensure spec decode works when running with TP > 1.""" - _skip_if_insufficient_gpus_for_tp(2) sd_case = ArgsTest( target_model="Qwen/Qwen3-1.7B", target_tensor_parallel_size=2, @@ -160,15 +161,13 @@ def test_draft_model_tensor_parallelism(): enforce_eager=False, expected_gsm8k_accuracy=0.5, ) - assert_draft_model_correctness(sd_case) + assert_draft_model_correctness(sd_case, vllm_runner) @multi_gpu_only(num_gpus=2) def test_draft_model_engine_args_tensor_parallelism(): """Ensure the vllm_config for the draft model is created correctly, and independently of the target model (quantization, TP, etc.)""" - _skip_if_insufficient_gpus_for_tp(2) - engine_args = EngineArgs( model="Qwen/Qwen3-1.7B-FP8", # <<< tgt quantized tensor_parallel_size=2, @@ -212,13 +211,24 @@ def _apply_draft_moe_backend(vllm_config: VllmConfig) -> VllmConfig: return vllm_config -def test_draft_model_moe_backend_override(): +def _platform_moe_backend(monkeypatch: pytest.MonkeyPatch) -> MoEBackend: + if current_platform.is_rocm(): + monkeypatch.setenv("VLLM_ROCM_USE_AITER", "1") + monkeypatch.setenv("VLLM_ROCM_USE_AITER_MOE", "1") + return "aiter" + if current_platform.is_cuda(): + return "flashinfer_trtllm" + return "auto" + + +def test_draft_model_moe_backend_override(monkeypatch: pytest.MonkeyPatch): """When moe_backend is set in speculative_config, the draft VllmConfig should use it while the target keeps its own setting.""" + target_moe_backend = _platform_moe_backend(monkeypatch) engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", tensor_parallel_size=1, - moe_backend="flashinfer_trtllm", + moe_backend=target_moe_backend, speculative_config={ "model": "Qwen/Qwen3-0.6B", "method": "draft_model", @@ -227,22 +237,23 @@ def test_draft_model_moe_backend_override(): }, ) tgt_config: VllmConfig = engine_args.create_engine_config() - assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm" + assert tgt_config.kernel_config.moe_backend == target_moe_backend assert tgt_config.speculative_config.moe_backend == "triton" draft_config = _apply_draft_moe_backend(tgt_config) assert draft_config.kernel_config.moe_backend == "triton" # Target config must be unaffected. - assert tgt_config.kernel_config.moe_backend == "flashinfer_trtllm" + assert tgt_config.kernel_config.moe_backend == target_moe_backend -def test_draft_model_moe_backend_inherits_target(): +def test_draft_model_moe_backend_inherits_target(monkeypatch: pytest.MonkeyPatch): """When moe_backend is not set in speculative_config, the draft should inherit the target's moe_backend.""" + target_moe_backend = _platform_moe_backend(monkeypatch) engine_args = EngineArgs( model="Qwen/Qwen3-1.7B", tensor_parallel_size=1, - moe_backend="flashinfer_cutlass", + moe_backend=target_moe_backend, speculative_config={ "model": "Qwen/Qwen3-0.6B", "method": "draft_model", @@ -250,11 +261,11 @@ def test_draft_model_moe_backend_inherits_target(): }, ) tgt_config: VllmConfig = engine_args.create_engine_config() - assert tgt_config.kernel_config.moe_backend == "flashinfer_cutlass" + assert tgt_config.kernel_config.moe_backend == target_moe_backend assert tgt_config.speculative_config.moe_backend is None draft_config = _apply_draft_moe_backend(tgt_config) - assert draft_config.kernel_config.moe_backend == "flashinfer_cutlass" + assert draft_config.kernel_config.moe_backend == target_moe_backend assert draft_config is tgt_config @@ -293,19 +304,23 @@ def test_draft_model_engine_args_rejects_invalid_tp_argname(): "tensor_parallel_size": 1, # <<< invalid arg name }, ) - with pytest.raises(ValueError): + with pytest.raises(ValueError, match="draft_tensor_parallel_size"): engine_args.create_engine_config() -def assert_draft_model_correctness(args: ArgsTest): +def assert_draft_model_correctness(args: ArgsTest, vllm_runner): """Compare the outputs using and not using speculative decoding. In the greedy decoding case, the outputs must match EXACTLY.""" test_prompts: list[Messages] = get_messages( dataset=args.dataset, n=args.num_prompts ) - spec_llm = LLM( - model=args.target_model, + with vllm_runner( + args.target_model, + block_size=None, + trust_remote_code=False, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), speculative_config={ "model": args.draft_model, "method": "draft_model", @@ -321,30 +336,41 @@ def assert_draft_model_correctness(args: ArgsTest): tensor_parallel_size=args.target_tensor_parallel_size, enforce_eager=args.enforce_eager, disable_log_stats=False, # enables get_metrics() - ) - - # we don't check the outputs, only check the metrics - spec_llm.chat(test_prompts, args.sampling_config) - metrics = spec_llm.get_metrics() - acceptance_rate: float = compute_acceptance_rate(metrics) - acceptance_len: float = compute_acceptance_len(metrics) + ) as spec_runner: + spec_llm = spec_runner.llm + + # we don't check the outputs, only check the metrics + spec_llm.chat(test_prompts, args.sampling_config) + metrics = spec_llm.get_metrics() + acceptance_rate: float = compute_acceptance_rate(metrics) + acceptance_len: float = compute_acceptance_len(metrics) + + # Need to evaluate after getting metrics to avoid polluting the AR + evaluate_llm_for_gsm8k( + spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy + ) - # Need to evaluate after getting metrics to avoid polluting the AR - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=args.expected_gsm8k_accuracy - ) + print( + f"spec-decode: target={args.target_model}, draft={args.draft_model}, " + f"temperature={args.sampling_config.temperature:.2f}, " + f"acceptance_rate={acceptance_rate:.2f}, " + f"acceptance_len={acceptance_len:.2f}, " + ) - print( - f"spec-decode: target={args.target_model}, draft={args.draft_model}, " - f"temperature={args.sampling_config.temperature:.2f}, " - f"acceptance_rate={acceptance_rate:.2f}, " - f"acceptance_len={acceptance_len:.2f}, " - ) + context = ( + f"draft_model target={args.target_model}, draft={args.draft_model}, " + f"eager={args.enforce_eager}, " + f"acceptance_rate={acceptance_rate:.3f}, " + f"acceptance_len={acceptance_len:.3f}" + ) + assert acceptance_rate >= args.expected_acceptance_rate, ( + f"{context}; expected acceptance_rate >= " + f"{args.expected_acceptance_rate:.3f}" + ) + assert acceptance_len >= args.expected_acceptance_len, ( + f"{context}; expected acceptance_len >= {args.expected_acceptance_len:.3f}" + ) + # draft_model supports async scheduling; assert it is active by default. + has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - assert acceptance_rate >= args.expected_acceptance_rate - assert acceptance_len >= args.expected_acceptance_len - has_async = spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - del spec_llm # CLEANUP - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() assert has_async, "Expected async_scheduling=True for draft_model spec decode" diff --git a/tests/v1/e2e/spec_decode/draft_model/test_lora.py b/tests/v1/e2e/spec_decode/draft_model/test_lora.py index 5decfbc00862..4ffc07f88cdd 100644 --- a/tests/v1/e2e/spec_decode/draft_model/test_lora.py +++ b/tests/v1/e2e/spec_decode/draft_model/test_lora.py @@ -8,12 +8,14 @@ import pytest import torch -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig from vllm.lora.request import LoRARequest from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed +from ..utils import assert_request_outputs_match + LORA_TEST_PROMPT_MAP: dict[str, str] = {} LORA_TEST_PROMPT_MAP["premjatin/qwen-linear-algebra-coder"] = """ @@ -36,7 +38,6 @@ SEED = 42 -@pytest.mark.skipif(not current_platform.is_cuda(), reason="CUDA not available") @pytest.mark.parametrize( "model_setup", [ @@ -49,9 +50,13 @@ ) ], ) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), reason="Requires CUDA or ROCm" +) def test_batch_inference_correctness( monkeypatch: pytest.MonkeyPatch, model_setup: tuple[str, str, str, str, int], + vllm_runner, ): """ Compare the outputs of a LLM with only Lora and a LLM with both SD and Lora. @@ -60,10 +65,11 @@ def test_batch_inference_correctness( """ with monkeypatch.context() as m: # Disable randomness - m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8") + if current_platform.is_cuda(): + m.setenv("CUBLAS_WORKSPACE_CONFIG", ":4096:8") set_random_seed(SEED) - torch.backends.cudnn.benchmark = False - torch.backends.cudnn.deterministic = True + m.setattr(torch.backends.cudnn, "benchmark", False) + m.setattr(torch.backends.cudnn, "deterministic", True) method, model_name, spec_model_name, lora_path, tp_size = model_setup @@ -74,8 +80,9 @@ def test_batch_inference_correctness( ) # without speculative decoding - ref_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + block_size=None, trust_remote_code=True, tensor_parallel_size=tp_size, max_model_len=2048, @@ -84,18 +91,16 @@ def test_batch_inference_correctness( max_loras=1, max_cpu_loras=1, max_lora_rank=16, - ) - try: - ref_outputs = ref_llm.generate( + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + ) as ref_runner: + ref_outputs = ref_runner.llm.generate( prompts, sampling_params, lora_request=lora_request ) - finally: - del ref_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - lora_spec_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + block_size=None, trust_remote_code=True, tensor_parallel_size=tp_size, speculative_config={ @@ -110,30 +115,22 @@ def test_batch_inference_correctness( max_loras=1, max_cpu_loras=1, max_lora_rank=16, - ) - try: - lora_spec_outputs = lora_spec_llm.generate( + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + ) as spec_runner: + lora_spec_outputs = spec_runner.llm.generate( prompts, sampling_params, lora_request=lora_request ) - matches = 0 - for ref_output, spec_output in zip(ref_outputs, lora_spec_outputs): - if ref_output.outputs[0].text == spec_output.outputs[0].text: - matches += 1 - else: - print(f"ref_output: {ref_output.outputs[0].text}") - print(f"spec_output: {spec_output.outputs[0].text}") - - # Heuristic threshold: under greedy verification, the spec-decode - # output should equal the non-spec output (modulo FP noise from the - # target's verify-path matmul running at seqlen - # num_speculative_tokens+1 vs 1). 90% leaves slack for that noise. - threshold = int(0.90 * len(ref_outputs)) - print(f"match ratio: {matches}/{len(ref_outputs)}") - assert matches > threshold, ( - f"match ratio {matches}/{len(ref_outputs)} <= {threshold}" - ) - finally: - del lora_spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + # Under greedy verification, the spec-decode output should equal the + # non-spec output (modulo FP noise from the target's verify-path matmul + # running at seqlen num_speculative_tokens+1 vs 1). 90% leaves slack. + assert_request_outputs_match( + ref_outputs, + lora_spec_outputs, + required_matches=int(0.90 * len(ref_outputs)) + 1, + context=( + f"LoRA + {method}: target={model_name}, draft={spec_model_name}, " + f"adapter={lora_path}" + ), + ) diff --git a/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py index c5c32d80e594..bb8ff8e57f27 100644 --- a/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py +++ b/tests/v1/e2e/spec_decode/eagle/test_eagle_correctness.py @@ -52,6 +52,7 @@ def test_eagle_correctness_light( enable_chunked_prefill: bool, model_impl: str, attn_backend: str, + vllm_runner, ): _run_eagle_correctness( monkeypatch, @@ -62,6 +63,7 @@ def test_eagle_correctness_light( enable_chunked_prefill, model_impl, attn_backend, + vllm_runner, ) @@ -89,11 +91,12 @@ def test_eagle_correctness_light( "transformers", 0.8, # TODO(hmellor): figure out why memory usage is so high - marks=pytest.mark.skip( - reason="Feature is experimental and uses too much memory in CI", + marks=pytest.mark.skipif( + not current_platform.is_rocm(), + reason="Feature is experimental and uses too much memory in CUDA CI", ), ), - pytest.param( + ( ( "eagle3", "Qwen/Qwen3-VL-8B-Instruct", @@ -104,9 +107,6 @@ def test_eagle_correctness_light( False, "auto", 0.8, - marks=pytest.mark.skip( - reason="architecture of its eagle3 is LlamaForCausalLMEagle3" - ), ), pytest.param( ( @@ -119,8 +119,11 @@ def test_eagle_correctness_light( False, "auto", 0.7, + # TODO: Re-measure the reference threshold when re-enabling this case. + # Text-only mode starts but loses target correctness with ROCm/TRITON; + # full multimodal profiling currently fails during engine startup. marks=pytest.mark.skip( - reason="Skipping due to its head_dim not being a multiple of 32" + reason="Qwen2.5-VL Eagle3 is not yet reliable in this test setup" ), ), ( @@ -154,6 +157,7 @@ def test_eagle_correctness_medium( enable_chunked_prefill: bool, model_impl: str, attn_backend: str, + vllm_runner, ): _run_eagle_correctness( monkeypatch, @@ -164,4 +168,5 @@ def test_eagle_correctness_medium( enable_chunked_prefill, model_impl, attn_backend, + vllm_runner, ) diff --git a/tests/v1/e2e/spec_decode/eagle/utils.py b/tests/v1/e2e/spec_decode/eagle/utils.py index 998be71af685..29127df85f5c 100644 --- a/tests/v1/e2e/spec_decode/eagle/utils.py +++ b/tests/v1/e2e/spec_decode/eagle/utils.py @@ -1,16 +1,17 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from typing import Any + import pytest -import torch -from tests.utils import wait_for_rocm_memory_to_settle -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig from vllm.platforms import current_platform from ..utils import ( _skip_if_insufficient_gpus_for_tp, + assert_request_outputs_match, evaluate_llm_for_gsm8k, get_test_prompts, ) @@ -25,26 +26,24 @@ def _run_eagle_correctness( enable_chunked_prefill: bool, model_impl: str, attn_backend: str, + vllm_runner, ): """ Compare the outputs of an original LLM and a speculative LLM which should be the same when using eagle speculative decoding. """ - if model_impl == "transformers": - import transformers - from packaging.version import Version - - installed = Version(transformers.__version__) - required = Version("5.0.0") - if installed < required: - pytest.skip( - "Eagle3 with the Transformers modeling backend requires " - f"transformers>={required}, but got {installed}" - ) + method, model_name, spec_model_name, tp_size = model_setup + _skip_if_insufficient_gpus_for_tp(tp_size) test_prompts = get_test_prompts(mm_enabled) - if "Llama-4-Scout" in model_setup[1] and attn_backend == "FLASH_ATTN": + extra_kwargs: dict[str, Any] = {} + if not mm_enabled and "Qwen3-VL" in model_name: + # These cases only exercise text generation. Avoid profiling an unused + # vision tower, which adds substantial memory to both reference runs. + extra_kwargs["limit_mm_per_prompt"] = {"image": 0, "video": 0} + + if "Llama-4-Scout" in model_name and attn_backend == "FLASH_ATTN": if current_platform.is_rocm(): print( "FLASH_ATTN for spec_decode not supported on " @@ -66,38 +65,36 @@ def _run_eagle_correctness( m.setenv("VLLM_MLA_DISABLE", "1") if attn_backend == "ROCM_AITER_FA" and current_platform.is_rocm(): - if "deepseek" in model_setup[1].lower(): + if "deepseek" in model_name.lower(): m.setenv("VLLM_ROCM_USE_AITER", "1") m.delenv("VLLM_MLA_DISABLE", raising=False) attention_config = {"backend": "ROCM_AITER_MLA"} else: m.setenv("VLLM_ROCM_USE_AITER", "1") - method, model_name, spec_model_name, tp_size = model_setup - _skip_if_insufficient_gpus_for_tp(tp_size) - max_model_len = 2048 max_num_batched_tokens = 128 if enable_chunked_prefill else max_model_len - ref_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + block_size=None, + trust_remote_code=False, max_model_len=max_model_len, tensor_parallel_size=tp_size, attention_config=attention_config, - ) - evaluate_llm_for_gsm8k( - ref_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - ref_outputs = ref_llm.chat(test_prompts, sampling_config) - del ref_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - # ROCm frees VRAM lazily; wait so the spec engine started right after - # does not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle() - - spec_llm = LLM( - model=model_name, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + **extra_kwargs, + ) as ref_runner: + evaluate_llm_for_gsm8k( + ref_runner.llm, + expected_accuracy_threshold=expected_accuracy_threshold, + ) + ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config) + + with vllm_runner( + model_name, + block_size=None, trust_remote_code=True, tensor_parallel_size=tp_size, speculative_config={ @@ -111,27 +108,29 @@ def _run_eagle_correctness( enable_chunked_prefill=enable_chunked_prefill, model_impl=model_impl, attention_config=attention_config, + compilation_config=CompilationConfig(), + **extra_kwargs, + ) as spec_runner: + # EAGLE/EAGLE3 supports async scheduling by default. + has_async = ( + spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling + ) + assert has_async, ( + f"Expected async scheduling for {method}: target={model_name}, " + f"draft={spec_model_name}, backend={attn_backend}; got {has_async}" + ) + evaluate_llm_for_gsm8k( + spec_runner.llm, + expected_accuracy_threshold=expected_accuracy_threshold, + ) + spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config) + + assert_request_outputs_match( + ref_outputs, + spec_outputs, + required_matches=int(0.6 * len(ref_outputs)) + 1, + context=( + f"{method} target={model_name}, draft={spec_model_name}, " + f"backend={attn_backend}" + ), ) - # EAGLE/EAGLE3 supports async scheduling; assert it is active by default. - assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - spec_outputs = spec_llm.chat(test_prompts, sampling_config) - matches = 0 - misses = 0 - for ref_output, spec_output in zip(ref_outputs, spec_outputs): - if ref_output.outputs[0].text == spec_output.outputs[0].text: - matches += 1 - else: - misses += 1 - print(f"ref_output: {ref_output.outputs[0].text}") - print(f"spec_output: {spec_output.outputs[0].text}") - - assert matches > int(0.6 * len(ref_outputs)) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - # ROCm frees VRAM lazily; wait so the next parametrization's engine does - # not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle() diff --git a/tests/v1/e2e/spec_decode/mtp/test_mtp.py b/tests/v1/e2e/spec_decode/mtp/test_mtp.py index c70a6fd3fc8f..a13152dda4d4 100644 --- a/tests/v1/e2e/spec_decode/mtp/test_mtp.py +++ b/tests/v1/e2e/spec_decode/mtp/test_mtp.py @@ -4,15 +4,15 @@ from typing import Any import pytest -import torch from tests.utils import single_gpu_only -from vllm import LLM, SamplingParams -from vllm.distributed import cleanup_dist_env_and_memory +from vllm import SamplingParams +from vllm.config import CompilationConfig from vllm.platforms import current_platform from ..utils import ( _skip_if_insufficient_gpus_for_tp, + assert_request_outputs_match, evaluate_llm_for_gsm8k, get_test_prompts, ) @@ -21,9 +21,10 @@ @pytest.mark.parametrize( ["model_setup", "mm_enabled", "expected_accuracy_threshold"], [ - (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1), False, 0.5), # ref: 65%-70% + # Reference accuracy: 65%-70%. + (("mtp", "XiaomiMiMo/MiMo-7B-Base", 1, None), False, 0.5), pytest.param( - ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1), + ("mtp", "ZixiQi/DeepSeek-V3-4layers-MTP-FP8", 1, None), False, 0.0, marks=pytest.mark.skipif( @@ -32,7 +33,7 @@ ), ), # dummy model ( - ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1), + ("mtp", "Qwen/Qwen3.5-0.8B-Base", 1, None), False, 0.20, ), # hybrid + MTP, ref: ~34%-35% @@ -48,9 +49,10 @@ def test_mtp_correctness( monkeypatch: pytest.MonkeyPatch, sampling_config: SamplingParams, - model_setup: tuple[str, str, int] | tuple[str, str, int, str], + model_setup: tuple[str, str, int, str | None], mm_enabled: bool, expected_accuracy_threshold: float, + vllm_runner, ): """ Compare the outputs of a original LLM and a speculative LLM @@ -59,18 +61,14 @@ def test_mtp_correctness( 6/10 output tokens match exactly, and that the GSM8k accuracy is above a precomputed reference threshold for each model. """ + method, model_name, tp_size, draft_model = model_setup + _skip_if_insufficient_gpus_for_tp(tp_size) + # Generate test prompts inside the function instead of using fixture test_prompts = get_test_prompts(mm_enabled) with monkeypatch.context() as m: m.setenv("VLLM_MLA_DISABLE", "1") - if len(model_setup) == 4: - method, model_name, tp_size, draft_model = model_setup - else: - method, model_name, tp_size = model_setup - draft_model = None - _skip_if_insufficient_gpus_for_tp(tp_size) - attn_backend = "TRITON_ATTN" if current_platform.is_rocm() else "auto" # Skip multimodal profiling for models that don't need it in this test. @@ -90,21 +88,22 @@ def test_mtp_correctness( f"got {transformers.__version__}" ) - ref_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + block_size=None, max_model_len=2048, tensor_parallel_size=tp_size, trust_remote_code=True, attention_backend=attn_backend, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), **extra_kwargs, - ) - ref_outputs = ref_llm.chat(test_prompts, sampling_config) - evaluate_llm_for_gsm8k( - ref_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - del ref_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + ) as ref_runner: + ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config) + evaluate_llm_for_gsm8k( + ref_runner.llm, + expected_accuracy_threshold=expected_accuracy_threshold, + ) speculative_config: dict[str, Any] = { "method": method, @@ -115,34 +114,37 @@ def test_mtp_correctness( speculative_config["model"] = draft_model speculative_config["num_speculative_tokens"] = 2 - spec_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + block_size=None, trust_remote_code=True, tensor_parallel_size=tp_size, speculative_config=speculative_config, max_model_len=2048, attention_backend=attn_backend, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), **extra_kwargs, - ) - # MTP supports async scheduling; assert it is active by default. - assert spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - evaluate_llm_for_gsm8k( - spec_llm, expected_accuracy_threshold=expected_accuracy_threshold - ) - spec_outputs = spec_llm.chat(test_prompts, sampling_config) - matches = 0 - misses = 0 - for ref_output, spec_output in zip(ref_outputs, spec_outputs): - if ref_output.outputs[0].text == spec_output.outputs[0].text: - matches += 1 - else: - misses += 1 - print(f"ref_output: {ref_output.outputs[0].text}") - print(f"spec_output: {spec_output.outputs[0].text}") + ) as spec_runner: + # MTP supports async scheduling by default. + has_async = ( + spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling + ) + assert has_async, ( + f"Expected async scheduling for {method}: target={model_name}, " + f"draft={draft_model}; got {has_async}" + ) + evaluate_llm_for_gsm8k( + spec_runner.llm, + expected_accuracy_threshold=expected_accuracy_threshold, + ) + spec_outputs = spec_runner.llm.chat(test_prompts, sampling_config) # Heuristic: expect at least 80% of the prompts to match exactly # Upon failure, inspect the outputs to check for inaccuracy. - assert matches > int(0.8 * len(ref_outputs)) - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + assert_request_outputs_match( + ref_outputs, + spec_outputs, + required_matches=int(0.8 * len(ref_outputs)) + 1, + context=f"{method} target={model_name}, draft={draft_model}", + ) diff --git a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py index b6dc946a98d7..d3627a8703b3 100644 --- a/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py +++ b/tests/v1/e2e/spec_decode/ngram_suffix/test_ngram_suffix.py @@ -2,15 +2,17 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import pytest -import torch from tests.utils import single_gpu_only -from vllm import LLM, SamplingParams +from vllm import SamplingParams from vllm.config import CompilationConfig -from vllm.distributed import cleanup_dist_env_and_memory from vllm.platforms import current_platform -from ..utils import evaluate_llm_for_gsm8k, get_test_prompts +from ..utils import ( + evaluate_llm_for_gsm8k, + get_spec_decode_metric_value, + get_test_prompts, +) @pytest.fixture @@ -43,6 +45,7 @@ def test_ngram_and_suffix_correctness( ): with vllm_runner( model_name, + block_size=None, # Keep LLM defaults; VllmRunner only provides lifecycle cleanup here. trust_remote_code=False, enable_chunked_prefill=None, @@ -59,6 +62,7 @@ def test_ngram_and_suffix_correctness( @single_gpu_only def test_ngram_gpu_default_with_async_scheduling( async_scheduling: bool, + vllm_runner, ): """ Test ngram_gpu speculative decoding (k=3) correctness with and without @@ -66,8 +70,10 @@ def test_ngram_gpu_default_with_async_scheduling( Uses Qwen/Qwen3-8B (ref GSM8K accuracy: 87%-92%). """ qwen3_model = "Qwen/Qwen3-8B" - spec_llm = LLM( - model=qwen3_model, + with vllm_runner( + qwen3_model, + block_size=None, + trust_remote_code=False, speculative_config={ "method": "ngram_gpu", "prompt_lookup_max": 3, @@ -76,22 +82,22 @@ def test_ngram_gpu_default_with_async_scheduling( }, max_model_len=4096, async_scheduling=async_scheduling, - ) - # Assert the resolved async_scheduling config matches what was requested. - assert ( - spec_llm.llm_engine.vllm_config.scheduler_config.async_scheduling - == async_scheduling - ) - evaluate_llm_for_gsm8k(spec_llm, expected_accuracy_threshold=0.8) - del spec_llm - cleanup_dist_env_and_memory() + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + ) as spec_runner: + # Assert the resolved async_scheduling config matches what was requested. + assert ( + spec_runner.llm.llm_engine.vllm_config.scheduler_config.async_scheduling + == async_scheduling + ) + evaluate_llm_for_gsm8k(spec_runner.llm, expected_accuracy_threshold=0.8) @single_gpu_only def test_suffix_decoding_acceptance( - monkeypatch: pytest.MonkeyPatch, sampling_config: SamplingParams, model_name: str, + vllm_runner, ): """ Check that suffix decoding caching takes effect and improves acceptance @@ -99,8 +105,10 @@ def test_suffix_decoding_acceptance( """ test_prompts = get_test_prompts(mm_enabled=False) - spec_llm = LLM( - model=model_name, + with vllm_runner( + model_name, + block_size=None, + trust_remote_code=False, speculative_config={ "method": "suffix", "suffix_decoding_max_spec_factor": 2.0, @@ -108,40 +116,60 @@ def test_suffix_decoding_acceptance( }, max_model_len=1024, disable_log_stats=False, - ) - - # Run several times and check that the accepted tokens increase. - num_draft = [] - num_accept = [] - for i in range(10): # Run multiple times to warm up the cache. - spec_llm.chat(test_prompts, sampling_config) - # Collect draft and acceptance stats. - metrics = spec_llm.get_metrics() - for metric in metrics: - if metric.name == "vllm:spec_decode_num_draft_tokens": - num_draft.append(metric.value) - if metric.name == "vllm:spec_decode_num_accepted_tokens": - num_accept.append(metric.value) + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + ) as spec_runner: + # Run several times and check that the accepted tokens increase. + num_draft = [] + num_accept = [] + for _ in range(10): # Run multiple times to warm up the cache. + spec_runner.llm.chat(test_prompts, sampling_config) + # Collect draft and acceptance stats. + metrics = spec_runner.llm.get_metrics() + num_draft.append( + get_spec_decode_metric_value( + metrics, "vllm:spec_decode_num_draft_tokens" + ) + ) + num_accept.append( + get_spec_decode_metric_value( + metrics, "vllm:spec_decode_num_accepted_tokens" + ) + ) # Calculate the acceptance rates for the first and last runs. first_accept_tokens = num_accept[0] first_draft_tokens = num_draft[0] + assert first_draft_tokens > 0, ( + "Suffix decoder produced no draft tokens on the first run: " + f"accepted={first_accept_tokens}, drafted={first_draft_tokens}" + ) first_accept_rate = first_accept_tokens / first_draft_tokens # Take the diff since the stats are cumulative. last_accept_tokens = num_accept[-1] - num_accept[-2] last_draft_tokens = num_draft[-1] - num_draft[-2] + assert last_draft_tokens > 0, ( + "Suffix decoder produced no draft tokens on the last run: " + f"accepted_delta={last_accept_tokens}, drafted_delta={last_draft_tokens}; " + f"cumulative_drafted={num_draft[-2:]}" + ) last_accept_rate = last_accept_tokens / last_draft_tokens + summary = ( + f"first accepted/drafted={first_accept_tokens}/{first_draft_tokens} " + f"(rate={first_accept_rate:.3f}); last delta accepted/drafted=" + f"{last_accept_tokens}/{last_draft_tokens} (rate={last_accept_rate:.3f})" + ) # Expect the acceptance length to improve. - assert first_accept_tokens < last_accept_tokens + assert first_accept_tokens < last_accept_tokens, ( + f"Expected accepted tokens to increase after cache warmup; {summary}" + ) # Expect the acceptance rate to improve. - assert first_accept_rate < last_accept_rate + assert first_accept_rate < last_accept_rate, ( + f"Expected acceptance rate to increase after cache warmup; {summary}" + ) # Heuristic: expect at least 80.0% acceptance rate at the end. - assert last_accept_rate > 0.80 - - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + assert last_accept_rate > 0.80, f"Expected final acceptance rate > 0.80; {summary}" diff --git a/tests/v1/e2e/spec_decode/speculators/test_speculators.py b/tests/v1/e2e/spec_decode/speculators/test_speculators.py index 27fa21fd7669..990f6190751a 100644 --- a/tests/v1/e2e/spec_decode/speculators/test_speculators.py +++ b/tests/v1/e2e/spec_decode/speculators/test_speculators.py @@ -7,14 +7,20 @@ from vllm import SamplingParams from vllm.config import CompilationConfig -from ..utils import evaluate_llm_for_gsm8k, get_test_prompts +from ..utils import ( + assert_request_outputs_match, + evaluate_llm_for_gsm8k, + get_test_prompts, +) @pytest.mark.parametrize( ["model_path", "expected_accuracy_threshold"], [ - ("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.7), # ref: 75%-80% - ("RedHatAI/Qwen3-8B-speculator.eagle3", 0.8), # ref: 87%-92% + # Measured reference: 75%-80%. + ("RedHatAI/Llama-3.1-8B-Instruct-speculator.eagle3", 0.72), + # Measured reference: 87%-92%. + ("RedHatAI/Qwen3-8B-speculator.eagle3", 0.84), ], ids=["llama3_eagle3_speculator", "qwen3_eagle3_speculator"], ) @@ -49,6 +55,7 @@ def test_speculators_model_integration( # First run: Direct speculator model (simplified integration) with vllm_runner( model_path, + block_size=None, trust_remote_code=False, enable_chunked_prefill=None, compilation_config=CompilationConfig(), @@ -82,6 +89,7 @@ def test_speculators_model_integration( # Second run: Reference without speculative decoding with vllm_runner( verifier_model, + block_size=None, trust_remote_code=False, enable_chunked_prefill=None, compilation_config=CompilationConfig(), @@ -90,15 +98,10 @@ def test_speculators_model_integration( ) as ref_runner: ref_outputs = ref_runner.llm.chat(test_prompts, sampling_config) - # Compare outputs - matches = sum( - 1 - for ref, spec in zip(ref_outputs, spec_outputs) - if ref.outputs[0].text == spec.outputs[0].text - ) - # Heuristic: expect at least 66% of prompts to match exactly - assert matches >= int(0.66 * len(ref_outputs)), ( - f"Only {matches}/{len(ref_outputs)} outputs matched. " - f"Expected at least {int(0.66 * len(ref_outputs))} matches." + assert_request_outputs_match( + ref_outputs, + spec_outputs, + required_matches=int(0.66 * len(ref_outputs)), + context=f"speculator={model_path}, verifier={verifier_model}", ) diff --git a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py index 186faa6529b1..236997261675 100644 --- a/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py +++ b/tests/v1/e2e/spec_decode/test_mtp_parallel_load.py @@ -11,8 +11,10 @@ from tests.utils import multi_gpu_marks from vllm import LLM, SamplingParams +from vllm.config import CompilationConfig from vllm.distributed import cleanup_dist_env_and_memory from vllm.engine.arg_utils import AsyncEngineArgs +from vllm.platforms import current_platform from vllm.sampling_params import RequestOutputKind from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.metrics.reader import Metric @@ -59,7 +61,7 @@ class InlineConfig: pp_size=2, skip_reason=( "DeepSeek MTP pipeline-parallel support is in flight upstream; " - "see https://github.com/vllm-project/vllm/pull/38104" + "see https://github.com/vllm-project/vllm/pull/44698" ), ), ] @@ -109,7 +111,11 @@ def _generate_token_ids_inline(llm: LLM) -> tuple[int, ...]: SamplingParams(temperature=0.0, max_tokens=MAX_TOKENS, ignore_eos=True), ) assert outputs and outputs[0].outputs, "expected one completion" - return tuple(outputs[0].outputs[0].token_ids) + token_ids = tuple(outputs[0].outputs[0].token_ids) + assert len(token_ids) == MAX_TOKENS, ( + f"expected {MAX_TOKENS} generated tokens, got {len(token_ids)}: {token_ids}" + ) + return token_ids def _spec_decode_num_drafts(metrics: list[Metric]) -> int: @@ -121,6 +127,16 @@ def _spec_decode_num_drafts(metrics: list[Metric]) -> int: return int(counter.value) +def _enable_batch_invariance(monkeypatch: pytest.MonkeyPatch) -> None: + # This test exercises DeepSeek MLA prefill. ROCm's upstream FlashAttention + # varlen API lacks the num_splits argument used by batch invariance, so + # explicitly override any ambient setting on that path. + if current_platform.is_cuda(): + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + else: + monkeypatch.setenv("VLLM_BATCH_INVARIANT", "0") + + @pytest.mark.parametrize( "config", [ @@ -131,6 +147,7 @@ def _spec_decode_num_drafts(metrics: list[Metric]) -> int: def test_deepseek_mtp_load_inline( monkeypatch: pytest.MonkeyPatch, config: InlineConfig, + vllm_runner, ): """MTP loads and drafts under TP/EP/EPLB; spec output matches no-spec greedy.""" if config.skip_reason is not None: @@ -138,24 +155,30 @@ def test_deepseek_mtp_load_inline( # Reduces run-to-run nondeterminism so spec / no-spec stay close; # see tests/v1/distributed/test_eagle_dp.py for the same pattern. - monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") - - spec_llm = LLM(**_inline_kwargs(config, with_spec=True)) - try: - spec_tokens = _generate_token_ids_inline(spec_llm) - n_drafts = _spec_decode_num_drafts(spec_llm.get_metrics()) - finally: - del spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() - - no_spec_llm = LLM(**_inline_kwargs(config, with_spec=False)) - try: - no_spec_tokens = _generate_token_ids_inline(no_spec_llm) - finally: - del no_spec_llm - torch.accelerator.empty_cache() - cleanup_dist_env_and_memory() + _enable_batch_invariance(monkeypatch) + + spec_kwargs = _inline_kwargs(config, with_spec=True) + spec_model = spec_kwargs.pop("model") + with vllm_runner( + spec_model, + block_size=None, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + **spec_kwargs, + ) as spec_runner: + spec_tokens = _generate_token_ids_inline(spec_runner.llm) + n_drafts = _spec_decode_num_drafts(spec_runner.llm.get_metrics()) + + no_spec_kwargs = _inline_kwargs(config, with_spec=False) + no_spec_model = no_spec_kwargs.pop("model") + with vllm_runner( + no_spec_model, + block_size=None, + enable_chunked_prefill=None, + compilation_config=CompilationConfig(), + **no_spec_kwargs, + ) as no_spec_runner: + no_spec_tokens = _generate_token_ids_inline(no_spec_runner.llm) # Non-vacuity: a silently-broken MTP drafter that falls back to # verifier-only decoding would still produce matching output below. @@ -185,7 +208,7 @@ async def test_deepseek_mtp_load_dp(monkeypatch: pytest.MonkeyPatch, dp_size: in if torch.accelerator.device_count() < dp_size: pytest.skip(f"dp{dp_size} requires at least {dp_size} GPUs") - monkeypatch.setenv("VLLM_BATCH_INVARIANT", "1") + _enable_batch_invariance(monkeypatch) base_args = AsyncEngineArgs( model=DEEPSEEK_MTP_MAIN_RANDOM, @@ -225,7 +248,10 @@ async def _generate(args: AsyncEngineArgs, request_id: str) -> tuple[int, ...]: sampling_params=sampling_params, ): token_ids = tuple(out.outputs[0].token_ids) - assert len(token_ids) == MAX_TOKENS + assert len(token_ids) == MAX_TOKENS, ( + f"{request_id}: expected {MAX_TOKENS} generated tokens, " + f"got {len(token_ids)}: {token_ids}" + ) return token_ids raise AssertionError("AsyncLLM produced no output") finally: diff --git a/tests/v1/e2e/spec_decode/utils.py b/tests/v1/e2e/spec_decode/utils.py index ccd5232eaaea..0ba7edd1c900 100644 --- a/tests/v1/e2e/spec_decode/utils.py +++ b/tests/v1/e2e/spec_decode/utils.py @@ -2,7 +2,7 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import random -from collections.abc import Iterable +from collections.abc import Iterable, Sequence from typing import Any import pytest @@ -12,11 +12,13 @@ from vllm import LLM, SamplingParams from vllm.assets.base import VLLM_S3_BUCKET_URL from vllm.assets.image import VLM_IMAGES_DIR +from vllm.outputs import RequestOutput +from vllm.utils.torch_utils import set_random_seed from vllm.v1.metrics.reader import Metric def _skip_if_insufficient_gpus_for_tp(tp_size: int): - """Skip test if available GPUs < tp_size on ROCm.""" + """Skip if fewer than ``tp_size`` accelerators are available.""" available_gpus = torch.accelerator.device_count() if available_gpus < tp_size: pytest.skip( @@ -43,7 +45,7 @@ def get_test_prompts(mm_enabled: bool, num_prompts: int = 100) -> list[Messages] # Generate a mixed batch of prompts, some of which can be easily # predicted by n-gram matching and some which likely cannot. - random.seed(0) + set_random_seed(0) for _ in range(num_repeat_prompts): word_choices = ["test", "temp", "hello", "where"] word = random.choice(word_choices) @@ -120,20 +122,88 @@ def evaluate_llm_for_gsm8k(llm: LLM, expected_accuracy_threshold: float = 0.70) ) +def assert_request_outputs_match( + ref_outputs: Sequence[RequestOutput], + spec_outputs: Sequence[RequestOutput], + *, + required_matches: int, + context: str, + max_mismatches: int = 3, +) -> None: + """Assert a thresholded exact-text match with bounded failure diagnostics.""" + assert ref_outputs, f"{context}: no reference outputs" + assert len(ref_outputs) == len(spec_outputs), ( + f"{context}: output count differs: " + f"reference={len(ref_outputs)}, speculative={len(spec_outputs)}" + ) + assert 0 <= required_matches <= len(ref_outputs), ( + f"{context}: invalid required_matches={required_matches} for " + f"{len(ref_outputs)} outputs" + ) + + mismatches: list[str] = [] + matches = 0 + for index, (ref_output, spec_output) in enumerate(zip(ref_outputs, spec_outputs)): + assert ref_output.outputs, ( + f"{context}: reference output {index} has no candidate" + ) + assert spec_output.outputs, ( + f"{context}: speculative output {index} has no candidate" + ) + ref_candidate = ref_output.outputs[0] + spec_candidate = spec_output.outputs[0] + if ref_candidate.text == spec_candidate.text: + matches += 1 + elif len(mismatches) < max_mismatches: + mismatches.append( + f"[{index}] ref_text={ref_candidate.text[:240]!r}, " + f"spec_text={spec_candidate.text[:240]!r}\n" + f" ref_token_ids={list(ref_candidate.token_ids)[:64]}\n" + f" spec_token_ids={list(spec_candidate.token_ids)[:64]}" + ) + + print( + f"{context}: exact text matches={matches}/{len(ref_outputs)} " + f"(required={required_matches})" + ) + mismatch_summary = "\n".join(mismatches) or "no mismatches captured" + assert matches >= required_matches, ( + f"{context}: only {matches}/{len(ref_outputs)} outputs matched; " + f"required at least {required_matches}. First mismatches:\n" + f"{mismatch_summary}" + ) + + +def get_spec_decode_metric_value(metrics: Sequence[Metric], metric_name: str) -> float: + """Get a spec-decode metric with an actionable error when stats are absent.""" + name2metric = {metric.name: metric for metric in metrics} + metric = name2metric.get(metric_name) + assert metric is not None, ( + f"Missing metric {metric_name!r}. Ensure disable_log_stats=False. " + "Available spec-decode metrics: " + f"{sorted(name for name in name2metric if 'spec_decode' in name) or ['']}" + ) + return float(metric.value) + + def compute_acceptance_rate( metrics: list[Metric], prev_metrics: list[Metric] | None = None ) -> float: - name2metric = {metric.name: metric for metric in metrics} - n_draft_toks = name2metric["vllm:spec_decode_num_draft_tokens"].value + n_draft_toks = get_spec_decode_metric_value( + metrics, "vllm:spec_decode_num_draft_tokens" + ) if n_draft_toks == 0: return float("nan") - n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value + n_accepted_toks = get_spec_decode_metric_value( + metrics, "vllm:spec_decode_num_accepted_tokens" + ) if prev_metrics is not None: - prev_name2metric = {metric.name: metric for metric in prev_metrics} - n_draft_toks -= prev_name2metric["vllm:spec_decode_num_draft_tokens"].value - n_accepted_toks -= prev_name2metric[ - "vllm:spec_decode_num_accepted_tokens" - ].value + n_draft_toks -= get_spec_decode_metric_value( + prev_metrics, "vllm:spec_decode_num_draft_tokens" + ) + n_accepted_toks -= get_spec_decode_metric_value( + prev_metrics, "vllm:spec_decode_num_accepted_tokens" + ) if n_draft_toks <= 0: return float("nan") return n_accepted_toks / n_draft_toks @@ -142,17 +212,19 @@ def compute_acceptance_rate( def compute_acceptance_len( metrics: list[Metric], prev_metrics: list[Metric] | None = None ) -> float: - name2metric = {metric.name: metric for metric in metrics} - n_drafts = name2metric["vllm:spec_decode_num_drafts"].value - n_accepted_toks = name2metric["vllm:spec_decode_num_accepted_tokens"].value + n_drafts = get_spec_decode_metric_value(metrics, "vllm:spec_decode_num_drafts") + n_accepted_toks = get_spec_decode_metric_value( + metrics, "vllm:spec_decode_num_accepted_tokens" + ) if n_drafts == 0: return 1 if prev_metrics is not None: - prev_name2metric = {metric.name: metric for metric in prev_metrics} - n_drafts -= prev_name2metric["vllm:spec_decode_num_drafts"].value - n_accepted_toks -= prev_name2metric[ - "vllm:spec_decode_num_accepted_tokens" - ].value + n_drafts -= get_spec_decode_metric_value( + prev_metrics, "vllm:spec_decode_num_drafts" + ) + n_accepted_toks -= get_spec_decode_metric_value( + prev_metrics, "vllm:spec_decode_num_accepted_tokens" + ) if n_drafts <= 0: return 1 return 1 + (n_accepted_toks / n_drafts) diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 262675671f4f..fa8c282007b2 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -121,7 +121,7 @@ class SpeculativeConfig: attention_backend: AttentionBackendEnum | None = None """Attention backend to use for the draft model. When `None`, the backend is automatically selected. Useful when the drafter requires a different attention - backend (e.g. DFlash needs a non-causal-capable backend like FLASH_ATTN).""" + backend (e.g. DFlash needs a backend that supports non-causal attention).""" kv_cache_dtype: CacheDType | None = None """KV cache dtype for the draft model. When `None`, the draft inherits the target model's `--kv-cache-dtype`.""" From 33584901bab6900045b504843a57e172205d3ff1 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 11 Aug 2026 10:00:12 +0800 Subject: [PATCH 22/30] =?UTF-8?q?[Frontend]=20Add=20content=5Fparts=20to?= =?UTF-8?q?=20/inference/v1/generate=20for=20raw=20multim=E2=80=A6=20(#514?= =?UTF-8?q?78)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Allen Shen Co-authored-by: Claude Opus 4.6 (1M context) Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- .../server/src/routes/inference/generate.rs | 25 +++++- .../src/routes/inference/generate/convert.rs | 7 +- .../src/routes/inference/generate/types.rs | 3 + rust/src/server/src/routes/render.rs | 1 + .../test_serving_multimodal_tokens.py | 84 +++++++++++++++++++ .../scale_out/token_in_token_out/protocol.py | 10 +++ .../scale_out/token_in_token_out/serving.py | 27 +++++- 7 files changed, 150 insertions(+), 7 deletions(-) diff --git a/rust/src/server/src/routes/inference/generate.rs b/rust/src/server/src/routes/inference/generate.rs index 57d2724cb68c..61fed421d38c 100644 --- a/rust/src/server/src/routes/inference/generate.rs +++ b/rust/src/server/src/routes/inference/generate.rs @@ -45,14 +45,31 @@ use crate::utils::resolve_request_context; pub async fn generate( State(state): State>, headers: HeaderMap, - ValidatedJson(body): ValidatedJson, + ValidatedJson(mut body): ValidatedJson, ) -> Response { let request_context = resolve_request_context(&headers, body.request_id.as_deref()); let lora_resolution = state.resolve_model_with_loras(body.model.as_deref()).await; - let prepared = match prepare_generate_request(body, &lora_resolution, request_context) { - Ok(prepared) => prepared, - Err(error) => return error.into_response(), + + let mm_features = if let Some(parts) = body.content_parts.take() { + match state.chat.prepare_media(parts, &mut body.token_ids).await { + Ok(features) => features, + Err(e) => { + return ApiError::invalid_request( + format!("failed to resolve content_parts: {}", e.as_report()), + Some("content_parts"), + ) + .into_response(); + } + } + } else { + None }; + + let prepared = + match prepare_generate_request(body, &lora_resolution, request_context, mm_features) { + Ok(prepared) => prepared, + Err(error) => return error.into_response(), + }; let request_span = tracing::info_span!( "generate", request_id = %prepared.request_id, diff --git a/rust/src/server/src/routes/inference/generate/convert.rs b/rust/src/server/src/routes/inference/generate/convert.rs index 10361e679ad8..5c141079890d 100644 --- a/rust/src/server/src/routes/inference/generate/convert.rs +++ b/rust/src/server/src/routes/inference/generate/convert.rs @@ -1,6 +1,7 @@ // SPDX-License-Identifier: Apache-2.0 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project +use vllm_engine_core_client::protocol::multimodal::MmFeatures; use vllm_text::{Prompt, TextDecodeOptions, TextRequest}; use super::types::GenerateRequest; @@ -37,6 +38,7 @@ pub(super) fn prepare_generate_request( request: GenerateRequest, lora_resolution: &LoraModelResolution, ctx: ResolvedRequestContext, + mm_features: Option, ) -> Result { validate::validate_request_compat(&request, &lora_resolution.model_names)?; @@ -67,7 +69,7 @@ pub(super) fn prepare_generate_request( let text_request = TextRequest { request_id: ctx.request_id.clone(), prompt: Prompt::TokenIds(request.token_ids), - mm_features: None, + mm_features, sampling_params, decode_options: TextDecodeOptions::default(), intermediate: false, @@ -134,6 +136,7 @@ mod tests { request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]), ResolvedRequestContext::default(), + None, ) .expect("prepare"); @@ -175,6 +178,7 @@ mod tests { request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]), ResolvedRequestContext::default(), + None, ) .expect("prepare"); @@ -204,6 +208,7 @@ mod tests { request, &served(&["Qwen/Qwen1.5-0.5B-Chat"]), ResolvedRequestContext::default(), + None, ) .expect("prepare"); diff --git a/rust/src/server/src/routes/inference/generate/types.rs b/rust/src/server/src/routes/inference/generate/types.rs index 18095d528853..200209e095e1 100644 --- a/rust/src/server/src/routes/inference/generate/types.rs +++ b/rust/src/server/src/routes/inference/generate/types.rs @@ -3,6 +3,7 @@ use std::collections::HashMap; +use llm_multimodal::MediaContentPart; use serde::{Deserialize, Serialize}; use serde_json::{Map, Value}; use validator::Validate; @@ -26,6 +27,8 @@ pub struct GenerateRequest { pub priority: i32, pub kv_transfer_params: Option>, pub ec_transfer_params: Option>, + /// Raw multimodal input; server resolves media. Mutually exclusive with `features`. + pub content_parts: Option>, #[serde(flatten)] pub other: Map, } diff --git a/rust/src/server/src/routes/render.rs b/rust/src/server/src/routes/render.rs index e5a66cec9592..1972c02d0976 100644 --- a/rust/src/server/src/routes/render.rs +++ b/rust/src/server/src/routes/render.rs @@ -93,6 +93,7 @@ fn lower_render_request( priority: text_request.priority, kv_transfer_params: None, ec_transfer_params: None, + content_parts: None, other: Default::default(), }; validate_generate_request(&request, &state.served_model_names)?; diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py index e13dd425af14..b4811666d4d7 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_multimodal_tokens.py @@ -156,3 +156,87 @@ async def test_render_to_generate_roundtrip(client, test_image): assert "red" in text.lower(), ( f"Expected model to identify the red image, got: {text!r}" ) + + +@pytest.mark.asyncio +async def test_content_parts_generates_tokens(client, test_image): + """content_parts with raw media should produce output tokens.""" + data_url = encode_image_url(test_image, format="PNG") + + render_resp = await client.post( + RENDER_ENDPOINT, + json={ + "model": MODEL_NAME, + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_url}}, + {"type": "text", "text": "What color is this?"}, + ], + } + ], + }, + ) + render_resp.raise_for_status() + token_ids = render_resp.json()["token_ids"] + + gen_resp = await client.post( + GEN_ENDPOINT, + json={ + "token_ids": token_ids, + "content_parts": [{"type": "image_url", "url": data_url}], + "sampling_params": {"max_tokens": 10, "temperature": 0.0}, + }, + ) + gen_resp.raise_for_status() + gen_data = gen_resp.json() + + assert "choices" in gen_data + choice = gen_data["choices"][0] + assert "token_ids" in choice + assert len(choice["token_ids"]) > 0 + + +@pytest.mark.asyncio +async def test_content_parts_streaming(client, test_image): + """content_parts should work with streaming.""" + import json as json_mod + + data_url = encode_image_url(test_image, format="PNG") + + render_resp = await client.post( + RENDER_ENDPOINT, + json={ + "model": MODEL_NAME, + "messages": [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": data_url}}, + {"type": "text", "text": "Describe this."}, + ], + } + ], + }, + ) + render_resp.raise_for_status() + token_ids = render_resp.json()["token_ids"] + + async with client.stream( + "POST", + GEN_ENDPOINT, + json={ + "token_ids": token_ids, + "content_parts": [{"type": "image_url", "url": data_url}], + "sampling_params": {"max_tokens": 8, "temperature": 0.0}, + "stream": True, + }, + ) as resp: + resp.raise_for_status() + chunks = [] + async for line in resp.aiter_lines(): + if line.startswith("data: ") and line != "data: [DONE]": + chunks.append(json_mod.loads(line[6:])) + + assert len(chunks) > 0 diff --git a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py index 86c0a8c8b177..606989c31d8c 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/protocol.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/protocol.py @@ -104,6 +104,16 @@ def validate_token_ids(cls, v: list[int]) -> list[int]: features: MultiModalFeatures | None = None """Multimodal hashes and placeholder positions (populated for MM inputs).""" + content_parts: list[dict[str, Any]] | None = None + """Raw multimodal input; server resolves media. Mutually exclusive + with ``features``.""" + + @model_validator(mode="after") + def _check_mm_fields_exclusive(self) -> "GenerateRequest": + if self.content_parts and self.features: + raise ValueError("content_parts and features are mutually exclusive") + return self + sampling_params: SamplingParams """The sampling parameters for the model.""" diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 7b5e2bbb3d92..c5b0ffea69b0 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -14,6 +14,7 @@ from fastapi import Request from vllm.engine.protocol import EngineClient +from vllm.entrypoints.chat_utils import AsyncMultiModalItemTracker from vllm.entrypoints.generate.base.serving import ( GenerateBaseServing, clamp_prompt_logprobs, @@ -33,7 +34,7 @@ from vllm.entrypoints.openai.models.serving import OpenAIServingModels from vllm.entrypoints.serve.utils.api_utils import get_max_tokens, should_include_usage from vllm.entrypoints.serve.utils.request_logger import RequestLogger -from vllm.inputs import EngineInput, mm_input +from vllm.inputs import EngineInput, TokensPrompt, mm_input from vllm.logger import init_logger from vllm.logprobs import Logprob from vllm.multimodal.inputs import ( @@ -143,7 +144,29 @@ async def serve_tokens( return self.create_error_response(e) engine_input: EngineInput - if features := request.features: + if request.content_parts: + tracker = AsyncMultiModalItemTracker(self.model_config) + mm_parser = tracker.create_parser() + for part in request.content_parts: + ptype = part.get("type", "") + url = part.get("url") + uuid = part.get("uuid") + if ptype == "image_url": + mm_parser.parse_image(url, uuid) + elif ptype == "audio_url": + mm_parser.parse_audio(url, uuid) + elif ptype == "video_url": + mm_parser.parse_video(url, uuid) + mm_data, mm_uuids = await tracker.resolve_items() + prompt = TokensPrompt(prompt_token_ids=request.token_ids) + if mm_data: + prompt["multi_modal_data"] = mm_data + if mm_uuids: + prompt["multi_modal_uuids"] = mm_uuids + (engine_input,) = await self.online_renderer.renderer.render_cmpl_async( + [prompt] + ) + elif features := request.features: # Convert PlaceholderRangeInfo → PlaceholderRange per modality. mm_placeholders: dict[str, list[PlaceholderRange]] = { modality: [ From f37dff22fc77e21ac857f2023faa7d5b8f9c9591 Mon Sep 17 00:00:00 2001 From: acmore Date: Tue, 11 Aug 2026 10:05:27 +0800 Subject: [PATCH 23/30] [Bugfix] Import each packed IPC export once on the consumer side (#51259) Signed-off-by: acmore --- .../distributed/weight_transfer/ipc_engine.py | 13 ++++- .../weight_transfer/packed_tensor.py | 50 +++++++++++++++++-- 2 files changed, 59 insertions(+), 4 deletions(-) diff --git a/vllm/distributed/weight_transfer/ipc_engine.py b/vllm/distributed/weight_transfer/ipc_engine.py index 8a7d8744c89d..ed7fb86fa26e 100644 --- a/vllm/distributed/weight_transfer/ipc_engine.py +++ b/vllm/distributed/weight_transfer/ipc_engine.py @@ -27,6 +27,7 @@ from vllm.config import VllmConfig from vllm.distributed.weight_transfer.packed_tensor import ( DEFAULT_PACKED_BUFFER_SIZE_BYTES, + PackedBufferImporter, packed_ipc_consumer, packed_ipc_producer, ) @@ -143,6 +144,9 @@ def __init__( # Set from the trainer-supplied init info at the handshake; defaults are # only for the (unreachable) receive-before-init case. self.packed = False + # Shared across all chunks of every packed transfer this engine + # receives; see PackedBufferImporter for the refcount contract. + self._packed_importer = PackedBufferImporter() def init_transfer_engine(self, init_info: IPCWeightTransferInitInfo) -> None: """ @@ -170,6 +174,12 @@ def finish_weight_update(self) -> None: ) finalize_layerwise_reload(self.model, self.model_config) + # Every reduce_tensor call is a fresh export with its own refcount + # slot, so releasing once per update always balances this update's + # export and lets the trainer reclaim its staging buffer. Callers + # that skip finish are still covered by the replace-on-next-export + # path inside the importer. + self._packed_importer.close() def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: """ @@ -202,6 +212,7 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: dtype_names=update_info.dtype_names, tensor_sizes=update_info.tensor_sizes, device_index=device_index, + importer=self._packed_importer, ) else: assert isinstance(update_info.ipc_handles, list) @@ -236,7 +247,7 @@ def receive_weights(self, update_info: IPCWeightTransferUpdateInfo) -> None: self.model.load_weights(weights) def shutdown(self) -> None: - pass + self._packed_importer.close() @staticmethod def trainer_send_weights(*args: Any, **kwargs: Any) -> None: diff --git a/vllm/distributed/weight_transfer/packed_tensor.py b/vllm/distributed/weight_transfer/packed_tensor.py index 59bdbafd40cd..79641a8395db 100644 --- a/vllm/distributed/weight_transfer/packed_tensor.py +++ b/vllm/distributed/weight_transfer/packed_tensor.py @@ -381,6 +381,45 @@ def packed_ipc_producer( ) +class PackedBufferImporter: + """One-slot cache for the consumer-side mapping of a packed IPC buffer. + + torch's cross-process refcount pairs one ``reduce_tensor`` export with + one consumer rebuild-release cycle. The producer exports its buffer once + and ships the same args with every chunk, so rebuilding and releasing + per chunk decrements the counter once per chunk, drives it negative, + and the dropped buffer is then never reclaimable by ``ipc_collect`` on + the producer side (it leaks one buffer per transfer). Caching the + rebuilt buffer keyed by the rebuild args restores the pairing: chunks + of one transfer reuse the mapping (it aliases producer memory, so it + always reads the current chunk's bytes), and replacing the entry when + the next export arrives — or ``close()`` — releases the previous + mapping exactly once. + + The importer that consumes a multi-chunk transfer must be one object + reused across chunks; a fresh importer per chunk reintroduces the + per-chunk release. ``IPCWeightTransferEngine`` owns one instance per + engine and closes it on shutdown. + """ + + def __init__(self) -> None: + self._entry: tuple[tuple, torch.Tensor] | None = None + + def rebuild(self, list_args: list) -> torch.Tensor: + from torch.multiprocessing.reductions import rebuild_cuda_tensor + + key = tuple(list_args) + if self._entry is not None and self._entry[0] == key: + return self._entry[1] + packed = rebuild_cuda_tensor(*list_args) + self._entry = (key, packed) + return packed + + def close(self) -> None: + """Release the held mapping (one producer-side refcount decrement).""" + self._entry = None + + def packed_ipc_consumer( ipc_handle: dict[str, tuple], names: list[str], @@ -388,6 +427,7 @@ def packed_ipc_consumer( dtype_names: list[str], tensor_sizes: list[int], device_index: int, + importer: "PackedBufferImporter | None" = None, ) -> list[tuple[str, torch.Tensor]]: """Unpack a single packed IPC chunk into named tensors. @@ -410,9 +450,11 @@ def packed_ipc_consumer( dtype_names: Parameter dtype name strings (e.g. "float16") tensor_sizes: Size in bytes of each parameter in the packed buffer device_index: Local CUDA device index + importer: Import cache that must be shared across all chunks of one + export so the buffer mapping is opened and released exactly once + (see ``PackedBufferImporter``). ``None`` builds an ephemeral one, + which is only safe for single-chunk consumption. """ - from torch.multiprocessing.reductions import rebuild_cuda_tensor - props = torch.cuda.get_device_properties(device_index) physical_gpu_id = str(props.uuid) @@ -425,7 +467,9 @@ def packed_ipc_consumer( args = ipc_handle[physical_gpu_id] list_args = list(args) list_args[6] = device_index - packed = rebuild_cuda_tensor(*list_args) + if importer is None: + importer = PackedBufferImporter() + packed = importer.rebuild(list_args) content_size = sum(tensor_sizes) packed = packed[:content_size] From 99e62b802c2a2fcf8393e0ca87d8b7d8f00297d3 Mon Sep 17 00:00:00 2001 From: aoshen02 Date: Tue, 11 Aug 2026 10:16:57 +0800 Subject: [PATCH 24/30] [Bugfix][Model Loader] Defer post-load attention weight processing (#49519) Signed-off-by: aoshen02 Co-authored-by: Roger Wang --- .../model_loader/test_reload.py | 77 +++++++++++++++++++ .../layers/attention/__init__.py | 12 +++ .../model_loader/reload/layerwise.py | 31 ++++---- vllm/model_executor/model_loader/utils.py | 9 +-- 4 files changed, 105 insertions(+), 24 deletions(-) diff --git a/tests/model_executor/model_loader/test_reload.py b/tests/model_executor/model_loader/test_reload.py index c21d9b1e02f8..4963f3ae5b42 100644 --- a/tests/model_executor/model_loader/test_reload.py +++ b/tests/model_executor/model_loader/test_reload.py @@ -9,6 +9,9 @@ from torch.nn.parameter import UninitializedParameter import vllm.model_executor.model_loader.reload.meta as reload_meta +from vllm.config import ModelConfig +from vllm.model_executor.layers.attention import MMEncoderAttention +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase from vllm.model_executor.layers.linear import QKVParallelLinear from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.reload.layerwise import ( @@ -101,6 +104,39 @@ def __init__(self): self.register_buffer("scale", torch.tensor(0.25), persistent=False) +class _ReloadableMMEncoderAttention(MMEncoderAttention): + """Minimal stand-in to test reload lifecycle without encoder initialization.""" + + def __init__(self): + torch.nn.Module.__init__(self) + self.weight = torch.nn.Parameter(torch.ones(2, 2)) + self.weight.weight_loader = default_weight_loader + self.post_load_called = False + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + self.post_load_called = True + + +class _ReloadableAttentionLayer( + torch.nn.Module, + AttentionLayerBase, +): + def __init__(self): + super().__init__() + self.weight = torch.nn.Parameter(torch.ones(2, 2)) + self.weight.weight_loader = default_weight_loader + self.post_load_called = False + + def get_attn_backend(self): + raise NotImplementedError + + def get_kv_cache_spec(self, vllm_config): + return None + + def process_weights_after_loading(self, act_dtype: torch.dtype) -> None: + self.post_load_called = True + + def test_move_metatensors(): tensor = torch.empty((1, 2, 3)) meta_tensor = to_meta_tensor(tensor) @@ -115,6 +151,47 @@ def test_move_metatensors(): assert tensor.__dict__ == meta_tensor.__dict__ == materialized_tensor.__dict__ +@pytest.mark.parametrize( + "layer_cls", + [_ReloadableMMEncoderAttention, _ReloadableAttentionLayer], +) +def test_attention_reload_defers_post_load(default_vllm_config, layer_cls): + default_vllm_config.model_config = ModelConfig() + layer = layer_cls() + model = torch.nn.Sequential(layer) + loaded_weight = torch.full_like(layer.weight, 7.0) + + record_metadata_for_reloading(model) + initialize_layerwise_reload(model) + layer.weight.weight_loader(layer.weight, loaded_weight) + + assert not layer.post_load_called + + finalize_layerwise_reload(model, default_vllm_config.model_config) + + assert layer.post_load_called + assert torch.equal(layer.weight, loaded_weight) + + +@pytest.mark.parametrize( + "layer_cls", + [_ReloadableMMEncoderAttention, _ReloadableAttentionLayer], +) +def test_attention_first_load_processes_weights(default_vllm_config, layer_cls): + default_vllm_config.model_config = ModelConfig() + layer = layer_cls() + model = torch.nn.Sequential(layer) + loaded_weight = torch.full_like(layer.weight, 7.0) + + initialize_online_processing(layer) + layer.weight.weight_loader(layer.weight, loaded_weight) + + finalize_layerwise_reload(model, default_vllm_config.model_config) + + assert layer.post_load_called + assert torch.equal(layer.weight, loaded_weight) + + def test_reload_lifecycle(): layer = torch.nn.Linear(2, 3) info = LayerReloadingInfo( diff --git a/vllm/model_executor/layers/attention/__init__.py b/vllm/model_executor/layers/attention/__init__.py index c9e477fb114f..4b04c841032c 100644 --- a/vllm/model_executor/layers/attention/__init__.py +++ b/vllm/model_executor/layers/attention/__init__.py @@ -1,6 +1,8 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import torch + from vllm.model_executor.layers.attention.attention import Attention from vllm.model_executor.layers.attention.chunked_local_attention import ( ChunkedLocalAttention, @@ -18,6 +20,15 @@ from vllm.model_executor.layers.attention.static_sink_attention import ( StaticSinkAttention, ) +from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase + + +def is_deferred_attention_layer(layer: torch.nn.Module) -> bool: + """Whether an attention-like layer requires deferred post-load processing.""" + return isinstance(layer, (AttentionLayerBase, MMEncoderAttention)) and callable( + getattr(layer, "process_weights_after_loading", None) + ) + __all__ = [ "Attention", @@ -29,4 +40,5 @@ "PrefillPrefixLMAttention", "RSWAAttention", "StaticSinkAttention", + "is_deferred_attention_layer", ] diff --git a/vllm/model_executor/model_loader/reload/layerwise.py b/vllm/model_executor/model_loader/reload/layerwise.py index d6d4eb4b9308..5aa510a6242a 100644 --- a/vllm/model_executor/model_loader/reload/layerwise.py +++ b/vllm/model_executor/model_loader/reload/layerwise.py @@ -9,7 +9,7 @@ from vllm.config import ModelConfig from vllm.logger import init_logger -from vllm.model_executor.layers.attention import Attention, MLAAttention +from vllm.model_executor.layers.attention import is_deferred_attention_layer from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase from vllm.model_executor.model_loader.weight_utils import default_weight_loader @@ -196,7 +196,7 @@ def online_process_loader(*args, **kwargs): ) # Do not online process attention layers, must wait until finalize - if isinstance(layer, (Attention, MLAAttention)): + if is_deferred_attention_layer(layer): return ret # Log warnings allocating excessive buffers on device @@ -249,8 +249,8 @@ def finalize_layerwise_processing(model: torch.nn.Module, model_config: ModelCon info.reset() continue - # Attention/MLA layers are processed after all other layers - if isinstance(layer, (Attention, MLAAttention)): + # Deferred attention-like layers are processed after all other layers + if is_deferred_attention_layer(layer): deferred_attn.append((layer, info)) continue @@ -293,15 +293,13 @@ def finalize_layerwise_reload(*args, **kwargs): def _finalize_attention_layer( layer: torch.nn.Module, info: LayerReloadingInfo, model_config: ModelConfig ) -> None: - if info.load_numel > 0 and info.kernel_tensors is not None: + if info.kernel_tensors is None: + if info.load_numel > 0: + _layerwise_process(layer, info) + elif info.load_numel > 0: # Reload with new scale weights from checkpoint _place_kernel_tensors(layer, info) _reload_attention_scales(layer, info) - elif info.load_numel > 0 or info.kernel_tensors is None: - raise ValueError( - "Layerwise loading of attention layers is not supported. " - "Attention must always process after linears." - ) else: _place_kernel_tensors(layer, info) layer.process_weights_after_loading(model_config.dtype) @@ -315,19 +313,18 @@ def _reload_attention_scales(layer: torch.nn.Module, info: LayerReloadingInfo) - processing, since we use .data.copy_() to preserve kernel tensor references.""" quant_method = getattr(layer, "quant_method", None) - if quant_method is None: - return - - # Re-create scale Parameters with sentinel values so unloaded scales - # are correctly detected by process_weights_after_loading - quant_method.create_weights(layer) + if quant_method is not None: + # Re-create scale Parameters with sentinel values so unloaded scales + # are correctly detected by process_weights_after_loading + quant_method.create_weights(layer) for name, args in info.loaded_weights: param = getattr(layer, name) args.arguments["param"] = param _get_weight_loader(param)(*args.args, **args.kwargs) - quant_method.process_weights_after_loading(layer) + if quant_method is not None: + quant_method.process_weights_after_loading(layer) _copy_and_restore_kernel_tensors(layer, info) diff --git a/vllm/model_executor/model_loader/utils.py b/vllm/model_executor/model_loader/utils.py index a472c7543583..b7ccb78c04ee 100644 --- a/vllm/model_executor/model_loader/utils.py +++ b/vllm/model_executor/model_loader/utils.py @@ -14,10 +14,7 @@ import vllm.envs as envs from vllm.config import ModelConfig, VllmConfig, set_current_vllm_config from vllm.logger import init_logger -from vllm.model_executor.layers.attention import ( - MMEncoderAttention, -) -from vllm.model_executor.layers.attention_layer_base import AttentionLayerBase +from vllm.model_executor.layers.attention import is_deferred_attention_layer from vllm.model_executor.layers.hpc import HpcModule from vllm.model_executor.layers.quantization.base_config import ( QuantizationConfig, @@ -124,9 +121,7 @@ def process_weights_after_loading( # encoder. NOTE: Happens after other modules so we can easily decompress # weights. for _, module in model.named_modules(): - if isinstance(module, (AttentionLayerBase, MMEncoderAttention)) and hasattr( - module, "process_weights_after_loading" - ): + if is_deferred_attention_layer(module): # TODO(lucas): see if there is a way to unify the signatures # of process_weights_after_loading with device_loading_context(module, target_device): From 419b51b38517bd446c6e359d37b1ee8674cd7e84 Mon Sep 17 00:00:00 2001 From: Andreas Karatzas Date: Mon, 10 Aug 2026 21:51:39 -0500 Subject: [PATCH 25/30] [Bugfix][ROCm][CI] Stabilize build context and source caches (#51721) Signed-off-by: Andreas Karatzas Signed-off-by: Andreas Karatzas --- .buildkite/scripts/ci-bake-rocm.sh | 370 ++++++++++++++++++++++++----- docker/Dockerfile.rocm | 36 ++- 2 files changed, 345 insertions(+), 61 deletions(-) diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index b0f32cd56fc8..8ba73b536a8a 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -22,7 +22,7 @@ DEFAULT_CI_BASE_METADATA_VERSION="3" # ROCm CI forces REMOTE_VLLM=0, so content identity covers only the selected # local-source stages rather than unreachable remote-fetch alternatives. DEFAULT_ROCM_CSRC_CONTENT_FILES=".dockerignore requirements/common.txt requirements/rocm.txt pyproject.toml setup.py CMakeLists.txt cmake csrc vllm/envs.py vllm/__init__.py tools/build_rust.py" -DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies csrc-build" +DEFAULT_ROCM_CSRC_DOCKERFILE_STAGES="base fetch_vllm_0 fetch_vllm build_vllm_dependencies rocm-triton-kernels csrc-build" DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore requirements/build/rust.txt rust/Cargo.lock rust/Cargo.toml rust/proto rust/src rust-toolchain.toml tools/build_rust.py tools/install_protoc.sh build_rust.sh" DEFAULT_ROCM_RUST_DOCKERFILE_STAGES="base rust_toolchain_input_0 rust-toolchain-input rust_input_0 rust-input rust-toolchain rust-build" # Docker's 128-character tag limit minus the longest cache prefix @@ -36,9 +36,14 @@ CI_HCL_PATH="" CI_BASE_LABEL_OVERRIDE_PATH="" CSRC_CACHE_OVERRIDE_PATH="" ROCM_ARG_OVERRIDE_PATH="" +BUILD_CONTEXT_OVERRIDE_PATH="" SCRIPT_TMP_DIR="" BAKE_CONFIG_FILE="" +ROCM_BUILD_CONTEXT_ROOT="" +ROCM_BUILD_CONTEXT_INDEX="" +ROCM_BUILD_CONTEXT_COMMIT="" BAKE_FILES=() +BAKE_ALLOW_ARGS=() BAKE_TARGETS=() DEPENDENCY_CACHE_TARGETS=() @@ -192,16 +197,20 @@ get_buildkite_target_repo_url() { printf 'https://github.com/%s.git\n' "${DEFAULT_REPO_SLUG}" } -git_fetch_for_cache() { +git_fetch_with_timeout() { local timeout_secs="${ROCM_CACHE_GIT_FETCH_TIMEOUT:-60}" if command -v timeout >/dev/null 2>&1; then - timeout "${timeout_secs}s" git fetch "$@" 2>/dev/null + timeout "${timeout_secs}s" git fetch "$@" else - git fetch "$@" 2>/dev/null + git fetch "$@" fi } +git_fetch_for_cache() { + git_fetch_with_timeout "$@" 2>/dev/null +} + hash_string_short() { printf '%s' "$1" | sha256sum | cut -c1-16 } @@ -209,23 +218,51 @@ hash_string_short() { list_content_files() { # Hash the checkout inputs Docker can intentionally consume, not ignored # compiler/test debris left behind on a reused worker. - git ls-files -z --cached --others --exclude-standard -- "$1" \ - | LC_ALL=C sort -z + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" ]]; then + GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \ + git ls-files -z --cached -- "$1" | LC_ALL=C sort -z + else + git ls-files -z --cached --others --exclude-standard -- "$1" \ + | LC_ALL=C sort -z + fi } -hash_content_file() { +content_regular_file() { local file="$1" + local physical_file="" - if [[ -L "${file}" ]]; then - printf 'symlink:%s\ntarget:%s\n' "${file}" "$(readlink "${file}")" + physical_file="${file}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${file}" != /* ]]; then + physical_file="${ROCM_BUILD_CONTEXT_ROOT}/${file}" + fi + [[ -f "${physical_file}" ]] +} + +hash_content_file() { + local file="$1" + local physical_file="" + local checksum="" + local file_mode="" + + physical_file="${file}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${file}" != /* ]]; then + physical_file="${ROCM_BUILD_CONTEXT_ROOT}/${file}" + fi + if [[ -L "${physical_file}" ]]; then + printf 'symlink:%s\ntarget:' "${file}" + readlink -n -- "${physical_file}" || return $? + printf '\n' return fi - if [[ ! -f "${file}" ]]; then + if [[ ! -f "${physical_file}" ]]; then printf 'missing:%s\n' "${file}" return fi - printf 'file:%s\nmode:%s\n' "${file}" "$(stat -c '%a' "${file}")" - sha256sum "${file}" + file_mode=$(stat -c '%a' "${physical_file}") || return $? + printf 'file:%s\nmode:%s\n' "${file}" "${file_mode}" + checksum=$(sha256sum < "${physical_file}") || return $? + checksum="${checksum%% *}" + printf '%s %s\n' "${checksum}" "${file}" } hash_content_directory() { @@ -239,11 +276,16 @@ hash_content_directory() { compute_content_hash() { local path="" + local physical_path="" for path in "$@"; do - if [[ -L "${path}" || -f "${path}" ]]; then + physical_path="${path}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${path}" != /* ]]; then + physical_path="${ROCM_BUILD_CONTEXT_ROOT}/${path}" + fi + if [[ -L "${physical_path}" || -f "${physical_path}" ]]; then hash_content_file "${path}" || return $? - elif [[ -d "${path}" ]]; then + elif [[ -d "${physical_path}" ]]; then hash_content_directory "${path}" || return $? else printf 'missing:%s\n' "${path}" @@ -251,25 +293,69 @@ compute_content_hash() { done | sha256sum | cut -d' ' -f1 } -normalize_ci_worktree_modes() { +validate_ci_build_context_source() { + local source_root="$1" + local context_commit="" + local source_head="" local entry="" local metadata="" local mode="" local stage="" local path="" local index_modes_file="${SCRIPT_TMP_DIR}/git-index-modes" - local -a regular_files=() - local -a executable_files=() + local info_attributes="" - [[ "${BUILDKITE:-false}" == "true" ]] || return 0 - [[ "${REMOTE_VLLM:-0}" == "0" ]] || return 0 - if ! git ls-files --stage -z > "${index_modes_file}"; then - echo "Failed to read Git modes for the CI Docker context" >&2 + if ! context_commit=$(git -C "${source_root}" rev-parse \ + --verify "${BUILDKITE_COMMIT:-HEAD}^{commit}"); then + echo "Failed to resolve the CI Docker context revision" >&2 + return 1 + fi + if ! source_head=$(git -C "${source_root}" rev-parse --verify HEAD); then + echo "Failed to resolve the checked-out CI revision" >&2 + return 1 + fi + if [[ "${context_commit}" != "${source_head}" ]]; then + echo "BUILDKITE_COMMIT does not match the checked-out CI revision" >&2 + return 1 + fi + ROCM_BUILD_CONTEXT_COMMIT="${context_commit}" + if ! git -C "${source_root}" diff \ + --quiet --no-ext-diff --ignore-submodules=none --; then + echo "Tracked worktree changes cannot be omitted from the CI Docker context" \ + >&2 + return 1 + fi + if ! git -C "${source_root}" diff --cached \ + --quiet --no-ext-diff --ignore-submodules=none HEAD --; then + echo "Staged changes cannot be omitted from the CI Docker context" >&2 + return 1 + fi + # Untracked and ignored worker outputs are intentionally absent: the + # pinned Git tree, rather than mutable checkout contents, is the contract. + if ! info_attributes=$(git -C "${source_root}" \ + rev-parse --path-format=absolute --git-path info/attributes); then + echo "Failed to locate repository-local Git attributes" >&2 + return 1 + fi + if [[ -s "${info_attributes}" ]]; then + echo "Repository-local Git attributes cannot define the CI Docker context" \ + >&2 return 1 fi - # Some shared workspaces expose 0644/0755 files as 0664/0775. Restore the - # index modes so both the identity hash and Docker context are canonical. + ROCM_BUILD_CONTEXT_INDEX="${SCRIPT_TMP_DIR}/docker-context.index" + if ! GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \ + git -C "${source_root}" -c core.splitIndex=false \ + read-tree "${context_commit}^{tree}"; then + echo "Failed to create the CI Docker context index" >&2 + return 1 + fi + if ! GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" \ + git -C "${source_root}" ls-files --stage -z \ + > "${index_modes_file}"; then + echo "Failed to read the Git index for the CI Docker context" >&2 + return 1 + fi while IFS= read -r -d '' entry; do metadata="${entry%%$'\t'*}" mode="${metadata%% *}" @@ -280,13 +366,12 @@ normalize_ci_worktree_modes() { return 1 fi case "${mode}" in - 100644) - regular_files+=("${path}") + 100644|100755|120000) ;; - 100755) - executable_files+=("${path}") - ;; - 120000|160000) + 160000) + echo "Git submodule cannot be materialized in the CI Docker context: ${path}" \ + >&2 + return 1 ;; *) echo "Unsupported Git mode ${mode} for ${path}" >&2 @@ -294,10 +379,130 @@ normalize_ci_worktree_modes() { ;; esac done < "${index_modes_file}" +} - ((${#regular_files[@]} == 0)) || chmod 0644 -- "${regular_files[@]}" - ((${#executable_files[@]} == 0)) || chmod 0755 -- "${executable_files[@]}" - echo "Normalized Git file modes for the CI Docker context" +describe_ci_revision() { + git -C "$1" describe --tags --long --abbrev=10 \ + --match '*[0-9]*' "$2" 2>/dev/null +} + +write_ci_git_archival_metadata() { + local source_root="$1" + local context_root="$2" + local commit="${ROCM_BUILD_CONTEXT_COMMIT}" + local commit_date="" + local describe="" + local is_shallow="false" + + commit_date=$(git -C "${source_root}" show -s --format=%cI "${commit}") \ + || return $? + is_shallow=$(git -C "${source_root}" rev-parse --is-shallow-repository) \ + || return $? + if git -C "${source_root}" remote get-url origin >/dev/null 2>&1; then + # AMD agents use shallow, no-tag clones. Reach a release tag so one + # commit cannot acquire different versions from different depths. + # Versioning only needs tags and commits, not historical source trees. + echo "Synchronizing version tags for the canonical CI Docker context" + (cd "${source_root}" && git_fetch_with_timeout --quiet \ + --filter=tree:0 --prune origin '+refs/tags/*:refs/tags/*') \ + || return $? + if [[ "${is_shallow}" == "true" ]] \ + && ! describe_ci_revision "${source_root}" "${commit}" >/dev/null; then + echo "Deepening history to reach a version tag" + (cd "${source_root}" && git_fetch_with_timeout --quiet \ + --filter=tree:0 --no-tags --deepen=1000 origin "${commit}") \ + || return $? + is_shallow=$(git -C "${source_root}" \ + rev-parse --is-shallow-repository) || return $? + if [[ "${is_shallow}" == "true" ]] \ + && ! describe_ci_revision "${source_root}" "${commit}" >/dev/null; then + echo "No version tag within 1,000 commits; fetching full history" + (cd "${source_root}" && git_fetch_with_timeout --quiet \ + --filter=tree:0 --no-tags --unshallow origin "${commit}") \ + || return $? + fi + fi + fi + if ! describe=$(describe_ci_revision "${source_root}" "${commit}"); then + echo "No numeric version tag is reachable from the CI revision" >&2 + return 1 + fi + if [[ -e "${context_root}/.git_archival.txt" \ + || -L "${context_root}/.git_archival.txt" ]]; then + echo "The source tree already contains .git_archival.txt" >&2 + return 1 + fi + printf 'node: %s\nnode-date: %s\ndescribe-name: %s\n' \ + "${commit}" "${commit_date}" "${describe}" \ + > "${context_root}/.git_archival.txt" \ + && chmod 0644 -- "${context_root}/.git_archival.txt" +} + +write_build_context_override() { + local escaped_context="" + + [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" ]] || return 0 + escaped_context=$(hcl_escape_string "${ROCM_BUILD_CONTEXT_ROOT}") || return $? + if ! { + printf 'target "_common-rocm" {\n' + printf ' context = "%s"\n' "${escaped_context}" + printf '}\n' + } > "${BUILD_CONTEXT_OVERRIDE_PATH}"; then + echo "Failed to write the CI Docker context override" >&2 + return 1 + fi + BAKE_FILES+=(-f "${BUILD_CONTEXT_OVERRIDE_PATH}") +} + +prepare_ci_build_context() { + local source_root="" + local context_root="" + + [[ "${BUILDKITE:-false}" == "true" ]] || return 0 + [[ "${REMOTE_VLLM:-0}" == "0" ]] || return 0 + + # BuildKit includes file modes in cache keys. Export the pinned revision to + # an owned context instead of changing modes in the shared checkout. + if ! source_root=$(git rev-parse --show-toplevel); then + echo "Failed to locate the CI source checkout" >&2 + return 1 + fi + validate_ci_build_context_source "${source_root}" || return $? + context_root="${SCRIPT_TMP_DIR}/docker-context" + if ! mkdir -m 0700 -- "${context_root}"; then + echo "Failed to create the owned CI Docker context" >&2 + return 1 + fi + if ! ( + umask 0022 + unset GIT_LFS_SKIP_SMUDGE + export GIT_ATTR_NOSYSTEM=1 + export GIT_INDEX_FILE="${ROCM_BUILD_CONTEXT_INDEX}" + git -C "${source_root}" \ + -c core.attributesFile=/dev/null \ + -c core.autocrlf=false \ + -c core.eol=lf \ + -c core.symlinks=true \ + --work-tree="${context_root}" \ + checkout-index --all + ); then + echo "Failed to materialize the CI Docker context from the Git index" >&2 + return 1 + fi + + # setuptools-scm understands Git's stable archive format, so full-source + # wheels retain their exact version without copying mutable Git history. + if ! is_ci_base_target; then + write_ci_git_archival_metadata "${source_root}" "${context_root}" \ + || return $? + fi + if [[ -e "${context_root}/.git" ]]; then + echo "Canonical CI Docker context unexpectedly contains .git" >&2 + return 1 + fi + ROCM_BUILD_CONTEXT_ROOT="${context_root}" + BAKE_ALLOW_ARGS+=(--allow "fs.read=${ROCM_BUILD_CONTEXT_ROOT}") + echo "Using canonical CI Docker context: ${ROCM_BUILD_CONTEXT_ROOT}" } compose_dependency_cache_key() { @@ -312,6 +517,12 @@ compose_dependency_cache_key() { hash_dockerfile_stages() { local dockerfile="$1" local stages="$2" + local physical_dockerfile="" + + physical_dockerfile="${dockerfile}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then + physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" + fi awk -v wanted_stages="${stages}" ' BEGIN { @@ -335,14 +546,19 @@ hash_dockerfile_stages() { emit { print } - ' "${dockerfile}" + ' "${physical_dockerfile}" } discover_dockerfile_stage_args() { local dockerfile="$1" local stages="$2" + local physical_dockerfile="" - [[ -f "${dockerfile}" ]] || return 0 + physical_dockerfile="${dockerfile}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then + physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" + fi + [[ -f "${physical_dockerfile}" ]] || return 0 awk -v wanted_stages="${stages}" ' function add_arg(name) { @@ -406,7 +622,7 @@ discover_dockerfile_stage_args() { } } } - ' "${dockerfile}" + ' "${physical_dockerfile}" } get_content_arg_names() { @@ -446,7 +662,7 @@ compute_ci_base_content_hash() { || return 1 if [[ -n "${stages}" ]]; then printf 'dockerfile-stages:%s\n' "${stages}" - if [[ -f "${dockerfile}" ]]; then + if content_regular_file "${dockerfile}"; then hash_dockerfile_stages "${dockerfile}" "${stages}" else printf 'missing:%s\n' "${dockerfile}" @@ -459,9 +675,14 @@ compute_ci_base_content_hash() { extract_dockerfile_arg_default() { local dockerfile="$1" local arg_name="$2" + local physical_dockerfile="" + physical_dockerfile="${dockerfile}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile}" != /* ]]; then + physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile}" + fi sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${arg_name}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ - "${dockerfile}" | head -1 + "${physical_dockerfile}" | head -1 } resolve_image_digest() { @@ -523,7 +744,7 @@ resolve_dockerfile_arg_value() { if [[ -z "${value}" && "${env_name}" != "${arg_name}" ]]; then value="${!arg_name:-}" fi - if [[ -z "${value}" && -f "${dockerfile}" ]]; then + if [[ -z "${value}" ]] && content_regular_file "${dockerfile}"; then value=$(extract_dockerfile_arg_default "${dockerfile}" "${arg_name}") fi @@ -741,6 +962,7 @@ init_config() { CI_BASE_LABEL_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/ci-base-label-override.hcl" CSRC_CACHE_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-csrc-cache-override.hcl" ROCM_ARG_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/rocm-arg-override.hcl" + BUILD_CONTEXT_OVERRIDE_PATH="${SCRIPT_TMP_DIR}/build-context-override.hcl" BAKE_CONFIG_FILE="bake-config-build-${BUILDKITE_BUILD_NUMBER:-local}.json" } @@ -1270,7 +1492,8 @@ prepare_git_cache_metadata() { if [[ -z "${PARENT_COMMIT:-}" || -z "${VLLM_MERGE_BASE_COMMIT:-}" ]] \ && git rev-parse --is-shallow-repository 2>/dev/null | grep -q "true"; then echo "Shallow clone detected - deepening for cache key computation" - git_fetch_for_cache --deepen=1 origin || true + git_fetch_for_cache --filter=tree:0 --no-tags --deepen=1 \ + origin "$(git rev-parse HEAD)" || true fi if [[ -z "${PARENT_COMMIT:-}" ]]; then @@ -1524,7 +1747,7 @@ compute_rocm_csrc_content_hash() { printf 'resolved-build-args:\n' hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}" printf 'dockerfile-stages:%s\n' "${stages}" - if [[ -f "${dockerfile_rocm}" ]]; then + if content_regular_file "${dockerfile_rocm}"; then hash_dockerfile_stages "${dockerfile_rocm}" "${stages}" else printf 'missing:%s\n' "${dockerfile_rocm}" @@ -1574,7 +1797,7 @@ compute_rocm_rust_content_hash() { printf 'resolved-build-args:\n' hash_dockerfile_arg_values "${dockerfile_rocm}" "${content_args[@]}" printf 'dockerfile-stages:%s\n' "${stages}" - if [[ -f "${dockerfile_rocm}" ]]; then + if content_regular_file "${dockerfile_rocm}"; then hash_dockerfile_stages "${dockerfile_rocm}" "${stages}" else printf 'missing:%s\n' "${dockerfile_rocm}" @@ -1736,24 +1959,24 @@ should_export_content_cache_ref() { case "${mode}" in always) - echo "${cache_name} content cache export mode is always; exporting ${cache_ref}" + echo "${cache_name} content cache export mode is always" return 0 ;; never) - echo "${cache_name} content cache export mode is never; not exporting ${cache_ref}" + echo "${cache_name} content cache export mode is never" return 1 ;; missing|"") - if remote_image_exists "${trusted_ref}"; then - echo "${cache_name} trusted content cache exists; not exporting ${cache_ref}" + if registry_ref_exists_with_retry imagetools "${trusted_ref}"; then + echo "${cache_name} trusted content cache is visible: ${trusted_ref}" return 1 fi if [[ "${cache_ref}" != "${trusted_ref}" ]] \ - && remote_image_exists "${cache_ref}"; then - echo "${cache_name} scoped content cache exists; not re-exporting ${cache_ref}" + && registry_ref_exists_with_retry imagetools "${cache_ref}"; then + echo "${cache_name} scoped content cache is visible: ${cache_ref}" return 1 fi - echo "${cache_name} content cache missing; will export ${cache_ref}" + echo "${cache_name} content cache is missing: ${cache_ref}" return 0 ;; *) @@ -1791,7 +2014,7 @@ write_rocm_cache_override() { validate_cache_export_mode "${rocm_cache_to_mode}" "ROCM_FINAL_CACHE_TO_MODE" echo "ROCm content cache export mode: ${content_cache_export_mode}" echo "ROCm csrc cache export mode: ${csrc_cache_to_mode}" - echo "ROCm Rust cache export mode: ${rust_cache_to_mode}" + echo "ROCm Rust fallback cache export mode: ${rust_cache_to_mode}" echo "ROCm final image cache export mode: ${rocm_cache_to_mode}" if [[ -n "${ROCM_CSRC_CONTENT_CACHE_REF:-}" ]]; then @@ -1825,14 +2048,20 @@ write_rocm_cache_override() { "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF}" ) fi - if should_export_content_cache_ref \ + # Legacy commit/branch exports are only needed while the exact-input + # content ref is absent. The exact ref itself is refreshed below. + if ! should_export_content_cache_ref \ "${ROCM_RUST_CONTENT_CACHE_REF}" "ROCm Rust" \ "${ROCM_RUST_TRUSTED_CONTENT_CACHE_REF}"; then + export_rust_cache=0 + fi + if [[ "${content_cache_export_mode}" != "never" ]]; then + # Refresh the exact-input ref in the original solve regardless of + # which local or remote cache supplied the Rust result. rust_cache_to+=( - "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=${rust_cache_to_mode},ignore-error=true" + "type=registry,ref=${ROCM_RUST_CONTENT_CACHE_REF},mode=min,ignore-error=true" ) - else - export_rust_cache=0 + echo "ROCm Rust exact-input cache will be refreshed (mode=min): ${ROCM_RUST_CONTENT_CACHE_REF}" fi fi @@ -1872,6 +2101,16 @@ write_rocm_cache_override() { ) fi + # Standalone image/wheel targets reach rust-build but not the cache-only + # exporter unless it is requested explicitly. + if ((${#rust_cache_to[@]} > 0)); then + case "${TARGET}" in + test-rocm-ci|export-wheel-rocm) + BAKE_TARGETS=("rust-rocm-ci" "${BAKE_TARGETS[@]}") + ;; + esac + fi + if [[ "${TARGET}" == "test-rocm-ci-with-wheel" ]]; then export_wheel_cache_to=() else @@ -1937,12 +2176,17 @@ EOF extract_dependency_pins() { local bake_dir="" local dockerfile_rocm="" + local physical_dockerfile="" local var="" local val="" bake_dir=$(dirname "${VLLM_BAKE_FILE}") dockerfile_rocm="${bake_dir}/Dockerfile.rocm" - if [[ ! -f "${dockerfile_rocm}" ]]; then + physical_dockerfile="${dockerfile_rocm}" + if [[ -n "${ROCM_BUILD_CONTEXT_ROOT:-}" && "${dockerfile_rocm}" != /* ]]; then + physical_dockerfile="${ROCM_BUILD_CONTEXT_ROOT}/${dockerfile_rocm}" + fi + if [[ ! -f "${physical_dockerfile}" ]]; then return 0 fi @@ -1954,7 +2198,7 @@ extract_dependency_pins() { val=$( sed -n -E "s/^[[:space:]]*ARG[[:space:]]+${var}=\"?([^\"[:space:]]+)\"?.*/\\1/p" \ - "${dockerfile_rocm}" | head -1 + "${physical_dockerfile}" | head -1 ) if [[ -n "${val}" ]]; then export "${var}=${val}" @@ -2162,7 +2406,8 @@ print_bake_config() { echo "--- :page_facing_up: Resolved bake configuration" mapfile -t print_targets < <(bake_config_targets) - docker buildx bake "${BAKE_FILES[@]}" --print "${print_targets[@]}" | tee "${BAKE_CONFIG_FILE}" + docker buildx bake "${BAKE_ALLOW_ARGS[@]}" \ + "${BAKE_FILES[@]}" --print "${print_targets[@]}" | tee "${BAKE_CONFIG_FILE}" if command -v buildkite-agent >/dev/null 2>&1 && [[ -n "${BUILDKITE_BUILD_NUMBER:-}" ]]; then buildkite-agent artifact upload "${BAKE_CONFIG_FILE}" || true @@ -2234,6 +2479,7 @@ seed_dependency_caches_if_needed() { echo "--- :docker: Seeding ${target}" echo "Expected cache ref: ${cache_ref}" docker buildx bake \ + "${BAKE_ALLOW_ARGS[@]}" \ "${BAKE_FILES[@]}" \ --progress "${BUILDKIT_PROGRESS:-plain}" \ "${target}" @@ -2250,6 +2496,7 @@ run_bake() { echo "--- :docker: Building ${TARGET}" docker buildx bake \ + "${BAKE_ALLOW_ARGS[@]}" \ "${BAKE_FILES[@]}" \ --progress "${BUILDKIT_PROGRESS:-plain}" \ "${BAKE_TARGETS[@]}" @@ -2358,14 +2605,21 @@ main() { configure_ci_base_write_scope print_header validate_inputs - normalize_ci_worktree_modes load_ci_hcl init_bake_files + if is_ci_base_target; then + prepare_ci_build_context + fi compute_ci_base_hash_if_needed configure_ci_base_image_refs maybe_skip_existing_image setup_builder prepare_git_cache_metadata + # Non-ci_base builds may deepen a shallow checkout above. Derive archival + # version metadata only after that lookup sees the available tag history. + if ! is_ci_base_target; then + prepare_ci_build_context + fi extract_dependency_pins write_rocm_build_arg_override compute_dependency_cache_keys @@ -2373,6 +2627,8 @@ main() { compute_rocm_csrc_content_hash_if_needed compute_rocm_rust_content_hash_if_needed write_rocm_cache_override + # Keep the context override last so every bake target uses the owned tree. + write_build_context_override resolve_ci_base_dependency_targets print_bake_config if [[ "${BAKE_PRINT_ONLY:-0}" == "1" ]]; then diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index f612a0bd0ec9..4595daa56f3e 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -5,6 +5,7 @@ ARG VLLM_BRANCH="main" ARG COMMON_WORKDIR=/app ARG BASE_IMAGE=rocm/vllm-dev:base ARG CI_BASE_IMAGE=rocm/vllm-dev:ci_base +ARG ROCM_TRITON_KERNELS_COMMIT=0f380657dbf3ee86eb57558ff71df24f03b5d4e7 # NIC backend for MoRI RDMA support. # By default (all), drivers and userspace libraries for all supported NIC types # (ainic and bnxt) are installed; MoRI selects the appropriate one at runtime. @@ -272,6 +273,20 @@ COPY --from=fetch_vllm \ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ uv pip install --system -r requirements/rocm.txt +# Fetch only the pinned ROCm Triton kernels package. Keeping this before the +# vLLM source COPYs makes the download reusable when native sources change. +FROM base AS rocm-triton-kernels +ARG ROCM_TRITON_KERNELS_COMMIT +ADD https://codeload.github.com/ROCm/triton/tar.gz/${ROCM_TRITON_KERNELS_COMMIT} \ + /tmp/rocm-triton.tar.gz +RUN mkdir -p /opt/rocm-triton-kernels \ + && tar --extract --gzip --file /tmp/rocm-triton.tar.gz \ + --directory /opt/rocm-triton-kernels --strip-components=4 \ + --no-same-owner \ + "triton-${ROCM_TRITON_KERNELS_COMMIT}/python/triton_kernels/triton_kernels" \ + && test -f /opt/rocm-triton-kernels/__init__.py \ + && rm /tmp/rocm-triton.tar.gz + # ----------------------- # vLLM native build stages # @@ -279,6 +294,11 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ # compilation. That keeps unrelated CI/test/docs edits from invalidating the # expensive HIP/C++ build layer. FROM build_vllm_dependencies AS csrc-build +ARG ROCM_TRITON_KERNELS_COMMIT +COPY --from=rocm-triton-kernels \ + /opt/rocm-triton-kernels /opt/rocm-triton-kernels +ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels + # pyproject.toml is bind-mounted so it is not persisted in the csrc-build image, # but it still participates in the BuildKit cache key for this RUN step. COPY setup.py CMakeLists.txt ./ @@ -292,13 +312,15 @@ ENV SETUPTOOLS_SCM_PRETEND_VERSION="0.0.0+rocm.csrc.build" RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproject.toml \ --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ - export CCACHE_BASEDIR="$PWD" \ + grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + cmake/external_projects/triton_kernels.cmake \ + && export CCACHE_BASEDIR="$PWD" \ && echo "=== ccache stats before ROCm native build ===" \ && (ccache --show-stats || true) \ && (ccache --zero-stats || true) \ && EFFECTIVE_MAX_JOBS="${MAX_JOBS:-$(nproc)}" \ && echo "Building ROCm native extension wheel with MAX_JOBS=${EFFECTIVE_MAX_JOBS}" \ - && LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ + && PYTHONDONTWRITEBYTECODE=1 LDFLAGS="-fuse-ld=mold" MAX_JOBS="${EFFECTIVE_MAX_JOBS}" python3 setup.py bdist_wheel --dist-dir=dist \ && test -d dist \ && ls dist/*.whl >/dev/null \ && echo "=== ccache stats after ROCm native build ===" \ @@ -322,7 +344,7 @@ RUN cd "${COMMON_WORKDIR}/vllm" \ && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ && export VLLM_DOCKER_BUILD_CONTEXT=1 \ && echo "Packaging vLLM ROCm wheel using precompiled extensions from ${VLLM_PRECOMPILED_WHEEL_LOCATION}" \ - && python3 setup.py bdist_wheel --dist-dir=dist \ + && PYTHONDONTWRITEBYTECODE=1 python3 setup.py bdist_wheel --dist-dir=dist \ && test -d dist \ && ls dist/*.whl >/dev/null FROM scratch AS export_vllm @@ -551,6 +573,10 @@ RUN /bin/bash -lc 'set -euo pipefail; \ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR +ARG ROCM_TRITON_KERNELS_COMMIT +COPY --from=rocm-triton-kernels \ + /opt/rocm-triton-kernels /opt/rocm-triton-kernels +ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -590,7 +616,9 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ # This ensures setuptools_scm sees clean repo state for version detection RUN --mount=type=bind,source=.git,target=vllm/.git \ --mount=type=cache,target=/root/.cache/uv \ - cd vllm \ + grep -Fq "set(TRITON_KERNELS_TAG \"${ROCM_TRITON_KERNELS_COMMIT}\")" \ + vllm/cmake/external_projects/triton_kernels.cmake \ + && cd vllm \ && uv pip install --system setuptools_scm regex \ && VLLM_VERSION=$(python3 -c "import setuptools_scm; print(setuptools_scm.get_version())") \ && echo "Detected vLLM version: ${VLLM_VERSION}" \ From 0f2ea973beebfe2959c74f06e2a412c55c1c7007 Mon Sep 17 00:00:00 2001 From: Itay Etelis <92247226+Etelis@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:03:59 +0300 Subject: [PATCH 26/30] [KV Connector][Offloading] Keep per-layer KV registration when canonical_layout is requested (#51688) Signed-off-by: Itay Etelis Co-authored-by: Itay Etelis --- .../unit/offloading_connector/test_config.py | 28 +++++++++++++++++++ .../kv_connector/v1/offloading_connector.py | 4 ++- 2 files changed, 31 insertions(+), 1 deletion(-) diff --git a/tests/v1/kv_connector/unit/offloading_connector/test_config.py b/tests/v1/kv_connector/unit/offloading_connector/test_config.py index cc58aaf7a7c9..0b0b6b874f74 100644 --- a/tests/v1/kv_connector/unit/offloading_connector/test_config.py +++ b/tests/v1/kv_connector/unit/offloading_connector/test_config.py @@ -661,6 +661,34 @@ def test_canonical_layout_certifies_v2_model_runner(): ).parallel.is_parallelism_agnostic +def test_prefer_cross_layer_blocks_yields_to_canonical_layout(): + """The connector must not request cross-layer blocks under + canonical_layout: cross-layer slabs have no per-layer refs to certify.""" + from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorRole + from vllm.distributed.kv_transfer.kv_connector.v1.offloading_connector import ( + OffloadingConnector, + ) + + kv_cache_config = _make_kv_cache_config() + connector_module = "vllm.distributed.kv_transfer.kv_connector.v1" + + def make_connector(extra_config: dict[str, Any] | None) -> OffloadingConnector: + with ( + patch(f"{connector_module}.offloading_connector.OffloadingSpecFactory"), + patch( + f"{connector_module}.offloading_connector.OffloadingConnectorScheduler" + ), + ): + return OffloadingConnector( + _make_vllm_config(extra_config=extra_config), + KVConnectorRole.SCHEDULER, + kv_cache_config, + ) + + assert make_connector(None).prefer_cross_layer_blocks + assert not make_connector({"canonical_layout": True}).prefer_cross_layer_blocks + + def test_parallelism_agnostic_disabled_on_v2_model_runner(): config = _make_vllm_config() config.use_v2_model_runner = True diff --git a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py index a32d5d9cc8ea..f003d0d6f88d 100644 --- a/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py +++ b/vllm/distributed/kv_transfer/kv_connector/v1/offloading_connector.py @@ -49,7 +49,8 @@ class OffloadingConnector(KVConnectorBase_V1, SupportsHMA): @property def prefer_cross_layer_blocks(self) -> bool: - return True + # Cross-layer slabs have no per-layer refs to certify canonical mappings + return not self._canonical_layout @property def requires_kv_delivery(self) -> bool: @@ -66,6 +67,7 @@ def __init__( super().__init__(vllm_config, role, kv_cache_config) offloading_config = build_offloading_config(vllm_config, kv_cache_config) + self._canonical_layout = offloading_config.canonical_layout spec = OffloadingSpecFactory.create_spec(offloading_config) self.connector_scheduler: OffloadingConnectorScheduler | None = None From 0914ed2e816fa1987951f22afba1c00b7786f110 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Mon, 10 Aug 2026 23:52:21 -0400 Subject: [PATCH 27/30] [Perf] Adaptive budget for spec scheduled input tokens, ~60% better Kimi K3 DSpark TTFT (#51725) Signed-off-by: yewentao256 --- tests/v1/core/test_scheduler.py | 24 ++++++++++++++++++++++++ tests/v1/core/utils.py | 2 ++ vllm/config/vllm.py | 27 ++++++++------------------- vllm/v1/core/sched/scheduler.py | 25 ++++++++++++++++++++----- 4 files changed, 54 insertions(+), 24 deletions(-) diff --git a/tests/v1/core/test_scheduler.py b/tests/v1/core/test_scheduler.py index 0e46f54f6f59..ccd0cf3e9333 100644 --- a/tests/v1/core/test_scheduler.py +++ b/tests/v1/core/test_scheduler.py @@ -1280,6 +1280,30 @@ def test_reset_connector_cache_no_connector_is_no_op_success(): assert scheduler.reset_prefix_cache(reset_connector=True) is True +def test_draft_slots_budgeted_per_scheduled_request(tmp_path, monkeypatch): + monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "0") + (tmp_path / "config.json").write_text( + '{"architectures": ["OPTForCausalLM"], "model_type": "opt"}' + ) + scheduler = create_scheduler( + model=str(tmp_path), + max_num_seqs=16, + max_num_batched_tokens=20, + num_speculative_tokens=4, + parallel_drafting=True, + skip_tokenizer_init=True, + ) + speculative_config = scheduler.vllm_config.speculative_config + assert speculative_config is not None + assert scheduler.max_num_scheduled_tokens == 20 + assert speculative_config.max_num_new_slots_for_drafting == 3 + + for request in create_requests(num_requests=2, num_tokens=10): + scheduler.add_request(request) + + assert scheduler.schedule().num_scheduled_tokens == {"0": 10, "1": 4} + + # Note - these test cases mirror some of those in test_rejection_sampler.py @pytest.mark.parametrize( "spec_tokens,output_tokens,expected", diff --git a/tests/v1/core/utils.py b/tests/v1/core/utils.py index 5807574c6444..0baa68af2589 100644 --- a/tests/v1/core/utils.py +++ b/tests/v1/core/utils.py @@ -61,6 +61,7 @@ def create_scheduler( max_model_len: int | None = None, num_speculative_tokens: int | None = None, speculative_method: str | None = None, + parallel_drafting: bool = False, skip_tokenizer_init: bool = False, async_scheduling: bool = False, pipeline_parallel_size: int = 1, @@ -151,6 +152,7 @@ def create_scheduler( spec_kwargs["prompt_lookup_max"] = num_speculative_tokens spec_kwargs["prompt_lookup_min"] = 1 speculative_config = SpeculativeConfig(**spec_kwargs) + speculative_config.parallel_drafting = parallel_drafting ec_transfer_config = ( ECTransferConfig( diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 454514f0f205..343feba36c52 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -1782,21 +1782,15 @@ def update_sizes_for_sequence_parallelism(self, possible_sizes: list) -> list: def _set_max_num_scheduled_tokens(self): """ In most cases, the scheduler may schedule a batch with as many tokens as the - worker is configured to handle. However for some speculative decoding methods, - the drafter model may insert additional slots into the batch when drafting. - To account for this, we need to decrease the max_num_scheduled_tokens by an - upper bound on the number of slots that can be added. + worker is configured to handle. """ if self.speculative_config is not None: scheduled_token_delta = ( self.speculative_config.max_num_new_slots_for_drafting - * self.scheduler_config.max_num_seqs ) max_num_batched_tokens = self.scheduler_config.max_num_batched_tokens if self.scheduler_config.max_num_scheduled_tokens is None: - self.scheduler_config.max_num_scheduled_tokens = ( - max_num_batched_tokens - scheduled_token_delta - ) + self.scheduler_config.max_num_scheduled_tokens = max_num_batched_tokens if self.scheduler_config.max_num_scheduled_tokens <= 0: raise ValueError( @@ -1805,7 +1799,7 @@ def _set_max_num_scheduled_tokens(self): " the speculative decoding settings, which does not allow" " any tokens to be scheduled. Increase max_num_batched_tokens" " to accommodate the additional draft token slots, or decrease" - " num_speculative_tokens or max_num_seqs." + " num_speculative_tokens." ) if self.scheduler_config.max_num_scheduled_tokens < 8192: logger.warning_once( @@ -1814,19 +1808,14 @@ def _set_max_num_scheduled_tokens(self): " the speculative decoding settings. This may lead to suboptimal" " performance. Consider increasing max_num_batched_tokens to" " accommodate the additional draft token slots, or decrease" - " num_speculative_tokens or max_num_seqs.", + " num_speculative_tokens.", ) - max_num_scheduled_tokens = self.scheduler_config.max_num_scheduled_tokens - if max_num_batched_tokens < max_num_scheduled_tokens + ( - self.speculative_config.max_num_new_slots_for_drafting - * self.scheduler_config.max_num_seqs - ): + if max_num_batched_tokens <= scheduled_token_delta: raise ValueError( - f"VllmConfig received max_num_scheduled_tokens but it does not have" - " enough slots to support the speculative decoding settings." - f" It should be greater by at least {scheduled_token_delta}, but" - f" got {max_num_batched_tokens=} and {max_num_scheduled_tokens=}." + "VllmConfig does not have enough slots to schedule a token and" + " support the speculative decoding settings." + f" Got {max_num_batched_tokens=} and {scheduled_token_delta=}." ) def _set_cudagraph_sizes(self): diff --git a/vllm/v1/core/sched/scheduler.py b/vllm/v1/core/sched/scheduler.py index c08a1302131e..a781b4580149 100644 --- a/vllm/v1/core/sched/scheduler.py +++ b/vllm/v1/core/sched/scheduler.py @@ -456,6 +456,9 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_to_new_blocks: dict[str, KVCacheBlocks] = {} num_scheduled_tokens: dict[str, int] = {} token_budget = self.max_num_scheduled_tokens + spec = self.vllm_config.speculative_config + draft_slots = spec.max_num_new_slots_for_drafting if spec is not None else 0 + input_budget = self.scheduler_config.max_num_batched_tokens if self._pause_state == PauseState.PAUSED_ALL: # Do not schedule any requests when paused. token_budget = 0 @@ -483,6 +486,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_index = 0 while req_index < len(self.running) and token_budget > 0: request = self.running[req_index] + if input_budget <= draft_slots: + break if ( request.num_output_placeholders > 0 @@ -519,7 +524,9 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: ) if 0 < self.scheduler_config.long_prefill_token_threshold < num_new_tokens: num_new_tokens = self.scheduler_config.long_prefill_token_threshold - num_new_tokens = min(num_new_tokens, token_budget) + num_new_tokens = min( + num_new_tokens, token_budget, input_budget - draft_slots + ) # Make sure the input position does not exceed the max model len. # This is necessary when using spec decoding. @@ -605,7 +612,9 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: if preempted_req in scheduled_running_reqs: preempted_req_id = preempted_req.request_id scheduled_running_reqs.remove(preempted_req) - token_budget += num_scheduled_tokens.pop(preempted_req_id) + restored = num_scheduled_tokens.pop(preempted_req_id) + token_budget += restored + input_budget += restored + draft_slots req_to_new_blocks.pop(preempted_req_id) scheduled_spec_decode_tokens.pop(preempted_req_id, None) preempted_encoder_inputs = scheduled_encoder_inputs.pop( @@ -643,6 +652,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: req_to_new_blocks[request_id] = new_blocks num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens + input_budget -= num_new_tokens + draft_slots req_index += 1 # Speculative decode related. @@ -693,6 +703,8 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: step_skipped_waiting = create_request_queue(self.policy) while (self.waiting or self.skipped_waiting) and token_budget > 0: + if input_budget <= draft_slots: + break # Paused streaming sessions (WAITING_FOR_STREAMING_REQ) are not # in `running` but still hold a model-runner request slot. num_running = len(self.running) + self.num_waiting_for_streaming_input @@ -866,6 +878,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # compute to a cadence-aligned step. break else: + request_token_budget = min(token_budget, input_budget - draft_slots) # Number of tokens to be scheduled. # We use `request.num_tokens` instead of # `request.num_prompt_tokens` to consider the resumed @@ -883,7 +896,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: ): num_new_tokens = 1 + self.num_spec_tokens if ( - num_new_tokens > token_budget + num_new_tokens > request_token_budget or num_computed_tokens + num_new_tokens > self.max_model_len ): # Prefer to not schedule than schedule un-padded here. @@ -898,13 +911,13 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: # pooling requests to be chunked if ( not self.scheduler_config.enable_chunked_prefill - and num_new_tokens > token_budget + and num_new_tokens > request_token_budget ): # If chunked_prefill is disabled, # we can stop the scheduling here. break - num_new_tokens = min(num_new_tokens, token_budget) + num_new_tokens = min(num_new_tokens, request_token_budget) assert num_new_tokens > 0 # Apply Mamba alignment before encoder caps. @@ -1065,6 +1078,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: ) num_scheduled_tokens[request_id] = num_new_tokens token_budget -= num_new_tokens + input_budget -= num_new_tokens + draft_slots request.status = RequestStatus.RUNNING request.num_computed_tokens = num_computed_tokens if pad_spec_decode: @@ -1104,6 +1118,7 @@ def schedule(self, throttle_prefills: bool = False) -> SchedulerOutput: assert total_num_scheduled_tokens <= self.max_num_scheduled_tokens assert token_budget >= 0 + assert input_budget >= 0 assert len(self.running) <= self.max_num_running_reqs # Since some requests in the RUNNING queue may not be scheduled in # this step, the total number of scheduled requests can be smaller than From 07443bea29db509cfc84b060317c7077b343376a Mon Sep 17 00:00:00 2001 From: Ming Huang Date: Mon, 10 Aug 2026 20:56:25 -0700 Subject: [PATCH 28/30] [BugFix][Core] free_blocks: restore prepend (LIFO) reuse order when prefix caching is off (#51482) Signed-off-by: theminghuang Signed-off-by: Ming Huang --- vllm/v1/core/block_pool.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/vllm/v1/core/block_pool.py b/vllm/v1/core/block_pool.py index a64b7e321596..6b73f91afe0c 100644 --- a/vllm/v1/core/block_pool.py +++ b/vllm/v1/core/block_pool.py @@ -732,7 +732,7 @@ def free_blocks(self, ordered_blocks: Iterable[KVCacheBlock]) -> None: if block.ref_cnt == 0 and not block.is_null: # When caching is disabled we always append for better # GPU cache locality from reusing recently used blocks - if block.block_hash is None and self.enable_caching: + if block.block_hash is None or not self.enable_caching: blocks_without_hash.append(block) else: blocks_with_hash.append(block) From adc1200597c0e8e024173bf38270a660b176642d Mon Sep 17 00:00:00 2001 From: Frank <147031176+frank-suwen@users.noreply.github.com> Date: Mon, 10 Aug 2026 21:12:22 -0700 Subject: [PATCH 29/30] [Misc] Use VLLMValidationError in scoring input validation (#51753) Signed-off-by: frank-suwen --- .../entrypoints/pooling/scoring/test_utils.py | 65 +++++++++++++++++++ vllm/entrypoints/pooling/scoring/utils.py | 11 ++-- 2 files changed, 72 insertions(+), 4 deletions(-) create mode 100644 tests/entrypoints/pooling/scoring/test_utils.py diff --git a/tests/entrypoints/pooling/scoring/test_utils.py b/tests/entrypoints/pooling/scoring/test_utils.py new file mode 100644 index 000000000000..35b205eebf07 --- /dev/null +++ b/tests/entrypoints/pooling/scoring/test_utils.py @@ -0,0 +1,65 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +import pytest + +from vllm.entrypoints.pooling.scoring.typing import ScoreInput +from vllm.entrypoints.pooling.scoring.utils import validate_score_input +from vllm.exceptions import VLLMValidationError + + +@pytest.mark.parametrize( + ("data_1", "data_2", "is_multimodal_model", "architecture", "message"), + [ + pytest.param( + {"content": []}, + "document", + False, + "TestModel", + "MultiModalParam is not supported for TestModel", + id="unsupported-multimodal-input", + ), + pytest.param( + ["query 1", "query 2"], + ["document"], + False, + "TestModel", + "Input lengths must be either 1:1, 1:N or N:N", + id="incompatible-input-lengths", + ), + pytest.param( + [], + ["document"], + False, + "TestModel", + "At least one text element must be given", + id="empty-first-input", + ), + pytest.param( + ["query"], + [], + False, + "TestModel", + "At least one text_pair element must be given", + id="empty-second-input", + ), + ], +) +def test_validate_score_input_rejects_invalid_inputs( + data_1: ScoreInput | list[ScoreInput], + data_2: ScoreInput | list[ScoreInput], + is_multimodal_model: bool, + architecture: str, + message: str, +): + with pytest.raises(VLLMValidationError) as exc_info: + validate_score_input( + data_1, + data_2, + is_multimodal_model=is_multimodal_model, + architecture=architecture, + ) + + assert str(exc_info.value) == message + assert exc_info.value.parameter is None + assert exc_info.value.value is None diff --git a/vllm/entrypoints/pooling/scoring/utils.py b/vllm/entrypoints/pooling/scoring/utils.py index 13db389b8233..148b429d3c5a 100644 --- a/vllm/entrypoints/pooling/scoring/utils.py +++ b/vllm/entrypoints/pooling/scoring/utils.py @@ -15,6 +15,7 @@ MultiModalItemTracker, _parse_chat_message_content_parts, ) +from vllm.exceptions import VLLMValidationError from vllm.inputs import MultiModalDataDict, MultiModalUUIDDict from .typing import ( @@ -85,7 +86,9 @@ def _validate_mm_score_input( out.append(d) else: if not is_multimodal_model: - raise ValueError(f"MultiModalParam is not supported for {architecture}") + raise VLLMValidationError( + f"MultiModalParam is not supported for {architecture}" + ) content = cast(list[ScoreContentPartParam], d.get("content", [])) out.append(content) return out @@ -99,11 +102,11 @@ def _validate_score_input_lens( len_2 = len(data_2) if len_1 > 1 and len_1 != len_2: - raise ValueError("Input lengths must be either 1:1, 1:N or N:N") + raise VLLMValidationError("Input lengths must be either 1:1, 1:N or N:N") if len_1 == 0: - raise ValueError("At least one text element must be given") + raise VLLMValidationError("At least one text element must be given") if len_2 == 0: - raise ValueError("At least one text_pair element must be given") + raise VLLMValidationError("At least one text_pair element must be given") def validate_score_input( From 90fd4a333fd62b09310924dcb04d8e827bcbcd9b Mon Sep 17 00:00:00 2001 From: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Date: Tue, 11 Aug 2026 00:33:10 -0400 Subject: [PATCH 30/30] Preserve revision pins in secondary artifact loaders (#51446) Signed-off-by: Clinton Thomas <1033162+KernelClint@users.noreply.github.com> Co-authored-by: Lucas Bourtoule <35483370+dhalf@users.noreply.github.com> --- tests/config/test_model_arch_config.py | 27 ++++++++++ tests/test_quantization_revision_pin.py | 40 +++++++++++++++ tests/test_tensorizer_revision_pin.py | 51 +++++++++++++++++++ tests/transformers_utils/test_config.py | 47 +++++++++++++++++ vllm/config/model.py | 8 ++- vllm/config/vllm.py | 1 + .../model_executor/model_loader/tensorizer.py | 11 +++- vllm/model_executor/models/funaudiochat.py | 9 +++- vllm/renderers/hf.py | 15 +++++- vllm/transformers_utils/config.py | 2 + .../model_arch_config_convertor.py | 28 +++++++--- 11 files changed, 226 insertions(+), 13 deletions(-) create mode 100644 tests/test_quantization_revision_pin.py create mode 100644 tests/test_tensorizer_revision_pin.py diff --git a/tests/config/test_model_arch_config.py b/tests/config/test_model_arch_config.py index a4475de7cdac..cf65c0e159b7 100644 --- a/tests/config/test_model_arch_config.py +++ b/tests/config/test_model_arch_config.py @@ -7,6 +7,7 @@ from pathlib import Path from types import SimpleNamespace from typing import cast +from unittest.mock import Mock import pytest from transformers import PretrainedConfig @@ -138,6 +139,32 @@ def test_head_size_falls_back_when_head_dim_is_zero(): assert convertor.get_head_size() == 128 +def test_qk_rope_head_dim_recovery_uses_model_revision(monkeypatch: pytest.MonkeyPatch): + hf_config = PretrainedConfig( + model_type="deepseek_v2", + qk_rope_head_dim=128, + qk_nope_head_dim=128, + ) + hf_config.name_or_path = "org/model" + get_hf_file_to_dict = Mock(return_value={"qk_rope_head_dim": 64}) + + monkeypatch.setattr( + "vllm.transformers_utils.repo_utils.get_hf_file_to_dict", + get_hf_file_to_dict, + ) + + convertor = ModelArchConfigConvertorBase( + hf_config, + hf_config, + revision="pinned-revision", + ) + + assert convertor._get_qk_rope_head_dim() == 64 + get_hf_file_to_dict.assert_called_once_with( + "config.json", "org/model", "pinned-revision" + ) + + @pytest.mark.parametrize( "attribute", [ diff --git a/tests/test_quantization_revision_pin.py b/tests/test_quantization_revision_pin.py new file mode 100644 index 000000000000..6c121ba03179 --- /dev/null +++ b/tests/test_quantization_revision_pin.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm.config.vllm import VllmConfig + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + + +def test_quantization_config_forwards_model_revision(monkeypatch: pytest.MonkeyPatch): + quant_config = Mock() + quant_config.get_min_capability.return_value = 0 + quant_config.get_supported_act_dtypes.return_value = [torch.float32] + model_config = SimpleNamespace( + quantization="auto_gptq", + model="repo/model", + hf_config=object(), + dtype=torch.float32, + revision="pinned-revision", + ) + + monkeypatch.setattr( + "vllm.model_executor.model_loader.weight_utils.get_quant_config", + lambda model_config, load_config: quant_config, + ) + monkeypatch.setattr( + "vllm.platforms.current_platform.get_device_capability", lambda: None + ) + + VllmConfig._get_quantization_config(model_config, SimpleNamespace()) + + quant_config.maybe_update_config.assert_called_once_with( + "repo/model", + hf_config=model_config.hf_config, + revision="pinned-revision", + ) diff --git a/tests/test_tensorizer_revision_pin.py b/tests/test_tensorizer_revision_pin.py new file mode 100644 index 000000000000..5c2507afd126 --- /dev/null +++ b/tests/test_tensorizer_revision_pin.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import contextlib +import io +from types import SimpleNamespace +from unittest.mock import Mock + +import pytest +import torch + +from vllm.model_executor.model_loader import tensorizer + +pytestmark = [pytest.mark.cpu_test, pytest.mark.skip_global_cleanup] + + +def test_serialize_vllm_model_forwards_model_revision( + monkeypatch: pytest.MonkeyPatch, +): + @contextlib.contextmanager + def fake_open_stream(*args, **kwargs): + yield io.BytesIO() + + snapshot_download = Mock() + monkeypatch.setattr(tensorizer, "open_stream", fake_open_stream) + monkeypatch.setattr(tensorizer, "TensorSerializer", Mock()) + monkeypatch.setattr( + tensorizer, + "hf_api", + lambda: SimpleNamespace(snapshot_download=snapshot_download), + ) + + model = torch.nn.Module() + tensorizer_config = SimpleNamespace( + _construct_tensorizer_args=lambda: SimpleNamespace( + tensorizer_uri="/tmp/model.tensors", + tensorizer_dir="/tmp/out", + stream_kwargs={}, + ), + encryption_keyfile=None, + _is_sharded=False, + serialization_kwargs=None, + ) + model_config = SimpleNamespace( + served_model_name="repo/model", + revision="pinned-revision", + ) + + tensorizer.serialize_vllm_model(model, tensorizer_config, model_config) + + assert snapshot_download.call_args.args == ("repo/model",) + assert snapshot_download.call_args.kwargs["revision"] == "pinned-revision" diff --git a/tests/transformers_utils/test_config.py b/tests/transformers_utils/test_config.py index 5a7421b6aa36..42261feb536f 100644 --- a/tests/transformers_utils/test_config.py +++ b/tests/transformers_utils/test_config.py @@ -6,7 +6,15 @@ `BaseRenderer.get_eos_token_id`. """ +from types import SimpleNamespace +from typing import cast +from unittest.mock import patch + +from transformers import PretrainedConfig + +from vllm.config.model import ModelConfig from vllm.tokenizers import get_tokenizer +from vllm.transformers_utils import config as config_module from vllm.transformers_utils.config import try_get_generation_config @@ -30,3 +38,42 @@ def test_get_blip2_eos_token(): generation_config = try_get_generation_config(model_name, trust_remote_code=False) assert generation_config is not None assert generation_config.eos_token_id == 50118 + + +def test_model_config_generation_fallback_forwards_code_revision(): + model_config = cast( + ModelConfig, + SimpleNamespace( + generation_config="auto", + hf_config_path=None, + model="org/model", + trust_remote_code=True, + revision="model-pin", + code_revision="code-pin", + config_format="auto", + hf_token=None, + ), + ) + + with ( + patch.object( + config_module.GenerationConfig, + "from_pretrained", + side_effect=OSError, + ), + patch.object( + config_module, + "get_config", + return_value=PretrainedConfig(), + ) as get_config, + ): + ModelConfig.try_get_generation_config(model_config) + + get_config.assert_called_once_with( + "org/model", + trust_remote_code=True, + revision="model-pin", + code_revision="code-pin", + config_format="auto", + token=None, + ) diff --git a/vllm/config/model.py b/vllm/config/model.py index e1d0e2568f34..c09fdd7d76d3 100644 --- a/vllm/config/model.py +++ b/vllm/config/model.py @@ -890,7 +890,11 @@ def get_model_arch_config(self) -> ModelArchitectureConfig: convertor_cls = MODEL_ARCH_CONFIG_CONVERTORS.get( self.hf_config.model_type, ModelArchConfigConvertorBase ) - convertor = convertor_cls(self.hf_config, self.hf_text_config) + convertor = convertor_cls( + self.hf_config, + self.hf_text_config, + revision=getattr(self, "revision", None), + ) return convertor.convert( supports_multimodal=self._supports_multimodal_for_mm_prefix() ) @@ -1651,6 +1655,7 @@ def try_get_generation_config(self) -> dict[str, Any]: self.hf_config_path or self.model, trust_remote_code=self.trust_remote_code, revision=self.revision, + code_revision=self.code_revision, config_format=self.config_format, hf_token=self.hf_token, ) @@ -1658,6 +1663,7 @@ def try_get_generation_config(self) -> dict[str, Any]: config = try_get_generation_config( self.generation_config, trust_remote_code=self.trust_remote_code, + code_revision=self.code_revision, config_format=self.config_format, hf_token=self.hf_token, ) diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index 343feba36c52..91372385e014 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -779,6 +779,7 @@ def _get_quantization_config( quant_config.maybe_update_config( model_config.model, hf_config=model_config.hf_config, + revision=model_config.revision, ) return quant_config return None diff --git a/vllm/model_executor/model_loader/tensorizer.py b/vllm/model_executor/model_loader/tensorizer.py index befe560929d1..3ef9b6bf2a59 100644 --- a/vllm/model_executor/model_loader/tensorizer.py +++ b/vllm/model_executor/model_loader/tensorizer.py @@ -621,7 +621,9 @@ def is_vllm_tensorized(tensorizer_config: "TensorizerConfig") -> bool: def serialize_extra_artifacts( - tensorizer_args: TensorizerArgs, served_model_name: str | list[str] | None + tensorizer_args: TensorizerArgs, + served_model_name: str | list[str] | None, + revision: str | None = None, ) -> None: if not isinstance(served_model_name, str): raise ValueError( @@ -632,6 +634,7 @@ def serialize_extra_artifacts( with tempfile.TemporaryDirectory() as tmpdir: hf_api().snapshot_download( served_model_name, + revision=revision, local_dir=tmpdir, ignore_patterns=[ "*.pt", @@ -693,7 +696,11 @@ def serialize_vllm_model( serializer.write_module(model) serializer.close() - serialize_extra_artifacts(tensorizer_args, model_config.served_model_name) + serialize_extra_artifacts( + tensorizer_args, + model_config.served_model_name, + revision=model_config.revision, + ) logger.info("Successfully serialized model to %s", str(output_file)) return model diff --git a/vllm/model_executor/models/funaudiochat.py b/vllm/model_executor/models/funaudiochat.py index 2d3334f27513..0a9e286783c4 100644 --- a/vllm/model_executor/models/funaudiochat.py +++ b/vllm/model_executor/models/funaudiochat.py @@ -519,12 +519,17 @@ class FunAudioChatProcessingInfo(BaseProcessingInfo): @cached_property def feature_extractor(self) -> WhisperFeatureExtractor: - return WhisperFeatureExtractor.from_pretrained(self.model_id) + return WhisperFeatureExtractor.from_pretrained( + self.model_id, + revision=self.ctx.model_config.revision, + ) @cached_property def speech_tokenizer(self) -> TokenizersBackend: return TokenizersBackend.from_pretrained( - self.model_id, subfolder="speech_tokenizer" + self.model_id, + subfolder="speech_tokenizer", + revision=self.ctx.model_config.tokenizer_revision, ) def get_feature_extractor(self) -> WhisperFeatureExtractor: diff --git a/vllm/renderers/hf.py b/vllm/renderers/hf.py index 95210e4a47b4..6afd245cc143 100644 --- a/vllm/renderers/hf.py +++ b/vllm/renderers/hf.py @@ -210,7 +210,7 @@ def _build_mixed_prompt_embeds( return full_embeds, is_token_ids.tolist() -_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, bool], str | None]() +_PROCESSOR_CHAT_TEMPLATES = dict[tuple[str, str | None, str | None, bool], str | None]() """ Used in `_try_get_processor_chat_template` to avoid calling `cached_get_processor` again if the processor fails to be loaded. @@ -222,9 +222,16 @@ def _build_mixed_prompt_embeds( def _try_get_processor_chat_template( tokenizer: HfTokenizer, *, + revision: str | None, + code_revision: str | None, trust_remote_code: bool, ) -> str | None: - cache_key = (tokenizer.name_or_path, trust_remote_code) + cache_key = ( + tokenizer.name_or_path, + revision, + code_revision, + trust_remote_code, + ) if cache_key in _PROCESSOR_CHAT_TEMPLATES: return _PROCESSOR_CHAT_TEMPLATES[cache_key] @@ -234,6 +241,8 @@ def _try_get_processor_chat_template( processor = cached_get_processor( tokenizer.name_or_path, processor_cls=(PythonBackend, TokenizersBackend, ProcessorMixin), + revision=revision, + code_revision=code_revision, trust_remote_code=trust_remote_code, ) if ( @@ -271,6 +280,8 @@ def resolve_chat_template( if tools is None: chat_template = _try_get_processor_chat_template( tokenizer, + revision=model_config.revision, + code_revision=model_config.code_revision, trust_remote_code=model_config.trust_remote_code, ) if chat_template is not None: diff --git a/vllm/transformers_utils/config.py b/vllm/transformers_utils/config.py index 89da81ea918f..ad8f0535b2ea 100644 --- a/vllm/transformers_utils/config.py +++ b/vllm/transformers_utils/config.py @@ -1071,6 +1071,7 @@ def try_get_generation_config( model: str, trust_remote_code: bool, revision: str | None = None, + code_revision: str | None = None, config_format: str | ConfigFormat = "auto", hf_token: bool | str | None = None, ) -> GenerationConfig | None: @@ -1086,6 +1087,7 @@ def try_get_generation_config( model, trust_remote_code=trust_remote_code, revision=revision, + code_revision=code_revision, config_format=config_format, token=hf_token, ) diff --git a/vllm/transformers_utils/model_arch_config_convertor.py b/vllm/transformers_utils/model_arch_config_convertor.py index 2e89c6aa3281..2e6bd05cd992 100644 --- a/vllm/transformers_utils/model_arch_config_convertor.py +++ b/vllm/transformers_utils/model_arch_config_convertor.py @@ -24,9 +24,15 @@ class ModelArchConfigConvertorBase: - def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig): + def __init__( + self, + hf_config: PretrainedConfig, + hf_text_config: PretrainedConfig, + revision: str | None = None, + ): self.hf_config = hf_config self.hf_text_config = hf_text_config + self.revision = revision def get_per_layer_hf_configs( self, @@ -122,7 +128,7 @@ def _get_qk_rope_head_dim(self) -> int: model_path = self.hf_config.name_or_path if not model_path: return qk_rope_head_dim - raw = get_hf_file_to_dict("config.json", model_path) + raw = get_hf_file_to_dict("config.json", model_path, self.revision) if raw and "qk_rope_head_dim" in raw: correct = raw["qk_rope_head_dim"] if correct != qk_rope_head_dim: @@ -595,16 +601,26 @@ def _strip_mimo_v2_attention_chunk_size( class MimoV2ModelArchConfigConvertor(ModelArchConfigConvertorBase): - def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig): + def __init__( + self, + hf_config: PretrainedConfig, + hf_text_config: PretrainedConfig, + revision: str | None = None, + ): if getattr(hf_config, "vision_config", None): hf_config.architectures = ["MiMoV2OmniForCausalLM"] - super().__init__(hf_config, hf_text_config) + super().__init__(hf_config, hf_text_config, revision) _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) class MimoV2MTPModelArchConfigConvertor(ModelArchConfigConvertorBase): - def __init__(self, hf_config: PretrainedConfig, hf_text_config: PretrainedConfig): - super().__init__(hf_config, hf_text_config) + def __init__( + self, + hf_config: PretrainedConfig, + hf_text_config: PretrainedConfig, + revision: str | None = None, + ): + super().__init__(hf_config, hf_text_config, revision) _strip_mimo_v2_attention_chunk_size(hf_config, hf_text_config) def get_num_hidden_layers(self) -> int: