From c296cf8259ed834b7174c59a2ee1f2f25967a87b Mon Sep 17 00:00:00 2001 From: Jakub Byczkowski Date: Tue, 18 Aug 2026 01:54:05 +0200 Subject: [PATCH 01/15] [Bugfix] Add forward_xpu to XDRotaryEmbedding for HunyuanOCR on XPU (#52174) Signed-off-by: Jakub Byczkowski Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- .../model_executor/layers/rotary_embedding/xdrope.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/vllm/model_executor/layers/rotary_embedding/xdrope.py b/vllm/model_executor/layers/rotary_embedding/xdrope.py index dab7aad9759a..a799b10a190e 100644 --- a/vllm/model_executor/layers/rotary_embedding/xdrope.py +++ b/vllm/model_executor/layers/rotary_embedding/xdrope.py @@ -137,6 +137,18 @@ def forward_cuda( key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape) return query, key + def forward_xpu( + self, + positions: torch.Tensor, + query: torch.Tensor, + key: torch.Tensor | None = None, + offsets: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + # No fused XPU kernel for XD-sectioned RoPE; the base + # RotaryEmbedding.forward_xpu calls the generic C++ kernel, which + # rejects the 2-D [4, num_tokens] xdrope positions. Use the native path. + return self.forward_native(positions, query, key, offsets) + @staticmethod def get_next_input_positions( context_len: int, From 58aa1e3d2692250c6bb24a0f852e3cb8e809a14f Mon Sep 17 00:00:00 2001 From: Tommy Asai <147074725+tommy-asai-sonarsource@users.noreply.github.com> Date: Tue, 18 Aug 2026 01:58:02 +0200 Subject: [PATCH 02/15] [Bugfix][SM120][MLA] Disable dense prefill for FlashInfer sparse MLA (#51395) Signed-off-by: tommy-asai-sonarsource Co-authored-by: Roger Wang --- tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py | 7 +++++++ .../attention/backends/mla/flashinfer_mla_sparse_sm120.py | 1 + 2 files changed, 8 insertions(+) diff --git a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py index 12a8cf5cbc39..724c8874b931 100644 --- a/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py +++ b/tests/v1/attention/test_flashinfer_sparse_mla_sm120_api.py @@ -34,6 +34,13 @@ def test_sm120_backend_uses_dedicated_backend_name() -> None: ) +def test_sm120_backend_uses_sparse_mqa_for_prefill() -> None: + impl_cls = FlashInferMLASparseSM120Backend.get_impl_cls() + + assert impl_cls.is_sparse + assert not impl_cls.supports_dense_mha_prefill + + def test_v32_glm_sm120_backend_accepts_glm_block_size( monkeypatch, ) -> None: diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py index 9db696a39c43..4de2c2f53f3b 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse_sm120.py @@ -33,6 +33,7 @@ class FlashInferMLASparseSM120Impl(MLAAttentionImpl[FlashInferMLASparseMetadata] """SM120 FlashInfer sparse-MLA implementation.""" is_sparse = True + supports_dense_mha_prefill = False def __init__( self, From 0e8989b41648e99176a53480e42c447f8bf95620 Mon Sep 17 00:00:00 2001 From: Kunshang Ji Date: Tue, 18 Aug 2026 08:05:48 +0800 Subject: [PATCH 03/15] [ROCm] gaurd on_gfx1250 call with rocm platform (#52625) Signed-off-by: Kunshang Ji --- .../model_executor/layers/quantization/utils/fp8_utils.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index 7910db192564..5df6bcb73421 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -881,9 +881,13 @@ def w8a8_triton_block_scaled_mm( torch.Tensor: The result of matmul. """ - from vllm.platforms.rocm import on_gfx1250 + _on_gfx1250 = False + if current_platform.is_rocm(): + from vllm.platforms.rocm import on_gfx1250 - if on_gfx1250(): + _on_gfx1250 = on_gfx1250() + + if _on_gfx1250: # Torch upcast reference: dequantize A,B to fp32 and matmul in fp32. # Avoids the gfx1250 native-fp8 block GEMM NaN bug. Correct but slow. _bn, _bk = block_size[0], block_size[1] From 0db502c8d8a620f29dd3560e69a22f0a7ed3477c Mon Sep 17 00:00:00 2001 From: Summer Yang Date: Mon, 17 Aug 2026 17:05:51 -0700 Subject: [PATCH 04/15] [Kimi-K3] support DCP partial prefix cache hit (#50493) Signed-off-by: Summer Yang Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> Co-authored-by: Yifan Qiao --- .../test_kimi_linear_context_parallel.py | 111 ++++++++++ tests/v1/attention/test_mla_backends.py | 43 ++++ .../test_partial_prefix_cache_hits.py | 204 +++++++++++++++++- .../test_partial_prefix_cache_primitives.py | 44 ++-- tests/v1/worker/test_gpu_model_runner_v2.py | 118 ++++++++++ vllm/v1/core/kv_cache_coordinator.py | 15 +- vllm/v1/worker/gpu/block_table.py | 9 +- vllm/v1/worker/gpu/model_runner.py | 16 +- 8 files changed, 523 insertions(+), 37 deletions(-) create mode 100644 tests/v1/worker/test_gpu_model_runner_v2.py diff --git a/tests/distributed/test_kimi_linear_context_parallel.py b/tests/distributed/test_kimi_linear_context_parallel.py index d37e19d0bd3b..06348074bb14 100644 --- a/tests/distributed/test_kimi_linear_context_parallel.py +++ b/tests/distributed/test_kimi_linear_context_parallel.py @@ -2,12 +2,15 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project import math +from pathlib import Path import pytest from tests.conftest import VllmRunner from tests.utils import create_new_process_for_each_test from vllm import SamplingParams, TokensPrompt +from vllm.config import CUDAGraphMode +from vllm.transformers_utils.configs.kimi_k3 import KimiK3Config MODEL = "moonshotai/Kimi-Linear-48B-A3B-Instruct" @@ -126,3 +129,111 @@ def test_kimi_linear_dcp_tiny( logprob_drifts.append(abs(baseline_logprob - dcp_logprob)) assert max(logprob_drifts) <= 1e-2 + + +def _make_tiny_k3_config(model_dir: Path) -> str: + config = KimiK3Config( + text_config=_make_tiny_overrides(), + vision_config={ + "vt_num_attention_heads": 2, + "vt_num_hidden_layers": 1, + "vt_hidden_size": 32, + "vt_intermediate_size": 64, + "qkv_hidden_size": 48, + }, + architectures=["KimiK3ForConditionalGeneration"], + ) + config.save_pretrained(model_dir) + return str(model_dir) + + +def _run_k3_partial_prefix_reuse( + vllm_runner: type[VllmRunner], model_name: str, dcp_size: int +) -> tuple[list[int], list[float], int]: + block_size = 256 + common_prefix = [3 + (token_idx % 251) for token_idx in range(block_size)] + prime_prompt = TokensPrompt(prompt_token_ids=common_prefix + [17]) + replay_prompt = TokensPrompt(prompt_token_ids=common_prefix + [18, 19]) + sampling_params = SamplingParams( + temperature=0, + max_tokens=8, + seed=0, + ignore_eos=True, + logprobs=20, + ) + + with vllm_runner( + model_name=model_name, + skip_tokenizer_init=True, + load_format="dummy", + language_model_only=True, + tensor_parallel_size=2, + decode_context_parallel_size=dcp_size, + cp_kv_cache_interleave_size=1, + distributed_executor_backend="mp", + dtype="bfloat16", + seed=0, + max_model_len=512, + max_num_seqs=4, + max_num_batched_tokens=256, + gpu_memory_utilization=0.85, + block_size=block_size, + enable_chunked_prefill=True, + enable_prefix_caching=True, + mamba_cache_mode="align", + ) as runner: + assert ( + runner.llm.llm_engine.vllm_config.compilation_config.cudagraph_mode + == CUDAGraphMode.FULL_AND_PIECEWISE + ) + runner.llm.generate( + [prime_prompt], + sampling_params, + use_tqdm=False, + ) + replay_output = runner.llm.generate( + [replay_prompt], + sampling_params, + use_tqdm=False, + )[0] + + completion = replay_output.outputs[0] + token_ids = list(completion.token_ids) + assert len(token_ids) == sampling_params.max_tokens + assert completion.logprobs is not None + assert len(completion.logprobs) == sampling_params.max_tokens + selected_logprobs = [ + step_logprobs[token_id].logprob + for token_id, step_logprobs in zip(token_ids, completion.logprobs) + ] + return token_ids, selected_logprobs, replay_output.num_cached_tokens + + +@create_new_process_for_each_test() +@pytest.mark.distributed(num_gpus=2) +def test_kimi_k3_dcp_partial_prefix_reuse( + vllm_runner: type[VllmRunner], + num_gpus_available: int, + tmp_path: Path, +) -> None: + if num_gpus_available < 2: + pytest.skip("Need at least 2 GPUs") + + model_name = _make_tiny_k3_config(tmp_path / "tiny-kimi-k3") + baseline_tokens, baseline_logprobs, baseline_cached = _run_k3_partial_prefix_reuse( + vllm_runner, model_name, dcp_size=1 + ) + dcp_tokens, dcp_logprobs, dcp_cached = _run_k3_partial_prefix_reuse( + vllm_runner, model_name, dcp_size=2 + ) + + assert baseline_cached == 256 + assert dcp_cached == 256 + assert dcp_tokens == baseline_tokens + assert len(dcp_logprobs) == len(baseline_logprobs) == 8 + logprob_drifts = [] + 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/v1/attention/test_mla_backends.py b/tests/v1/attention/test_mla_backends.py index 72773adcb674..cd164250a2fc 100644 --- a/tests/v1/attention/test_mla_backends.py +++ b/tests/v1/attention/test_mla_backends.py @@ -27,6 +27,7 @@ MLAAttention, QueryLenSupport, _DecodeConcatQuantFP8, + build_mla_chunked_context_metadata, ) from vllm.model_executor.layers.quantization.utils.quant_utils import GroupShape from vllm.platforms import current_platform @@ -93,6 +94,48 @@ def test_mla_kv_cache_spec_uses_layer_cache_dtype( assert spec.page_size_bytes == 64 * 656 +@pytest.mark.cpu_test +def test_dcp_chunked_context_accepts_non_virtual_block_aligned_prefix(): + context_lens_cpu = torch.tensor([65, 96], dtype=torch.int32) + prefill_query_start_loc_cpu = torch.tensor([0, 1, 2], dtype=torch.int32) + + metadata = build_mla_chunked_context_metadata( + context_lens_cpu=context_lens_cpu, + prefill_query_start_loc_cpu=prefill_query_start_loc_cpu, + chunked_prefill_workspace=torch.empty(0), + chunked_prefill_workspace_size=128, + block_size=64, + align_chunk_to_block=True, + device=torch.device("cpu"), + dcp_world_size=4, + dcp_local_block_size=1, + dcp_virtual_block_size=4, + ) + + assert metadata is not None + assert metadata.context_lens.tolist() == [65, 96] + assert [chunk.seq_lens.tolist() for chunk in metadata.chunks] == [[65], [96]] + assert [chunk.starts.tolist() for chunk in metadata.chunks] == [[0], [0]] + assert [chunk.local_context_lens_allranks for chunk in metadata.chunks] == [ + [[17, 16, 16, 16]], + [[24, 24, 24, 24]], + ] + assert [chunk.padded_local_seq_lens for chunk in metadata.chunks] == [ + [17], + [24], + ] + assert [chunk.padded_local_cu_seq_lens.tolist() for chunk in metadata.chunks] == [ + [0, 17], + [0, 24], + ] + assert [chunk.cu_seq_lens.tolist() for chunk in metadata.chunks] == [ + [0, 65], + [0, 96], + ] + assert [chunk.num_context_tokens for chunk in metadata.chunks] == [65, 96] + assert [chunk.num_local_context_tokens for chunk in metadata.chunks] == [17, 24] + + # Remove sm100 backends from the list if not using sm100 if not torch.cuda.is_available() or torch.cuda.get_device_properties(0).major < 10: BACKENDS_TO_TEST.remove(AttentionBackendEnum.CUTLASS_MLA) diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py index d237ce9f345b..da5bd5b5a1f3 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_hits.py @@ -4,6 +4,7 @@ "align") models: scheduler chunk splitting, partial tail registration, CoW on partial hits, and same-step deferral.""" +from math import lcm from types import SimpleNamespace from unittest.mock import MagicMock @@ -68,12 +69,62 @@ def test_capable_connector_uses_divergent_partial_hit_lookup(): manager.get_computed_blocks.assert_not_called() -def test_mamba_align_split_partial_tail_schedule(): +def make_full_mamba_manager( + *, + dcp_world_size: int, + hash_block_size: int = 2, + full_block_size: int = 4, + mamba_block_size: int = 4, + num_blocks: int = 32, + use_eagle: bool = False, +): + kv_cache_config = KVCacheConfig( + num_blocks=num_blocks, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec( + ["full"], + FullAttentionSpec( + block_size=full_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.float32, + ), + ), + KVCacheGroupSpec( + ["mamba"], + MambaSpec( + block_size=mamba_block_size, + shapes=(1, 1), + dtypes=(torch.float32,), + mamba_cache_mode="align", + ), + ), + ], + ) + scheduler_block_size = lcm( + full_block_size * dcp_world_size, + mamba_block_size, + ) + return make_kv_cache_manager( + kv_cache_config=kv_cache_config, + max_model_len=8192, + enable_caching=True, + dcp_world_size=dcp_world_size, + scheduler_block_size=scheduler_block_size, + hash_block_size=hash_block_size, + use_eagle=use_eagle, + ) + + +@pytest.mark.parametrize("dcp_world_size", [1, 4]) +def test_mamba_align_split_partial_tail_schedule(dcp_world_size: int): """Chunk ends with partial hits on: block-aligned chunks, one extra stop at the prompt's last hash boundary (registering the partial tail), then the remaining tokens. block=512, hash=32, prompt=10000, budget=8192: 0 -> 8192 -> 9728 -> 9984 -> 10000.""" block_size = 512 + scheduler_block_size = block_size * dcp_world_size hash_block_size = 32 mock = SimpleNamespace( cache_config=SimpleNamespace(block_size=block_size), @@ -81,6 +132,8 @@ def test_mamba_align_split_partial_tail_schedule(): scheduler_config=SimpleNamespace(long_prefill_token_threshold=0), use_eagle=False, hash_block_size=hash_block_size, + dcp_world_size=dcp_world_size, + scheduler_block_size=scheduler_block_size, mamba_partial_cache_hit=True, ) split = Scheduler._mamba_block_aligned_split @@ -109,6 +162,7 @@ def test_mamba_align_split_partial_tail_schedule(): # stops at the next block boundary (10240), later chunk ends re-align. req2 = make_request("1", [0] * 12000, hash_block_size, sha256) req2.num_computed_tokens = 9984 + assert req2.num_computed_tokens % scheduler_block_size != 0 assert split(self=mock, request=req2, num_new_tokens=2016) == 256 req2.num_computed_tokens = 10240 assert split(self=mock, request=req2, num_new_tokens=1000) == 512 @@ -1499,3 +1553,151 @@ def test_hybrid_sliding_window_group_disables_partial_hash_hits(): assert num_computed == mamba_block_size assert len(computed_blocks.blocks[0]) * hash_block_size == num_computed + + +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_hybrid_partial_hash_hit_uses_cow_under_dcp(dcp_world_size: int): + hash_block_size = 2 + physical_block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=dcp_world_size, + hash_block_size=hash_block_size, + full_block_size=physical_block_size, + mamba_block_size=physical_block_size, + ) + assert manager.coordinator.enable_partial_hash_hits + + req0 = make_request("dcp-owner", [0, 0, 1, 1, 2, 2], 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 6, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + partial_hash = req0.block_hashes[2] + partial_full_block = manager.block_pool.get_cached_block(partial_hash, [0]) + partial_mamba_block = manager.block_pool.get_cached_block(partial_hash, [1]) + assert partial_full_block is not None + assert partial_mamba_block is not None + + req1 = make_request("dcp-replay", [0, 0, 1, 1, 2, 2, 3, 3], 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 6 + full_block_size = physical_block_size * dcp_world_size + assert [len(group) for group in computed_blocks.blocks] == [ + (6 + full_block_size - 1) // full_block_size, + 2, + ] + + new_blocks = manager.allocate_slots(req1, 2, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_id = new_blocks.get_block_ids()[0][0] + mamba_new_block_id = new_blocks.get_block_ids()[1][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_full_block[0].block_id, full_new_block_id) in copies + assert ( + KVCacheBlockCopy(partial_mamba_block[0].block_id, mamba_new_block_id) in copies + ) + manager.block_pool.free_blocks(retained) + + +@pytest.mark.parametrize("dcp_world_size", [2, 4]) +def test_dcp_partial_hit_resumes_on_replicated_mamba_snapshot( + dcp_world_size: int, +): + block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=dcp_world_size, + hash_block_size=block_size, + full_block_size=block_size, + mamba_block_size=block_size, + ) + assert manager.coordinator.enable_partial_hash_hits + assert manager.coordinator._cache_hit_alignment_tokens == block_size + assert manager.coordinator.single_type_managers[0].block_size == ( + block_size * dcp_world_size + ) + assert manager.coordinator.single_type_managers[1].block_size == block_size + + prefix = list(range(12)) + req0 = make_request("snapshot-owner", prefix, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 12, num_computed, computed_blocks) is not None + manager.free(req0) + manager.new_step_starts() + + req1 = make_request( + "snapshot-replay", prefix + list(range(12, 16)), block_size, sha256 + ) + + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 12 + assert [len(group) for group in computed_blocks.blocks] == [ + (12 + block_size * dcp_world_size - 1) // (block_size * dcp_world_size), + 3, + ] + partial_full_block = computed_blocks.blocks[0][-1] + new_blocks = manager.allocate_slots(req1, 4, num_computed, computed_blocks) + assert new_blocks is not None + full_new_block_id = new_blocks.get_block_ids()[0][0] + copies, retained = manager.take_kv_cache_block_copies() + assert KVCacheBlockCopy(partial_full_block.block_id, full_new_block_id) in copies + assert all( + copy.src_block_id != computed_blocks.blocks[1][-1].block_id for copy in copies + ) + manager.block_pool.free_blocks(retained) + + +def test_dcp_joint_hit_is_bounded_by_replicated_mamba_snapshots(): + block_size = 4 + manager = make_full_mamba_manager( + dcp_world_size=2, + hash_block_size=block_size, + full_block_size=block_size, + mamba_block_size=block_size, + ) + prefix = list(range(12)) + req0 = make_request("joint-owner", prefix, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 8, num_computed, computed_blocks) is not None + manager.new_step_starts() + + replay_tokens = prefix + list(range(12, 16)) + req1 = make_request("joint-before", replay_tokens, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 8 + assert [len(group) for group in computed_blocks.blocks] == [1, 2] + + req0.num_computed_tokens = 8 + assert manager.allocate_slots(req0, 4) is not None + manager.new_step_starts() + + req2 = make_request("joint-after", replay_tokens, block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req2) + assert num_computed == 12 + assert [len(group) for group in computed_blocks.blocks] == [2, 3] + + +def test_dcp_partial_hit_with_eagle_rewinds_one_hash_unit(): + hash_block_size = 2 + manager = make_full_mamba_manager( + dcp_world_size=2, + hash_block_size=hash_block_size, + full_block_size=4, + mamba_block_size=4, + use_eagle=True, + ) + + req0 = make_request("eagle-owner", [7] * 6, hash_block_size, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req0) + assert manager.allocate_slots(req0, 4, num_computed, computed_blocks) is not None + req0.num_computed_tokens = 4 + manager.new_step_starts() + assert manager.allocate_slots(req0, 2) is not None + req0.num_computed_tokens = 6 + manager.new_step_starts() + + req1 = make_request("eagle-replay", [7] * 6 + [9] * 2, 2, sha256) + computed_blocks, num_computed, _ = manager.get_computed_blocks(req1) + assert num_computed == 4 + assert [len(group) for group in computed_blocks.blocks] == [1, 1] + assert manager.allocate_slots(req1, 4, num_computed, computed_blocks) is not None diff --git a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py index 225d87e86796..bbb7d2626017 100644 --- a/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py +++ b/tests/v1/core/prefix_cache/test_partial_prefix_cache_primitives.py @@ -55,9 +55,9 @@ def cache_full_block_and_partial_tail( token_ids: list[int], *, enable_kv_cache_events: bool = False, + block_size: int = 6, ) -> tuple[BlockPool, Request, list[KVCacheBlock], BlockHash]: hash_block_size = 2 - block_size = 6 kv_cache_group_id = 0 req = make_request("0", token_ids, hash_block_size, sha256) pool = BlockPool( @@ -389,27 +389,34 @@ def test_reset_prefix_cache_clears_partial_entry_metadata(): assert pool.cached_block_hashes_by_block == {} -def test_evict_cached_block_removes_full_hash_and_partial_entry(): - pool, req, blocks, partial_hash_10 = cache_full_block_and_partial_tail( - [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_evict_cached_block_removes_full_hash_and_partial_entry( + dcp_world_size: int, +): + block_size = 6 * dcp_world_size + partial_num_tokens = 2 * block_size - 2 + pool, req, blocks, partial_hash = cache_full_block_and_partial_tail( + list(range(partial_num_tokens)), block_size=block_size ) - full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, 6)[0] + full_hash = BlockHashListWithBlockSize(req.block_hashes, 2, block_size)[0] assert pool.get_cached_block(full_hash, [0]) == [blocks[0]] - assert pool.get_cached_block(partial_hash_10, [0]) == [blocks[1]] + assert pool.get_cached_block(partial_hash, [0]) == [blocks[1]] pool.evict_blocks({blocks[0].block_id, blocks[1].block_id}) assert pool.get_cached_block(full_hash, [0]) is None - assert pool.get_cached_block(partial_hash_10, [0]) is None + assert pool.get_cached_block(partial_hash, [0]) is None assert pool.cached_block_hashes_by_block == {} -def test_partial_block_promotes_to_direct_full_block_hash(): +@pytest.mark.parametrize("dcp_world_size", [1, 2, 4]) +def test_partial_block_promotes_to_direct_full_block_hash(dcp_world_size: int): hash_block_size = 2 - block_size = 6 + block_size = 6 * dcp_world_size kv_cache_group_id = 0 - token_ids = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4] + partial_num_tokens = 2 * block_size - hash_block_size + token_ids = list(range(partial_num_tokens)) req = make_request("0", token_ids, hash_block_size, sha256) pool = BlockPool( num_gpu_blocks=3, @@ -426,27 +433,22 @@ def test_partial_block_promotes_to_direct_full_block_hash(): block_size=block_size, kv_cache_group_id=kv_cache_group_id, ) - partial_hash_10 = boundary_hash(req, hash_block_size, 10) + partial_hash = boundary_hash(req, hash_block_size, partial_num_tokens) assert pool.cache_partial_block( request=req, block=blocks[1], - num_tokens=10, + num_tokens=partial_num_tokens, kv_cache_group_id=kv_cache_group_id, block_size=block_size, ) - assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) == [blocks[1]] + assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) == [blocks[1]] - req.append_output_token_ids([5, 5]) + req.append_output_token_ids(list(range(partial_num_tokens, 2 * block_size))) full_hashes = BlockHashListWithBlockSize( req.block_hashes, hash_block_size, block_size ) promoted_full_hash = full_hashes[1] - # The promoted full-block hash is the fine hash at the 12-token boundary, - # not a concatenation of the fine hashes inside the block. - assert promoted_full_hash == req.block_hashes[12 // hash_block_size - 1] - assert promoted_full_hash != BlockHash( - req.block_hashes[3] + req.block_hashes[4] + req.block_hashes[5] - ) + assert promoted_full_hash == req.block_hashes[2 * block_size // hash_block_size - 1] pool.cache_full_blocks( request=req, @@ -457,4 +459,4 @@ def test_partial_block_promotes_to_direct_full_block_hash(): kv_cache_group_id=kv_cache_group_id, ) assert pool.get_cached_block(promoted_full_hash, [kv_cache_group_id]) == [blocks[1]] - assert pool.get_cached_block(partial_hash_10, [kv_cache_group_id]) is None + assert pool.get_cached_block(partial_hash, [kv_cache_group_id]) is None diff --git a/tests/v1/worker/test_gpu_model_runner_v2.py b/tests/v1/worker/test_gpu_model_runner_v2.py new file mode 100644 index 000000000000..59c11e79c49e --- /dev/null +++ b/tests/v1/worker/test_gpu_model_runner_v2.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from types import SimpleNamespace + +import pytest +import torch + +import vllm.v1.worker.gpu.model_runner as model_runner_module +from vllm.v1.kv_cache_interface import ( + FullAttentionSpec, + KVCacheConfig, + KVCacheGroupSpec, + MambaSpec, +) +from vllm.v1.worker.gpu.block_table import BlockTables +from vllm.v1.worker.gpu.model_runner import GPUModelRunner + + +@pytest.mark.parametrize( + ("mamba_cache_mode", "num_speculative_blocks", "expected"), + [ + pytest.param("align", 0, 65_536, id="align-prefix-cache"), + pytest.param("none", 7, 8, id="no-prefix-cache-with-speculation"), + ], +) +def test_initialize_kv_cache_does_not_dcp_shard_mamba_block_table( + monkeypatch, + mamba_cache_mode: str, + num_speculative_blocks: int, + expected: int, +): + """Mamba/GDN block-table rows index global positions, unlike DCP KV.""" + + max_model_len = 1_048_576 + attention_block_size = 1_536 + mamba_block_size = 16 + dcp_size = 8 + full_attention_spec = FullAttentionSpec( + block_size=attention_block_size, + num_kv_heads=1, + head_size=1, + dtype=torch.bfloat16, + ) + mamba_spec = MambaSpec( + shapes=((1,),), + dtypes=(torch.bfloat16,), + block_size=mamba_block_size, + mamba_cache_mode=mamba_cache_mode, + num_speculative_blocks=num_speculative_blocks, + ) + kv_cache_config = KVCacheConfig( + num_blocks=1, + kv_cache_tensors=[], + kv_cache_groups=[ + KVCacheGroupSpec(["attention"], full_attention_spec), + KVCacheGroupSpec(["kda"], mamba_spec), + ], + ) + vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(decode_context_parallel_size=dcp_size), + cache_config=SimpleNamespace(mamba_cache_mode=mamba_cache_mode), + ) + runner = SimpleNamespace( + max_model_len=max_model_len, + is_encoder_decoder=False, + vllm_config=vllm_config, + ) + + class _CapturedWidths(Exception): + pass + + captured: list[int] = [] + + def capture_width(max_num_blocks: int, *_args, **_kwargs) -> int: + captured.append(max_num_blocks) + if len(captured) == 2: + raise _CapturedWidths + return max_num_blocks + + monkeypatch.setattr(model_runner_module, "get_block_table_width", capture_width) + + with pytest.raises(_CapturedWidths): + GPUModelRunner.initialize_kv_cache(runner, kv_cache_config) + + # Attention KV is local to one of eight DCP ranks; KDA state is replicated + # and therefore needs one table entry for every global 16-token page. + assert captured == [86, expected] + + +def test_append_block_ids_rejects_write_past_row_capacity(): + """Reject an oversized staged write before it can corrupt the next row.""" + + class _BlockTable: + gpu = torch.empty((2, 4), dtype=torch.int32) + + def stage_write(self, *_args): + pytest.fail("an oversized write must not be staged") + + block_tables = BlockTables.__new__(BlockTables) + block_tables.num_kv_cache_groups = 1 + block_tables.blocks_per_kv_block = [1] + block_tables.block_tables = [_BlockTable()] + block_tables.num_blocks = SimpleNamespace( + np=torch.tensor([[0, 3]], dtype=torch.int32) + ) + + with pytest.raises( + RuntimeError, + match=r"request 1, group 0 exceeds row capacity \(5 > 4\)", + ): + block_tables.append_block_ids( + req_index=1, + new_block_ids=([4, 5],), + overwrite=False, + ) + + assert block_tables.num_blocks.np[0, 1] == 3 diff --git a/vllm/v1/core/kv_cache_coordinator.py b/vllm/v1/core/kv_cache_coordinator.py index 0829ec0d2be8..767ea0033483 100644 --- a/vllm/v1/core/kv_cache_coordinator.py +++ b/vllm/v1/core/kv_cache_coordinator.py @@ -619,15 +619,22 @@ def __init__( "full-attention and Mamba groups, got: " f"{type(g.kv_cache_spec).__name__}." ) - # Fine-grained hash hits require Mamba "align", no context - # parallelism, and compatible cache managers in every group. + # Fine-grained hash hits require Mamba "align" and compatible cache + # managers in every group. TP needs hashing finer than the Mamba block; + # DCP accepts equality because it scales the effective full-attention + # block instead. has_partial_mamba_group = any( isinstance(g.kv_cache_spec, MambaSpec) and g.kv_cache_spec.mamba_cache_mode == "align" - and g.kv_cache_spec.block_size > hash_block_size + and ( + (dcp_world_size == 1 and g.kv_cache_spec.block_size > hash_block_size) + or ( + dcp_world_size > 1 and g.kv_cache_spec.block_size >= hash_block_size + ) + ) for g in kv_cache_config.kv_cache_groups ) - self.enable_partial_hash_hits = dcp_world_size == 1 and has_partial_mamba_group + self.enable_partial_hash_hits = has_partial_mamba_group if self.enable_partial_hash_hits: unsupported_partial_hit_managers = { type(manager).__name__ diff --git a/vllm/v1/worker/gpu/block_table.py b/vllm/v1/worker/gpu/block_table.py index 90fc104dc29f..fbc4f383cb35 100644 --- a/vllm/v1/worker/gpu/block_table.py +++ b/vllm/v1/worker/gpu/block_table.py @@ -116,8 +116,15 @@ def append_block_ids( bpk = self.blocks_per_kv_block[i] if bpk > 1: block_ids = [b * bpk + k for b in block_ids for k in range(bpk)] + end = start + len(block_ids) + row_capacity = self.block_tables[i].gpu.shape[1] + if end > row_capacity: + raise RuntimeError( + f"Block table write for request {req_index}, group {i} exceeds " + f"row capacity ({end} > {row_capacity})" + ) self.block_tables[i].stage_write(req_index, start, block_ids) - self.num_blocks.np[i, req_index] = start + len(block_ids) + self.num_blocks.np[i, req_index] = end def apply_staged_writes(self) -> None: if self.num_kv_cache_groups == 0: diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 77ae59f0ee38..4f65769c9580 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -60,7 +60,6 @@ ) from vllm.sequence import IntermediateTensors from vllm.tasks import SupportedTask -from vllm.utils.math_utils import cdiv from vllm.utils.mem_utils import DeviceMemoryProfiler, format_gib from vllm.utils.torch_utils import STR_DTYPE_TO_TORCH_DTYPE from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput @@ -517,17 +516,14 @@ def initialize_kv_cache(self, kv_cache_config: KVCacheConfig) -> None: for kv_cache_group in kv_cache_config.kv_cache_groups: spec = kv_cache_group.kv_cache_spec block_sizes.append(spec.block_size) - # When using DCP, each request's KV cache is sharded among different ranks. - # As a result, one block on the current rank covers `block_size * cp_size` - # tokens in the full, global (unsharded) sequence. - max_num_blocks = cdiv( - block_table_max_model_len, spec.block_size * self.dcp_size + # Let each cache type account for CP. Attention KV is DCP-sharded, + # while Mamba/GDN recurrent state is replicated across DCP ranks. + max_num_blocks = spec.max_num_blocks_per_req( + self.vllm_config, block_table_max_model_len ) - # For Mamba/Hybrid Model, KVCaches need extra blocks for speculative tokens + # Preserve each cache type's alignment requirements after applying + # its topology-aware block-table width. if isinstance(spec, MambaSpec): - max_num_blocks = ( - max_num_blocks if self.cache_config.enable_prefix_caching else 1 - ) + spec.num_speculative_blocks max_num_blocks = get_block_table_width( max_num_blocks, spec.block_size, token_alignment=None ) From c296851a7d173fa89d2eefbca0243be42ae9b5e0 Mon Sep 17 00:00:00 2001 From: Bo Li <22713281+bobboli@users.noreply.github.com> Date: Tue, 18 Aug 2026 08:18:14 +0800 Subject: [PATCH 05/15] [MoE] Refine FlashInfer one-sided All2All integration (#51924) Signed-off-by: Bo Li <22713281+bobboli@users.noreply.github.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com> --- docs/design/moe_kernel_features.md | 2 +- tests/distributed/test_mnnvl_alltoall.py | 30 ++++---- tests/kernels/moe/test_moe_layer.py | 2 +- vllm/config/parallel.py | 1 + .../device_communicators/all2all.py | 12 +-- .../layers/fused_moe/all2all_utils.py | 74 +++++++++++++------ .../fused_moe/experts/trtllm_fp8_moe.py | 28 ++++++- .../flashinfer_nvlink_one_sided.py | 48 ++++++------ 8 files changed, 129 insertions(+), 68 deletions(-) diff --git a/docs/design/moe_kernel_features.md b/docs/design/moe_kernel_features.md index bdf407510294..0b1a9151ec69 100644 --- a/docs/design/moe_kernel_features.md +++ b/docs/design/moe_kernel_features.md @@ -36,7 +36,7 @@ th { | deepep_high_throughput | standard | fp8 | G(128),A,T2 | Y | Y | [`DeepEPHTPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ht.DeepEPHTPrepareAndFinalize] | | deepep_low_latency | batched | fp8 | G(128),A,T3 | Y | Y | [`DeepEPLLPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.deepep_ll.DeepEPLLPrepareAndFinalize] | | flashinfer_nvlink_two_sided | standard | nvfp4,fp8 | G,A,T | N | N | [`FlashInferNVLinkTwoSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_two_sided.FlashInferNVLinkTwoSidedPrepareAndFinalize] | -| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | +| flashinfer_nvlink_one_sided | standard | nvfp4,bf16,mxfp8,fp8 | G,A,T | N | N | [`FlashInferNVLinkOneSidedPrepareAndFinalize`][vllm.model_executor.layers.fused_moe.prepare_finalize.flashinfer_nvlink_one_sided.FlashInferNVLinkOneSidedPrepareAndFinalize] | !!! info "Table key" 1. All types: mxfp4, nvfp4, int4, int8, fp8 diff --git a/tests/distributed/test_mnnvl_alltoall.py b/tests/distributed/test_mnnvl_alltoall.py index f9b76ea1b030..636f7cb525f5 100644 --- a/tests/distributed/test_mnnvl_alltoall.py +++ b/tests/distributed/test_mnnvl_alltoall.py @@ -342,6 +342,8 @@ def _one_sided_lifecycle_worker(rank, world_size): top_k=2, num_experts=world_size * 8, hidden_size=4096, + x_bytes_per_token=4096 * 2, + x_sf_bytes_per_token=0, ) # Initialize @@ -409,12 +411,12 @@ def _one_sided_workspace_grow_worker(rank, world_size): hidden_size=4096, ) nvfp4_kwargs = dict( - dispatch_dtype_bytes_per_elem=0, - dispatch_scale_bytes_per_token=base_kwargs["hidden_size"] // 16, + x_bytes_per_token=base_kwargs["hidden_size"] // 2, + x_sf_bytes_per_token=base_kwargs["hidden_size"] // 16, ) bf16_kwargs = dict( - dispatch_dtype_bytes_per_elem=2, - dispatch_scale_bytes_per_token=0, + x_bytes_per_token=base_kwargs["hidden_size"] * 2, + x_sf_bytes_per_token=0, ) # First init: NVFP4-like (hidden_bytes = hidden // 2 + hidden // 16). @@ -780,11 +782,12 @@ def _one_sided_data_worker(rank, world_size): top_k=experts_per_token, num_experts=num_experts, hidden_size=hidden_size, - # Account for the fp8 block-scale payload (a1q_scale: hidden//16 bytes + x_bytes_per_token=hidden_size // 2, + # Account for the fp8 block-scale payload (x_sf: hidden//16 bytes # per token) that is dispatched alongside the nvfp4 hidden states. # Without this the dispatch region is under-reserved and the combine # payload overflows the per-rank workspace. - dispatch_scale_bytes_per_token=hidden_size // 16, + x_sf_bytes_per_token=hidden_size // 16, ) assert manager.initialized assert manager.moe_alltoall is not None @@ -798,17 +801,17 @@ def _one_sided_data_worker(rank, world_size): # Create test data with raw tensors matching the nvfp4 payload # sizes the workspace was allocated for: - # a1q: (tokens, hidden_size // 2) — nvfp4 hidden states - # a1q_scale: (tokens, hidden_size // 16) — fp8 scaling factors + # x: (tokens, hidden_size // 2) — nvfp4 hidden states + # x_sf: (tokens, hidden_size // 16) — fp8 scaling factors torch.manual_seed(rank + 42) - a1q = torch.randint( + x = torch.randint( 0, 256, (tokens_per_rank, hidden_size // 2), device=device, dtype=torch.uint8, ) - a1q_scale = torch.randint( + x_sf = torch.randint( 0, 256, (tokens_per_rank, hidden_size // 16), @@ -830,15 +833,16 @@ def _one_sided_data_worker(rank, world_size): ) # --- One-sided dispatch --- - payloads = [a1q, a1q_scale, topk_ids, topk_weights] + payloads = [x, x_sf, topk_ids, topk_weights] recv_payloads = manager.moe_alltoall.dispatch( token_selected_experts=topk_ids, input_payloads=payloads, runtime_max_tokens_per_rank=runtime_max_tokens, ) assert len(recv_payloads) == 4 - recv_a1q, recv_scale, recv_ids, recv_weights = recv_payloads - assert recv_a1q.numel() > 0 + recv_x, recv_x_sf, recv_ids, recv_weights = recv_payloads + assert recv_x.numel() > 0 + assert recv_x_sf.numel() > 0 assert recv_ids.numel() > 0 # --- Round-trip exact verification --- diff --git a/tests/kernels/moe/test_moe_layer.py b/tests/kernels/moe/test_moe_layer.py index 1942ce96bcca..74554944a314 100644 --- a/tests/kernels/moe/test_moe_layer.py +++ b/tests/kernels/moe/test_moe_layer.py @@ -130,7 +130,7 @@ "mori_high_throughput": {None, "fp8", "modelopt_fp8"}, "mori_low_latency": {None, "fp8", "modelopt_fp8"}, "flashinfer_nvlink_two_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 - "flashinfer_nvlink_one_sided": {None, "modelopt_fp4"}, # noqa: E501 + "flashinfer_nvlink_one_sided": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "deepep_low_latency": {None, "fp8_blocked", "modelopt_fp4"}, # noqa: E501 "deepep_high_throughput": {None, "fp8_blocked", "modelopt_fp8", "modelopt_fp4"}, # noqa: E501 "nixl_ep": {None, "fp8_blocked", "modelopt_fp8"}, diff --git a/vllm/config/parallel.py b/vllm/config/parallel.py index 83c4e6fe53d7..eb346ed4a34a 100644 --- a/vllm/config/parallel.py +++ b/vllm/config/parallel.py @@ -677,6 +677,7 @@ def use_sequence_parallel_moe(self) -> bool: "allgather_reducescatter", "deepep_high_throughput", "deepep_low_latency", + "flashinfer_nvlink_one_sided", "mori_high_throughput", "mori_low_latency", "nixl_ep", diff --git a/vllm/distributed/device_communicators/all2all.py b/vllm/distributed/device_communicators/all2all.py index 5abe7568a292..657ad9fd0e13 100644 --- a/vllm/distributed/device_communicators/all2all.py +++ b/vllm/distributed/device_communicators/all2all.py @@ -734,17 +734,13 @@ def initialize( top_k: int, num_experts: int, hidden_size: int, - dispatch_dtype_bytes_per_elem: int = 0, - dispatch_scale_bytes_per_token: int = 0, + x_bytes_per_token: int, + x_sf_bytes_per_token: int, ): """Initialize (or grow) the MoeAlltoAll workspace.""" - if dispatch_dtype_bytes_per_elem == 0: - hidden_bytes = hidden_size // 2 - else: - hidden_bytes = hidden_size * dispatch_dtype_bytes_per_elem total_dispatch_payload_size_per_token = ( - hidden_bytes - + dispatch_scale_bytes_per_token + x_bytes_per_token + + x_sf_bytes_per_token + top_k * 4 # int32 topks ids + top_k * 4 # float32 topk weights ) diff --git a/vllm/model_executor/layers/fused_moe/all2all_utils.py b/vllm/model_executor/layers/fused_moe/all2all_utils.py index 29905927c106..7861345453a0 100644 --- a/vllm/model_executor/layers/fused_moe/all2all_utils.py +++ b/vllm/model_executor/layers/fused_moe/all2all_utils.py @@ -1,6 +1,7 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from dataclasses import dataclass from typing import Any import torch @@ -37,6 +38,52 @@ has_nixl_ep, ) + +@dataclass(frozen=True) +class FlashInferOneSidedDispatchLayout: + x_bytes_per_token: int + x_sf_bytes_per_token: int + + +def flashinfer_one_sided_dispatch_layout( + hidden_dim: int, quant_config: FusedMoEQuantConfig +) -> FlashInferOneSidedDispatchLayout: + """Return the one-sided activation payload layout.""" + if quant_config.quant_dtype is None: + return FlashInferOneSidedDispatchLayout(hidden_dim * 2, 0) + if quant_config.quant_dtype == "nvfp4": + scale_elems = hidden_dim // 16 + return FlashInferOneSidedDispatchLayout(hidden_dim // 2, scale_elems) + if quant_config.quant_dtype == "mxfp8": + align = quant_config.mx_alignment + padded_k = ( + ((hidden_dim + align - 1) // align) * align if align > 0 else hidden_dim + ) + scale_elems = padded_k // 32 + return FlashInferOneSidedDispatchLayout(hidden_dim, scale_elems) + if ( + quant_config.use_fp8_w8a8 + and quant_config.quant_dtype == current_platform.fp8_dtype() + and quant_config.block_shape == [128, 128] + ): + if hidden_dim % 128 != 0: + raise NotImplementedError( + "flashinfer_nvlink_one_sided DeepSeek Blockwise FP8 dispatch " + f"requires hidden_dim divisible by 128; got {hidden_dim}" + ) + scale_elems = hidden_dim // 128 + scale_bytes = scale_elems * torch.float32.itemsize + return FlashInferOneSidedDispatchLayout(hidden_dim, scale_bytes) + raise NotImplementedError( + "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " + "DeepSeek Blockwise FP8 (E4M3 with FP32 1x128 scales), and bf16 " + "(quant_dtype=None) today; got " + f"quant_dtype={quant_config.quant_dtype!r}, " + f"use_fp8_w8a8={quant_config.use_fp8_w8a8!r}, " + f"block_shape={quant_config.block_shape!r}" + ) + + logger = init_logger(__name__) if current_platform.is_cuda_alike(): @@ -282,34 +329,17 @@ def maybe_make_prepare_finalize( max_num_tokens = ( get_current_vllm_config().scheduler_config.max_num_batched_tokens ) - if quant_config.quant_dtype is None: - dispatch_dtype_bytes_per_elem = 2 - dispatch_scale_bytes_per_token = 0 - elif quant_config.quant_dtype == "nvfp4": - dispatch_dtype_bytes_per_elem = 0 - dispatch_scale_bytes_per_token = moe.hidden_dim // 16 - elif quant_config.quant_dtype == "mxfp8": - dispatch_dtype_bytes_per_elem = 1 - align = quant_config.mx_alignment - if align > 0: - padded_k = ((moe.hidden_dim + align - 1) // align) * align - else: - padded_k = moe.hidden_dim - dispatch_scale_bytes_per_token = padded_k // 32 - else: - raise NotImplementedError( - "flashinfer_nvlink_one_sided dispatch supports nvfp4, mxfp8, " - "and bf16 (quant_dtype=None) today; got " - f"quant_dtype={quant_config.quant_dtype!r}" - ) + dispatch_layout = flashinfer_one_sided_dispatch_layout( + moe.hidden_dim, quant_config + ) prepare_finalize = FlashInferNVLinkOneSidedPrepareAndFinalize( max_num_tokens=max_num_tokens, top_k=moe.experts_per_token, num_experts=moe.num_experts, hidden_size=moe.hidden_dim, num_dispatchers=all2all_manager.world_size, - dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, - dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token, + x_bytes_per_token=dispatch_layout.x_bytes_per_token, + x_sf_bytes_per_token=dispatch_layout.x_sf_bytes_per_token, ) elif moe.use_ag_rs_all2all_kernels and allow_new_interface: diff --git a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py index d32877eb9a95..00ab06fe952b 100644 --- a/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py +++ b/vllm/model_executor/layers/fused_moe/experts/trtllm_fp8_moe.py @@ -36,6 +36,28 @@ logger = init_logger(__name__) +def prepare_deepseek_fp8_x_sf(x: torch.Tensor, x_sf: torch.Tensor) -> torch.Tensor: + """Validate DeepSeek Blockwise FP8 tensors and return TRTLLM layout.""" + if x.dtype != current_platform.fp8_dtype(): + raise ValueError( + f"DeepSeekFp8 activations must use the platform E4M3 dtype; got {x.dtype}" + ) + if x.ndim != 2 or x.shape[1] % 128 != 0: + raise ValueError( + "DeepSeekFp8 activations must be [M,K] with K divisible by 128; " + f"got {tuple(x.shape)}" + ) + expected_shape = (x.shape[0], x.shape[1] // 128) + if x_sf.dtype != torch.float32 or tuple(x_sf.shape) != expected_shape: + raise ValueError( + "DeepSeekFp8 activation scales must be FP32 [M,K/128]; " + f"expected {expected_shape}, got dtype={x_sf.dtype}, " + f"shape={tuple(x_sf.shape)}" + ) + # FlashInfer TRTLLM-gen consumes [K/128,M] for DeepSeekFp8/BlockMajorK. + return x_sf.t().contiguous() + + class TrtLlmFp8ExpertsBase: """ Fp8 TRTLLM-Gen MoE kernels. Shared base for modular and monolithic @@ -239,6 +261,10 @@ def apply( # Pack topk ids and weights into format expected by the kernel. packed_topk_ids = trtllm_moe_pack_topk_ids_weights(topk_ids, topk_weights) + if a1q_scale is None: + raise RuntimeError( + "TRT-LLM FP8 experts require precomputed activation scales" + ) assert a1q_scale is not None is_mxfp8 = self.quant_config.block_shape == [1, 32] @@ -251,7 +277,7 @@ def apply( fp8_quant_type = Fp8QuantizationType.DeepSeekFp8 use_shuffled_weight = True weight_layout = WeightLayout.BlockMajorK - hidden_states_scale = a1q_scale.t().contiguous() + hidden_states_scale = prepare_deepseek_fp8_x_sf(hidden_states, a1q_scale) flashinfer.fused_moe.trtllm_fp8_block_scale_routed_moe( topk_ids=packed_topk_ids, diff --git a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py index 74341e7681f8..0b57afb6062c 100644 --- a/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py +++ b/vllm/model_executor/layers/fused_moe/prepare_finalize/flashinfer_nvlink_one_sided.py @@ -30,9 +30,9 @@ def __init__( top_k: int, num_experts: int, hidden_size: int, + x_bytes_per_token: int, + x_sf_bytes_per_token: int, num_dispatchers: int = 1, - dispatch_dtype_bytes_per_elem: int = 0, - dispatch_scale_bytes_per_token: int = 0, ): super().__init__() self.max_num_tokens = max_num_tokens @@ -40,7 +40,6 @@ def __init__( self.num_experts = num_experts self.hidden_size = hidden_size self.num_dispatchers_ = num_dispatchers - self.scale_elems_per_token = dispatch_scale_bytes_per_token device_communicator = get_ep_group().device_communicator assert device_communicator is not None @@ -52,8 +51,8 @@ def __init__( top_k=self.top_k, num_experts=self.num_experts, hidden_size=self.hidden_size, - dispatch_dtype_bytes_per_elem=dispatch_dtype_bytes_per_elem, - dispatch_scale_bytes_per_token=dispatch_scale_bytes_per_token, + x_bytes_per_token=x_bytes_per_token, + x_sf_bytes_per_token=x_sf_bytes_per_token, ) @property @@ -98,9 +97,9 @@ def prepare( ) if defer_input_quant: - a1q, a1q_scale = a1, None + dispatch_x, dispatch_x_sf = a1, None else: - a1q, a1q_scale = moe_kernel_quantize_input( + dispatch_x, dispatch_x_sf = moe_kernel_quantize_input( a1, quant_config.a1_gscale, quant_config.quant_dtype, @@ -110,10 +109,9 @@ def prepare( mx_alignment=quant_config.mx_alignment, ) - payloads = [] - payloads.append(a1q) - if a1q_scale is not None: - payloads.append(a1q_scale) + payloads = [dispatch_x] + if dispatch_x_sf is not None: + payloads.append(dispatch_x_sf) topk_ids_payload_index = len(payloads) payloads.append(topk_ids) payloads.append(topk_weights) @@ -126,23 +124,29 @@ def prepare( invalid_token_expert_id=-1, # Follow TRTLLM Pattern expert_id_payload_index=topk_ids_payload_index, ) - if a1q_scale is not None: - a1q_recv, a1q_scale_recv, topk_ids_recv, topk_weights_recv = recv_payloads + if dispatch_x_sf is not None: + recv_x, recv_x_sf, topk_ids_recv, topk_weights_recv = recv_payloads + x_sf_width = recv_x_sf.shape[-1] # Apply scale interleaving only for CUTLASS (not TRT-LLM) if quant_config.quant_dtype == "nvfp4" and quant_config.is_scale_swizzled: - a1q_scale_recv = a1q_scale_recv.view(-1, a1q_scale_recv.shape[-1]) - a1q_scale_recv = a1q_scale_recv.view(torch.uint8) - a1q_scale_recv = nvfp4_block_scale_interleave(a1q_scale_recv) - assert self.scale_elems_per_token > 0 - a1q_scale_recv = a1q_scale_recv.view(-1, self.scale_elems_per_token) + recv_x_sf = recv_x_sf.view(-1, x_sf_width) + recv_x_sf = recv_x_sf.view(torch.uint8) + recv_x_sf = nvfp4_block_scale_interleave(recv_x_sf) + recv_x_sf = recv_x_sf.view(-1, x_sf_width) else: - a1q_recv, topk_ids_recv, topk_weights_recv = recv_payloads - a1q_scale_recv = None - a1q_recv = a1q_recv.view(-1, a1q_recv.shape[-1]) + recv_x, topk_ids_recv, topk_weights_recv = recv_payloads + recv_x_sf = None + recv_x = recv_x.view(-1, recv_x.shape[-1]) topk_ids_recv = topk_ids_recv.view(-1, topk_ids_recv.shape[-1]) topk_weights_recv = topk_weights_recv.view(-1, topk_weights_recv.shape[-1]) - return a1q_recv, a1q_scale_recv, None, topk_ids_recv, topk_weights_recv + return ( + recv_x, + recv_x_sf, + None, + topk_ids_recv, + topk_weights_recv, + ) def finalize( self, From cdb8545a91be16f2b80234e0801458ad2a18cf1b Mon Sep 17 00:00:00 2001 From: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> Date: Tue, 18 Aug 2026 09:17:40 +0800 Subject: [PATCH 06/15] [Kernel][Perf] Support Qwen head ratios in fused GDN MTP (#52539) Signed-off-by: kiroxu <148877251+BabyDrangoner@users.noreply.github.com> --- .../gdn/fused_gdn_decode_kernel.cu | 72 ++++++++++++------- tests/kernels/mamba/test_gdn_fused_mtp.py | 42 +++++++++++ tests/kernels/test_fused_gdn_post_conv.py | 59 ++++++++++++--- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 3 +- 4 files changed, 142 insertions(+), 34 deletions(-) diff --git a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu index 31a16fcf9965..1166b41868e0 100644 --- a/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu +++ b/csrc/libtorch_stable/gdn/fused_gdn_decode_kernel.cu @@ -146,7 +146,7 @@ __device__ __forceinline__ Sum2 warp_reduce_sum_pair(float x, float y) { return {x, y}; } -template +template __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( const __nv_bfloat16* __restrict__ mixed_qkv, const __nv_bfloat16* __restrict__ a, const __nv_bfloat16* __restrict__ b, @@ -186,7 +186,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( return; } - const int key_head = value_head / 8; + const int key_head = value_head / ValueHeadsPerKeyHead; __shared__ StateT shared_state[2][kChunkV][kDimK]; __shared__ float shared_q[kMaxMtpTokens][kDimK]; __shared__ float shared_k[kMaxMtpTokens][kDimK]; @@ -370,7 +370,7 @@ __global__ __launch_bounds__(kThreads, 2) void gdn_decode_post_conv_mtp_kernel( } } -template +template void launch_gdn_decode_post_conv_mtp( torch::stable::Tensor const& mixed_qkv, torch::stable::Tensor const& a_log, torch::stable::Tensor const& dt_bias, @@ -395,17 +395,19 @@ void launch_gdn_decode_post_conv_mtp( get_current_cuda_stream(mixed_qkv.get_device_index()); const int num_requests = static_cast(state_indices.size(0)); const dim3 grid(num_requests, num_value_heads); - gdn_decode_post_conv_mtp_kernel<<>>( - static_cast(mixed_qkv.data_ptr()), a, b, - static_cast(a_log.data_ptr()), dt_bias.data_ptr(), - static_cast(state_indices.data_ptr()), - static_cast(cu_seqlens.data_ptr()), - static_cast(num_accepted_tokens.data_ptr()), - static_cast(state.data_ptr()), output_gate, - norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), - num_key_heads, num_value_heads, static_cast(state_indices.size(1)), - dt_bias_type, norm_weight.scalar_type() == ScalarType::BFloat16, - static_cast(scale), static_cast(norm_eps), strides); + gdn_decode_post_conv_mtp_kernel + <<>>( + static_cast(mixed_qkv.data_ptr()), a, b, + static_cast(a_log.data_ptr()), dt_bias.data_ptr(), + static_cast(state_indices.data_ptr()), + static_cast(cu_seqlens.data_ptr()), + static_cast(num_accepted_tokens.data_ptr()), + static_cast(state.data_ptr()), output_gate, + norm_weight.data_ptr(), static_cast<__nv_bfloat16*>(out.data_ptr()), + num_key_heads, num_value_heads, + static_cast(state_indices.size(1)), dt_bias_type, + norm_weight.scalar_type() == ScalarType::BFloat16, + static_cast(scale), static_cast(norm_eps), strides); const cudaError_t error = cudaGetLastError(); STD_TORCH_CHECK(error == cudaSuccess, "GDN decode MTP post-conv kernel launch failed: ", @@ -479,8 +481,12 @@ void fused_gdn_decode_post_conv_mtp( STD_TORCH_CHECK(key_width > 0 && key_width % (2 * kDimK) == 0, "mixed_qkv width is inconsistent with state"); const int num_key_heads = static_cast(key_width / (2 * kDimK)); - STD_TORCH_CHECK(num_value_heads == 8 * num_key_heads, - "GDN decode MTP fusion requires HV/H=8"); + const int value_heads_per_key_head = num_value_heads / num_key_heads; + STD_TORCH_CHECK( + num_value_heads % num_key_heads == 0 && + ((value_heads_per_key_head >= 1 && value_heads_per_key_head <= 4) || + value_heads_per_key_head == 8), + "GDN decode MTP fusion requires HV/H in {1, 2, 3, 4, 8}"); STD_TORCH_CHECK(state_indices.dim() == 2 && state_indices.size(0) > 0 && state_indices.size(1) > 0 && @@ -541,17 +547,35 @@ void fused_gdn_decode_post_conv_mtp( const auto* b_ptr = static_cast(b.data_ptr()); const auto* output_gate_ptr = static_cast(output_gate.data_ptr()); - if (state_scalar_type == ScalarType::Float) { - launch_gdn_decode_post_conv_mtp( - mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, - num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, - output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, - strides); - } else { - launch_gdn_decode_post_conv_mtp<__nv_bfloat16>( + const auto launch = [&]() { + launch_gdn_decode_post_conv_mtp( mixed_qkv, a_log, dt_bias, state_indices, cu_seqlens, num_accepted_tokens, state, norm_weight, out, a_ptr, b_ptr, output_gate_ptr, num_key_heads, num_value_heads, scale, norm_eps, strides); + }; + const auto dispatch_state_type = [&]() { + if (state_scalar_type == ScalarType::Float) { + launch.template operator()(); + } else { + launch.template operator()<__nv_bfloat16, ValueHeadsPerKeyHead>(); + } + }; + switch (value_heads_per_key_head) { + case 1: + dispatch_state_type.template operator()<1>(); + break; + case 2: + dispatch_state_type.template operator()<2>(); + break; + case 3: + dispatch_state_type.template operator()<3>(); + break; + case 4: + dispatch_state_type.template operator()<4>(); + break; + default: + dispatch_state_type.template operator()<8>(); + break; } } diff --git a/tests/kernels/mamba/test_gdn_fused_mtp.py b/tests/kernels/mamba/test_gdn_fused_mtp.py index 97ed4c48e787..0d78f8ceeba0 100644 --- a/tests/kernels/mamba/test_gdn_fused_mtp.py +++ b/tests/kernels/mamba/test_gdn_fused_mtp.py @@ -5,6 +5,7 @@ from __future__ import annotations import types +from typing import cast from unittest.mock import patch import pytest @@ -37,6 +38,7 @@ ) from vllm.utils.torch_utils import _encode_layer_name # noqa: E402 from vllm.v1.attention.backends.gdn_attn import ( # noqa: E402 + GDNAttentionMetadata, GDNAttentionMetadataBuilder, ) from vllm.v1.kv_cache_interface import MambaSpec # noqa: E402 @@ -139,6 +141,46 @@ def _build_layer( return layer +@pytest.mark.parametrize( + "num_v_heads,expected", + [ + pytest.param(2, True, id="ratio1"), + pytest.param(4, True, id="ratio2"), + pytest.param(6, True, id="ratio3"), + pytest.param(8, True, id="ratio4"), + pytest.param(16, True, id="ratio8"), + pytest.param(5, False, id="non-integral-ratio"), + pytest.param(10, False, id="unsupported-ratio5"), + ], +) +def test_fused_mtp_head_ratio_guard(num_v_heads: int, expected: bool) -> None: + if not hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp"): + pytest.skip("fused GDN decode MTP op is not built") + + layer = types.SimpleNamespace( + num_k_heads=2, + num_v_heads=num_v_heads, + kv_cache=(None, torch.empty(1, dtype=torch.float32, device="cuda")), + gdn_decode_kernel="cuda", + ) + attn_metadata = types.SimpleNamespace( + spec_state_indices_tensor=torch.ones( + 1, SPEC_TOKENS, dtype=torch.int32, device="cuda" + ), + spec_sequence_masks=object(), + num_decodes=0, + num_spec_decodes=1, + ) + + assert ( + QwenGatedDeltaNetAttention._can_use_fused_gdn_mtp_decode( + cast(QwenGatedDeltaNetAttention, layer), + cast(GDNAttentionMetadata, attn_metadata), + ) + is expected + ) + + @torch.inference_mode() def test_fused_forward_uses_packed_entrypoint() -> None: """Fused mode keeps projected QKVZ and BA packed through the model op.""" diff --git a/tests/kernels/test_fused_gdn_post_conv.py b/tests/kernels/test_fused_gdn_post_conv.py index b6dacf2e73f2..7cdeb55b9b57 100644 --- a/tests/kernels/test_fused_gdn_post_conv.py +++ b/tests/kernels/test_fused_gdn_post_conv.py @@ -217,18 +217,59 @@ def test_fused_post_conv_l0(): @pytest.mark.parametrize( - "tp_size,query_lengths,state_dtype,norm_dtype", + "head_ratio,tp_size,query_lengths,state_dtype,norm_dtype", [ - pytest.param(16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), - pytest.param(4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), - pytest.param(16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), - pytest.param(4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), - pytest.param(16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), - pytest.param(4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + pytest.param(8, 16, (4, 4), torch.bfloat16, torch.bfloat16, id="tp16-bf16"), + pytest.param(8, 4, (4, 4), torch.float32, torch.float32, id="tp4-fp32"), + pytest.param(8, 16, (4, 2, 0), torch.bfloat16, torch.float32, id="tp16-ragged"), + pytest.param(8, 4, (4, 2, 0), torch.float32, torch.bfloat16, id="tp4-ragged"), + pytest.param(8, 16, (8,), torch.float32, torch.bfloat16, id="tp16-max"), + pytest.param(8, 4, (8,), torch.bfloat16, torch.float32, id="tp4-max"), + pytest.param( + 1, + 1, + (4, 4), + torch.float32, + torch.bfloat16, + id="ratio1-tp1-fp32", + ), + pytest.param( + 2, + 1, + (4, 2, 0), + torch.bfloat16, + torch.bfloat16, + id="ratio2-tp1-ragged-bf16", + ), + pytest.param( + 2, + 4, + (8,), + torch.float32, + torch.float32, + id="ratio2-tp4-max-fp32", + ), + pytest.param( + 3, + 2, + (4, 2, 0), + torch.float32, + torch.float32, + id="ratio3-tp2-ragged-fp32", + ), + pytest.param( + 4, + 4, + (8,), + torch.float32, + torch.bfloat16, + id="ratio4-tp4-max-fp32", + ), ], ) @torch.inference_mode() -def test_fused_gdn_decode_post_conv_mtp_ratio8( +def test_fused_gdn_decode_post_conv_mtp_head_ratios( + head_ratio: int, tp_size: int, query_lengths: tuple[int, ...], state_dtype: torch.dtype, @@ -242,7 +283,7 @@ def test_fused_gdn_decode_post_conv_mtp_ratio8( torch.manual_seed(0) device = "cuda" H = 16 // tp_size - HV = 128 // tp_size + HV = head_ratio * H K = V = 128 num_reqs = len(query_lengths) state_width = max(query_lengths) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index 27dc4d90d7e0..d83b076509ee 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1805,7 +1805,8 @@ def _can_use_fused_gdn_mtp_decode( and attn_metadata.num_spec_decodes > 0 and self.kv_cache[1].dtype in FUSED_GDN_STATE_DTYPES and self.gdn_decode_kernel == "cuda" - and self.num_v_heads == 8 * self.num_k_heads + and self.num_v_heads % self.num_k_heads == 0 + and self.num_v_heads // self.num_k_heads in (1, 2, 3, 4, 8) and state_indices is not None and state_indices.size(1) <= MAX_FUSED_GDN_MTP_TOKENS and hasattr(torch.ops._C, "fused_gdn_decode_post_conv_mtp") From 69d3335066718b7e2b3d00713e02646874b424a9 Mon Sep 17 00:00:00 2001 From: Han <180024497+LH-and-FPGA@users.noreply.github.com> Date: Tue, 18 Aug 2026 03:38:42 +0200 Subject: [PATCH 07/15] [Bugfix][Quantization] Guard the MXFP8 FlashInfer path on FlashInfer availability (#52648) Signed-off-by: Han Li Signed-off-by: Misha Goin Co-authored-by: Claude Opus 5 (1M context) Co-authored-by: Misha Goin --- .../kernels/linear/mxfp8/flashinfer.py | 12 ++++++++---- .../layers/quantization/utils/mxfp8_utils.py | 3 ++- 2 files changed, 10 insertions(+), 5 deletions(-) diff --git a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py index d26e5579edbb..3150d93f99e1 100644 --- a/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py +++ b/vllm/model_executor/kernels/linear/mxfp8/flashinfer.py @@ -11,7 +11,7 @@ ) from vllm.platforms import current_platform from vllm.utils import flashinfer as vllm_flashinfer -from vllm.utils.flashinfer import has_flashinfer_cutedsl +from vllm.utils.flashinfer import has_flashinfer, has_flashinfer_cutedsl from .Mxfp8LinearKernel import Mxfp8LinearKernel, Mxfp8LinearLayerConfig @@ -23,9 +23,13 @@ class FlashInferCutlassMxfp8LinearKernel(Mxfp8LinearKernel): def is_supported( cls, compute_capability: int | None = None ) -> tuple[bool, str | None]: - if current_platform.has_device_capability(100): - return True, None - return False, "requires >=sm_100 (Blackwell)" + if not ( + current_platform.is_cuda() and current_platform.has_device_capability(100) + ): + return False, "requires >=sm_100 (Blackwell)" + if not has_flashinfer(): + return False, "requires FlashInfer" + return True, None @classmethod def can_implement(cls, c: Mxfp8LinearLayerConfig) -> tuple[bool, str | None]: diff --git a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py index 932bf9235289..1e45ea9ab03b 100644 --- a/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/mxfp8_utils.py @@ -177,8 +177,9 @@ def _mxfp8_e4m3_quantize_impl( alignment: int = 0, ) -> tuple[torch.Tensor, torch.Tensor]: from vllm.platforms import current_platform + from vllm.utils.flashinfer import has_flashinfer - if current_platform.has_device_capability(100): + if current_platform.has_device_capability(100) and has_flashinfer(): from flashinfer import mxfp8_quantize as flashinfer_mxfp8_quantize x_q, x_scales = flashinfer_mxfp8_quantize( From f4b161d7fca438bfe29509984759be1943a5aa88 Mon Sep 17 00:00:00 2001 From: Qiming Zhang Date: Mon, 17 Aug 2026 18:47:35 -0700 Subject: [PATCH 08/15] [XPU][UT] Fix OOM and skip graph case (#49287) Signed-off-by: mayuyuace Signed-off-by: Qiming Zhang Co-authored-by: Kunshang Ji --- tests/conftest.py | 37 +++++++++++++------ .../models/language/generation/test_hybrid.py | 3 ++ tests/models/language/pooling/test_colbert.py | 4 +- .../generation/test_voxtral_realtime.py | 8 ++-- tests/utils.py | 15 ++++++-- .../nixl_integration/run_accuracy_test.sh | 2 +- 6 files changed, 47 insertions(+), 22 deletions(-) diff --git a/tests/conftest.py b/tests/conftest.py index e277b35dd46b..bc6d9cb281e3 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -914,13 +914,13 @@ def __enter__(self): return self def __exit__(self, exc_type, exc_value, traceback): - from tests.utils import wait_for_rocm_memory_to_settle + from tests.utils import wait_for_memory_to_settle del self.model cleanup_dist_env_and_memory() # ROCm frees VRAM lazily; wait so a runner started right after this HF # model exits does not OOM on its startup memory guard. - wait_for_rocm_memory_to_settle( + wait_for_memory_to_settle( threshold_ratio=getattr(self, "threshold_ratios", None) ) if hasattr(self, "threshold_ratios"): @@ -998,9 +998,24 @@ def __init__( # V1 startup requires free_memory >= total * gpu_memory_utilization. # ROCm CI can hand a test a device that is still lazily releasing # VRAM from a previous process, so wait before constructing LLM. - from tests.utils import wait_for_rocm_memory_to_settle - - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + from tests.utils import wait_for_memory_to_settle + + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + elif current_platform.is_xpu(): + # The XPU/oneAPI runtime keeps ~1 GiB of context resident in the + # parent pytest process for its whole lifetime (grown by in-process + # HfRunner models), and distributed tests additionally allocate a + # CCL context in the engine subprocess. The default utilization of + # 0.92 leaves too little headroom for both, so lower it on XPU when + # the caller did not request an explicit value. + if "gpu_memory_utilization" not in kwargs: + kwargs["gpu_memory_utilization"] = 0.9 + gpu_memory_utilization = kwargs["gpu_memory_utilization"] + # XPU (Level Zero) can also release device memory lazily after a + # previous engine shuts down, so wait before constructing LLM. + from tests.utils import wait_for_memory_to_settle + + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) with init_ctx: self.llm = LLM( @@ -1326,14 +1341,14 @@ def collective_rpc(self, *args, **kwargs): def __enter__(self): return self - def _wait_for_rocm_memory_release(self, gpu_memory_utilization: float) -> None: - from tests.utils import wait_for_rocm_memory_to_settle + def _wait_for_memory_release(self, gpu_memory_utilization: float) -> None: + from tests.utils import wait_for_memory_to_settle # V1 startup requires free_memory >= total * gpu_memory_utilization. # Wait for the complementary used-memory ratio so the next runner does # not fail the startup guard immediately after this runner exits. The # wait is bounded so cleanup failures fail this test instead of hanging. - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) def __exit__(self, exc_type, exc_value, traceback): # Explicitly shutdown the engine core to release GPU resources @@ -1359,7 +1374,7 @@ def __exit__(self, exc_type, exc_value, traceback): del self.llm torch._dynamo.reset() cleanup_dist_env_and_memory() - self._wait_for_rocm_memory_release(gpu_memory_utilization) + self._wait_for_memory_release(gpu_memory_utilization) @pytest.fixture(scope="session") @@ -1746,7 +1761,7 @@ def clean_gpu_memory_between_tests(): import gc - from tests.utils import wait_for_gpu_memory_to_clear, wait_for_rocm_memory_to_settle + from tests.utils import wait_for_gpu_memory_to_clear, wait_for_memory_to_settle num_gpus = torch.accelerator.device_count() @@ -1755,7 +1770,7 @@ def _wait_for_settled_gpu_memory() -> None: return try: if current_platform.is_rocm(): - wait_for_rocm_memory_to_settle() + wait_for_memory_to_settle() else: wait_for_gpu_memory_to_clear( devices=list(range(num_gpus)), diff --git a/tests/models/language/generation/test_hybrid.py b/tests/models/language/generation/test_hybrid.py index bb73b1dea825..c5406fae48e0 100644 --- a/tests/models/language/generation/test_hybrid.py +++ b/tests/models/language/generation/test_hybrid.py @@ -9,6 +9,7 @@ from tests.models.registry import HF_EXAMPLE_MODELS from tests.utils import multi_gpu_test from vllm import LLM +from vllm.config import CUDAGraphMode from vllm.engine.arg_utils import EngineArgs from vllm.platforms import current_platform from vllm.sampling_params import SamplingParams @@ -210,6 +211,8 @@ def test_mamba_cache_cg_padding( cudagraph_dispatcher.initialize_cudagraph_keys( vllm_config.compilation_config.cudagraph_mode ) + if cudagraph_dispatcher.cudagraph_mode == CUDAGraphMode.NONE: + pytest.skip("CUDA/XPU graph is disabled.Please enable it to run this test. ") while ( len(example_prompts) == cudagraph_dispatcher.dispatch(len(example_prompts))[1].num_tokens diff --git a/tests/models/language/pooling/test_colbert.py b/tests/models/language/pooling/test_colbert.py index bb7afb23365b..dd4e1dea9772 100644 --- a/tests/models/language/pooling/test_colbert.py +++ b/tests/models/language/pooling/test_colbert.py @@ -11,7 +11,7 @@ import pytest import torch -from tests.utils import wait_for_rocm_memory_to_settle +from tests.utils import wait_for_memory_to_settle from vllm.distributed import cleanup_dist_env_and_memory from vllm.entrypoints.pooling.scoring.utils import compute_maxsim_score @@ -165,7 +165,7 @@ def _hf_colbert_model(model_name: str, hf_spec: dict, device: torch.device): finally: del hf_model, linear_weight cleanup_dist_env_and_memory() - wait_for_rocm_memory_to_settle() + wait_for_memory_to_settle() def _assert_embeddings_close(vllm_outputs, hf_embeddings): diff --git a/tests/models/multimodal/generation/test_voxtral_realtime.py b/tests/models/multimodal/generation/test_voxtral_realtime.py index df2d63ad84be..86fe732378be 100644 --- a/tests/models/multimodal/generation/test_voxtral_realtime.py +++ b/tests/models/multimodal/generation/test_voxtral_realtime.py @@ -104,9 +104,9 @@ async def async_engine(): from vllm.platforms import current_platform if current_platform.is_rocm(): - from tests.utils import wait_for_rocm_memory_to_settle + from tests.utils import wait_for_memory_to_settle - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) engine_args = AsyncEngineArgs(**ENGINE_CONFIG) llm = AsyncLLM.from_engine_args(engine_args) @@ -122,9 +122,9 @@ async def async_engine(): from vllm.distributed import cleanup_dist_env_and_memory cleanup_dist_env_and_memory() - from tests.utils import wait_for_rocm_memory_to_settle + from tests.utils import wait_for_memory_to_settle - wait_for_rocm_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) + wait_for_memory_to_settle(threshold_ratio=1.0 - gpu_memory_utilization) def test_voxtral_realtime_forward(audio_assets, tokenizer, vllm_runner, monkeypatch): diff --git a/tests/utils.py b/tests/utils.py index fac455ab738c..86774bf4bf82 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -1540,6 +1540,13 @@ def record_gpu_memory_usage_stats( mem_info = amdsmi_get_gpu_vram_usage(dev_handle) gb_used = mem_info["vram_used"] / 2**10 gb_total = mem_info["vram_total"] / 2**10 + elif current_platform.is_xpu(): + # nvml/amdsmi are unavailable on XPU. Query device memory through + # torch.accelerator.get_memory_info, which the XPU platform patches + # to return (free, total) bytes via Level Zero. + free_b, total_b = torch.accelerator.get_memory_info(device) + gb_used = (total_b - free_b) / 2**30 + gb_total = total_b / 2**30 else: dev_handle = get_nvml_device_handle(device) mem_info = nvmlDeviceGetMemoryInfo(dev_handle) @@ -1680,19 +1687,19 @@ def wait_for_gpu_memory_to_clear( time.sleep(poll_interval_s) -def wait_for_rocm_memory_to_settle( +def wait_for_memory_to_settle( *, threshold_ratio: float | dict[int, float] | None = 0.1, timeout_s: float = 240, ) -> None: - """Block until ROCm device VRAM usage drops below ``threshold_ratio``. + """Block until ROCm or XPU device VRAM usage drops below ``threshold_ratio``. - ROCm reclaims GPU memory more lazily than CUDA, so back-to-back model + ROCm and XPU reclaims GPU memory more lazily than CUDA, so back-to-back model loads in a single test process can OOM the *next* engine/model startup even after ``cleanup_dist_env_and_memory``. This gives the driver time to actually release VRAM before the next allocation. No-op off ROCm. """ - if not current_platform.is_rocm(): + if not current_platform.is_rocm() and not current_platform.is_xpu(): return num_gpus = current_platform.device_count() diff --git a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh index fc5c04a1ad07..b8a37e1e96fa 100755 --- a/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh +++ b/tests/v1/kv_connector/nixl_integration/run_accuracy_test.sh @@ -112,7 +112,7 @@ wait_for_server() { # Function to clean up previous instances wait_for_gpu_memory_release() { if [[ "$SMI_BIN" == *"rocm"* ]]; then - PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_rocm_memory_to_settle; wait_for_rocm_memory_to_settle()" + PYTHONPATH="${GIT_ROOT}" python3 -c "from tests.utils import wait_for_memory_to_settle; wait_for_memory_to_settle()" fi } From d5f5de7a7db4e13ea9f94ce8662b55b6b5e97d02 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E2=84=8D=F0=9D=95=A0=F0=9D=95=9D=F0=9D=95=9D=F0=9D=95=A0?= =?UTF-8?q?=F0=9D=95=A8=20=F0=9D=95=84=F0=9D=95=92=F0=9D=95=9F?= Date: Tue, 18 Aug 2026 05:55:37 +0300 Subject: [PATCH 09/15] [Bugfix] Fix DeepSeek V4 mHC broadcast buffer for weight sync (#52626) Signed-off-by: Hollow Man --- tests/kernels/test_mhc_kernels.py | 61 +++++++++++++++++++++++++ vllm/models/deepseek_v4/nvidia/model.py | 6 ++- 2 files changed, 66 insertions(+), 1 deletion(-) diff --git a/tests/kernels/test_mhc_kernels.py b/tests/kernels/test_mhc_kernels.py index 2bdce9f9c14d..a2070dd595e3 100644 --- a/tests/kernels/test_mhc_kernels.py +++ b/tests/kernels/test_mhc_kernels.py @@ -1,7 +1,10 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from types import SimpleNamespace + import pytest import torch +import torch.nn as nn import vllm.model_executor.kernels.mhc # noqa: F401 from vllm.model_executor.kernels.mhc.tilelang import ( @@ -9,6 +12,10 @@ _torch_hc_prenorm_gemm, ) from vllm.model_executor.layers.mhc import HAS_TILELANG_MHC +from vllm.models.deepseek_v4.nvidia.model import ( + DeepseekV4DecoderLayer, + DeepseekV4Model, +) from vllm.platforms import current_platform from vllm.utils.torch_utils import set_random_seed @@ -355,3 +362,57 @@ def test_hc_head_tilelang(num_tokens, hidden_size, hc_mult): out_ref = hc_head_ref(residual, fn, hc_scale, hc_base, rms_eps, hc_eps) torch.testing.assert_close(out, out_ref, atol=5e-2, rtol=1e-2) + + +def _make_mhc_decoder_layer(hc_mult: int, hidden_size: int) -> DeepseekV4DecoderLayer: + layer = DeepseekV4DecoderLayer.__new__(DeepseekV4DecoderLayer) + nn.Module.__init__(layer) + layer.hc_mult = hc_mult + layer.hidden_size = hidden_size + mix_hc = (2 + hc_mult) * hc_mult + layer.hc_attn_fn = nn.Parameter( + torch.randn(mix_hc, hc_mult * hidden_size, dtype=torch.float32), + requires_grad=False, + ) + layer.hc_attn_fn_broadcast = None + return layer + + +def _patch_first_rank_pp_group(monkeypatch): + monkeypatch.setattr( + "vllm.models.deepseek_v4.nvidia.model.get_pp_group", + lambda: SimpleNamespace(is_first_rank=True), + ) + + +def test_deepseek_v4_mhc_broadcast_finalize_sums_hc_streams(monkeypatch): + """First finalize (at the end of load_weights) allocates + hc_attn_fn_broadcast as hc_attn_fn summed over hc streams.""" + _patch_first_rank_pp_group(monkeypatch) + layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) + model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) + + DeepseekV4Model.finalize_mhc_broadcast_weights(model) + + assert layer.hc_attn_fn_broadcast is not None + expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) + assert torch.equal(layer.hc_attn_fn_broadcast, expected) + + +def test_deepseek_v4_mhc_broadcast_refit_refreshes_in_place(monkeypatch): + """Re-finalizing after a weight refit must copy into the existing + broadcast tensor so its address stays stable for captured CUDA graphs, + while picking up the new hc_attn_fn values.""" + _patch_first_rank_pp_group(monkeypatch) + layer = _make_mhc_decoder_layer(hc_mult=2, hidden_size=8) + model = SimpleNamespace(start_layer=0, end_layer=1, layers=[layer]) + + DeepseekV4Model.finalize_mhc_broadcast_weights(model) + buffer = layer.hc_attn_fn_broadcast + + layer.hc_attn_fn.add_(1.0) + DeepseekV4Model.finalize_mhc_broadcast_weights(model) + + assert layer.hc_attn_fn_broadcast is buffer + expected = layer.hc_attn_fn.detach().view(-1, 2, 8).sum(dim=1) + assert torch.equal(layer.hc_attn_fn_broadcast, expected) diff --git a/vllm/models/deepseek_v4/nvidia/model.py b/vllm/models/deepseek_v4/nvidia/model.py index ed7ee11ce575..17f29f4a1ae1 100644 --- a/vllm/models/deepseek_v4/nvidia/model.py +++ b/vllm/models/deepseek_v4/nvidia/model.py @@ -1372,11 +1372,15 @@ def finalize_mhc_broadcast_weights(self) -> None: return layer = self.layers[self.start_layer] if isinstance(layer, DeepseekV4DecoderLayer): - layer.hc_attn_fn_broadcast = ( + broadcast = ( layer.hc_attn_fn.detach() .view(-1, layer.hc_mult, layer.hidden_size) .sum(dim=1) ) + if layer.hc_attn_fn_broadcast is None: + layer.hc_attn_fn_broadcast = broadcast + else: + layer.hc_attn_fn_broadcast.copy_(broadcast) def _make_deepseek_v4_weights_mapper(expert_dtype: str) -> WeightsMapper: From e0e5a7fb2808504ba86c94f7b379e38496002fd0 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Mon, 17 Aug 2026 20:16:54 -0700 Subject: [PATCH 10/15] [Rust Frontend] Fix GLM-5.2 chat template rendering parity (#51426) Co-authored-by: Bugen Zhao Co-authored-by: OpenAI Codex Signed-off-by: Bugen Zhao Signed-off-by: Woosuk Kwon --- rust/Cargo.lock | 8 +-- rust/Cargo.toml | 4 +- rust/src/chat/src/renderer/hf/mod.rs | 66 +++++++++++++++++++++++ rust/src/chat/src/renderer/hf/template.rs | 20 +++++++ rust/src/chat/tests/roundtrip.rs | 14 +++++ 5 files changed, 106 insertions(+), 6 deletions(-) diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 092a79843124..6262bf2e783d 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -2402,9 +2402,9 @@ dependencies = [ [[package]] name = "minijinja" -version = "2.22.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ef84a52be188a1d4124bd717903fdde96ca4705f2b56adfe2d91fcc57fcb6987" +checksum = "86886cf6dbf4e614b19c9a1eec9775f021869d7eadde0fc73921a81b90c9b4c9" dependencies = [ "indexmap 2.13.0", "memo-map", @@ -2414,9 +2414,9 @@ dependencies = [ [[package]] name = "minijinja-contrib" -version = "2.22.0" +version = "2.24.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fd6e279dc925840c2d9ebc1e9f85410611d01292561297463ba3e516c3ad94dd" +checksum = "bd3e5f077bc2379f0f7d911e7cfdd921114ed99fc884533dca502944cb355b11" dependencies = [ "minijinja", "serde", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 405d6775bcf4..723e34806f78 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.22", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } -minijinja-contrib = { version = "2.22", features = ["pycompat"] } +minijinja = { version = "2.24", features = ["unstable_machinery", "json", "builtins", "loader", "loop_controls", "preserve_order"] } +minijinja-contrib = { version = "2.24", 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/src/renderer/hf/mod.rs b/rust/src/chat/src/renderer/hf/mod.rs index df0d94331a8b..7461acf2955a 100644 --- a/rust/src/chat/src/renderer/hf/mod.rs +++ b/rust/src/chat/src/renderer/hf/mod.rs @@ -290,6 +290,7 @@ struct TemplateToolDefinition { name: String, description: Option, parameters: JsonValue, + #[serde(skip_serializing_if = "Option::is_none")] strict: Option, } @@ -1280,6 +1281,71 @@ mod tests { assert_eq!(rendered, "get_weather|city"); } + #[test] + fn chat_template_preserves_openai_tool_field_order() { + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); + let tools = vec![ChatTool { + name: "get_weather".to_string(), + description: Some("Get weather".to_string()), + parameters: serde_json::json!({"type": "object"}), + strict: None, + }]; + request.tool_context = crate::request::ResolvedToolContext::new( + &request.messages, + tools, + Some(ChatToolChoice::Auto), + true, + ) + .expect("tool context should resolve"); + + let rendered = render( + Some("{% for key, value in tools[0].function.items() %}{{ key }}|{% endfor %}"), + &request, + ) + .unwrap(); + + assert_eq!(rendered, "name|description|parameters|"); + } + + #[test] + fn chat_template_preserves_python_optional_tool_fields() { + let mut request = sample_request(vec![ChatMessage::text(ChatRole::User, "hello")]); + let tools = vec![ + ChatTool { + name: "without_strict".to_string(), + description: None, + parameters: Value::Null, + strict: None, + }, + ChatTool { + name: "with_strict".to_string(), + description: Some("description".to_string()), + parameters: serde_json::json!({"type": "object"}), + strict: Some(false), + }, + ]; + request.tool_context = crate::request::ResolvedToolContext::new( + &request.messages, + tools, + Some(ChatToolChoice::Auto), + true, + ) + .expect("tool context should resolve"); + + let rendered = render( + Some( + "{% for tool in tools %}{% for key, value in tool.function.items() %}{{ key }}={{ value|tojson }}|{% endfor %};{% endfor %}", + ), + &request, + ) + .unwrap(); + + assert_eq!( + rendered, + "name=\"without_strict\"|description=null|parameters=null|;name=\"with_strict\"|description=\"description\"|parameters={\"type\": \"object\"}|strict=false|;" + ); + } + #[test] fn chat_template_exposes_assistant_tool_calls_and_tool_messages() { let request = sample_request(vec![ diff --git a/rust/src/chat/src/renderer/hf/template.rs b/rust/src/chat/src/renderer/hf/template.rs index 5a609621cc6e..b7de0172f36f 100644 --- a/rust/src/chat/src/renderer/hf/template.rs +++ b/rust/src/chat/src/renderer/hf/template.rs @@ -152,6 +152,26 @@ mod tests { assert_eq!(result, "[]"); } + #[test] + fn test_midchain_dotted_integer_lookup() { + let template = CompiledChatTemplate::new( + "{{ values.0.name }}".to_string(), + ChatTemplateContentFormatOption::Auto, + ) + .unwrap(); + let mut kwargs = HashMap::new(); + kwargs.insert("values".to_string(), serde_json::json!([{"name": "first"}])); + + let result = template + .apply(TemplateContext { + template_kwargs: Some(&kwargs), + ..Default::default() + }) + .unwrap(); + + assert_eq!(result, "first"); + } + #[test] fn test_chat_template_state_invalid_template() { let result = CompiledChatTemplate::new( diff --git a/rust/src/chat/tests/roundtrip.rs b/rust/src/chat/tests/roundtrip.rs index f8c3519bbfdd..049dc6a40308 100644 --- a/rust/src/chat/tests/roundtrip.rs +++ b/rust/src/chat/tests/roundtrip.rs @@ -185,6 +185,19 @@ impl RoundtripCase { } } + /// GLM-5.2 XML-like argument format with `` reasoning tags. + fn glm52() -> Self { + Self { + model_id: "zai-org/GLM-5.2-FP8", + assistant_stop_suffix: "", + tool_call_parser: ParserSelection::Auto, + reasoning_parser: ParserSelection::Auto, + thinking_behavior: ThinkingBehavior::Toggleable { default: true }, + json_fmt: compact_json_fmt(), + sort_json_keys: false, + } + } + /// Gemma4 channel reasoning with custom function-call arguments. fn gemma4() -> Self { Self { @@ -324,6 +337,7 @@ roundtrip_tests! { deepseek_v32 => [tool_call_mix], glm45 => [reasoning_and_content, tool_call_mix], glm47 => [reasoning_and_content, tool_call_mix], + glm52 => [reasoning_and_content, tool_call_mix], seed_oss => [reasoning_and_content, tool_call_mix], step3p5 => [reasoning_and_content], nemotron_v3 => [reasoning_and_content], From 101c4477dd4538f6fc32161d680ff2fba942b6d0 Mon Sep 17 00:00:00 2001 From: Do_it_now_! Date: Tue, 18 Aug 2026 12:46:08 +0800 Subject: [PATCH 11/15] [Bugfix] Handle DeepseekV4ForCausalLM in benchmark_moe get_model_params (#52044) Co-authored-by: Do_it_now_! <23432123@users.noreply.github.com> --- benchmarks/kernels/benchmark_moe.py | 1 + 1 file changed, 1 insertion(+) diff --git a/benchmarks/kernels/benchmark_moe.py b/benchmarks/kernels/benchmark_moe.py index 9ea22dd1304f..5d86a7599fca 100644 --- a/benchmarks/kernels/benchmark_moe.py +++ b/benchmarks/kernels/benchmark_moe.py @@ -766,6 +766,7 @@ def get_model_params(config): "DeepseekV2ForCausalLM", "DeepseekV3ForCausalLM", "DeepseekV32ForCausalLM", + "DeepseekV4ForCausalLM", "GlmMoeDsaForCausalLM", "Glm4MoeForCausalLM", "Glm4MoeLiteForCausalLM", From aa9903490c616dc6871e5acc62cec7bb1e5e9434 Mon Sep 17 00:00:00 2001 From: Komal Kumar Teru <162363718+kkt-cohere@users.noreply.github.com> Date: Tue, 18 Aug 2026 00:59:57 -0400 Subject: [PATCH 12/15] [Cohere] Misc changes to cohere model definitions (#50156) Co-authored-by: Cursor --- vllm/model_executor/models/cohere2_moe.py | 2 +- vllm/model_executor/models/commandr.py | 111 +++++++++++++++++----- 2 files changed, 87 insertions(+), 26 deletions(-) diff --git a/vllm/model_executor/models/cohere2_moe.py b/vllm/model_executor/models/cohere2_moe.py index 12291fb1423d..3873b6ee004f 100644 --- a/vllm/model_executor/models/cohere2_moe.py +++ b/vllm/model_executor/models/cohere2_moe.py @@ -208,7 +208,7 @@ def __init__( layer_types is not None and layer_types[self.layer_idx] == "sliding_attention" ): - self.sliding_window = config.sliding_window + self.sliding_window = config.sliding_window + 1 # Prefix-dense layers have full attention (no sliding window). When # prefix_dense_sliding_window_pattern == 1, they keep RoPE even though diff --git a/vllm/model_executor/models/commandr.py b/vllm/model_executor/models/commandr.py index 3c755c65e10c..5481fb0abcee 100644 --- a/vllm/model_executor/models/commandr.py +++ b/vllm/model_executor/models/commandr.py @@ -32,7 +32,11 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig -from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size +from vllm.distributed import ( + get_pp_group, + get_tensor_model_parallel_world_size, + tensor_model_parallel_all_reduce, +) from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.linear import ( @@ -44,9 +48,7 @@ from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import ( - row_parallel_weight_loader, -) +from vllm.model_executor.model_loader.weight_utils import row_parallel_weight_loader from vllm.model_executor.utils import set_weight_attrs from vllm.platforms import current_platform from vllm.sequence import IntermediateTensors @@ -87,18 +89,61 @@ def forward(self, hidden_states, residuals=None): return hidden_states, residuals +@torch.compile(backend=current_platform.simple_compile_backend) +def rms_norm_func(hidden_states, weight, variance_epsilon): + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + variance_epsilon) + + hidden_states = weight.to(torch.float32) * hidden_states + return hidden_states.to(input_dtype) + + +class RMSNorm(nn.Module): + def __init__(self, param_shape=None, eps=1e-6): + super().__init__() + self.weight = nn.Parameter(torch.ones(param_shape)) + self.variance_epsilon = eps + set_weight_attrs(self.weight, {"weight_loader": row_parallel_weight_loader}) + + def forward(self, hidden_states, residuals=None): + hidden_states = rms_norm_func(hidden_states, self.weight, self.variance_epsilon) + return hidden_states, residuals + + +def select_norm_impl(config: CohereConfig) -> tuple[type[nn.Module], float]: + """ + Returns the normalization layer class and epsilon value to use. + If `config.rms_norm_eps` is present, use RMSNorm. + Otherwise default to LayerNorm with `config.layer_norm_eps`. + """ + rms_eps = getattr(config, "rms_norm_eps", None) + if rms_eps is not None: + return RMSNorm, rms_eps + + return LayerNorm, config.layer_norm_eps + + # Copied from transformers.models.llama.modeling_llama.LlamaMLP Llama->Cohere class CohereMLP(nn.Module): def __init__( self, - config: CohereConfig | Cohere2Config, + config: CohereConfig, + intermediate_size: int | None = None, quant_config: QuantizationConfig | None = None, + # so we can reduce the attention and MLP outputs together + reduce_results: bool = False, prefix: str = "", ): super().__init__() self.config = config self.hidden_size = config.hidden_size - self.intermediate_size = config.intermediate_size + if intermediate_size is None: + self.intermediate_size = config.intermediate_size + else: + self.intermediate_size = intermediate_size self.gate_up_proj = MergedColumnParallelLinear( self.hidden_size, [self.intermediate_size] * 2, @@ -111,6 +156,7 @@ def __init__( self.hidden_size, bias=False, quant_config=quant_config, + reduce_results=reduce_results, prefix=f"{prefix}.down_proj", ) self.act_fn = SiluAndMul() @@ -133,11 +179,15 @@ def __init__( super().__init__() tp_size = get_tensor_model_parallel_world_size() self.config = config + self.layer_idx = extract_layer_index(prefix) self.attention_dropout = config.attention_dropout self.hidden_size = config.hidden_size self.total_num_heads = config.num_attention_heads self.num_heads = self.total_num_heads // tp_size - self.head_dim = self.hidden_size // self.total_num_heads + if hasattr(config, "head_dim"): + self.head_dim = config.head_dim + else: + self.head_dim = self.hidden_size // self.total_num_heads self.total_num_kv_heads = config.num_key_value_heads if self.total_num_kv_heads >= tp_size: # Number of KV heads is greater than TP size, so we partition @@ -171,6 +221,9 @@ def __init__( bias=False, quant_config=quant_config, prefix=f"{prefix}.o_proj", + # NOTE: reduction will happen in the decoder layer forward + # so we can combine the attention and MLP outputs together + reduce_results=False, ) self.rotary_emb = get_rope( self.head_dim, @@ -182,11 +235,14 @@ def __init__( # Model v2 has interleaved sliding windows, v1 does not self.v1 = isinstance(config, CohereConfig) + # cohere swa layer sees [pos - sliding_window, pos], i.e. sliding_window + 1 + # tokens. vLLM's FlashAttention backend does (value - 1, 0), so we pass + # sliding_window + 1 here to match the training convention. The same +1 + # propagates into the KV-cache eviction formula + # (single_type_kv_cache_manager.py), keeping both paths consistent. self.sliding_window = None - if not self.v1: - layer_idx = extract_layer_index(prefix) - if config.layer_types[layer_idx] == "sliding_attention": - self.sliding_window = config.sliding_window + if not self.v1 and config.layer_types[self.layer_idx] == "sliding_attention": + self.sliding_window = config.sliding_window + 1 self.attn = Attention( self.num_heads, @@ -198,13 +254,13 @@ def __init__( per_layer_sliding_window=self.sliding_window, prefix=f"{prefix}.attn", ) + norm_cls, norm_eps = select_norm_impl(config) if self.use_qk_norm: - self.q_norm = LayerNorm( - param_shape=(self.num_heads, self.head_dim), eps=config.layer_norm_eps + self.q_norm = norm_cls( + param_shape=(self.num_heads, self.head_dim), eps=norm_eps ) - self.k_norm = LayerNorm( - param_shape=(self.num_kv_heads, self.head_dim), - eps=config.layer_norm_eps, + self.k_norm = norm_cls( + param_shape=(self.num_kv_heads, self.head_dim), eps=norm_eps ) def _apply_qk_norm(self, q, k): @@ -242,18 +298,16 @@ def __init__( ): super().__init__() self.hidden_size = config.hidden_size - + self.tp_size = get_tensor_model_parallel_world_size() self.self_attn = CohereAttention( config, cache_config, quant_config=quant_config, prefix=f"{prefix}.self_attn", ) - self.mlp = CohereMLP(config, quant_config=quant_config, prefix=f"{prefix}.mlp") - self.input_layernorm = LayerNorm( - param_shape=(config.hidden_size), eps=config.layer_norm_eps - ) + norm_cls, norm_eps = select_norm_impl(config) + self.input_layernorm = norm_cls(param_shape=(config.hidden_size), eps=norm_eps) def forward( self, @@ -269,8 +323,14 @@ def forward( hidden_states=hidden_states, ) hidden_states_mlp = self.mlp(hidden_states) + parallel_block_output = hidden_states_attention + hidden_states_mlp # Add everything together - hidden_states = residual + hidden_states_attention + hidden_states_mlp + if self.tp_size > 1: + # do the reduction in 1 shot instead of 2 separate all reduce + parallel_block_output = tensor_model_parallel_all_reduce( + parallel_block_output + ) + hidden_states = residual + parallel_block_output return hidden_states, residual @@ -299,9 +359,10 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): ), prefix=f"{prefix}.layers", ) - self.norm = LayerNorm( - param_shape=(config.hidden_size), eps=config.layer_norm_eps - ) + + norm_cls, norm_eps = select_norm_impl(config) + self.norm = norm_cls(param_shape=(config.hidden_size), eps=norm_eps) + self.make_empty_intermediate_tensors = make_empty_intermediate_tensors_factory( ["hidden_states", "residual"], config.hidden_size ) From b0e9cff5e7f7dc4180ff0cd9380f9dcc2933c2f7 Mon Sep 17 00:00:00 2001 From: Yan Ma Date: Tue, 18 Aug 2026 14:27:41 +0800 Subject: [PATCH 13/15] [XPU] update xpu-manager to v2.1.0 (#52569) Signed-off-by: Yan Ma Co-authored-by: Kunshang Ji --- docker/Dockerfile.xpu | 7 ++----- 1 file changed, 2 insertions(+), 5 deletions(-) diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 31efca624151..ea7159f8c453 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -87,11 +87,8 @@ RUN apt-get update -y && \ apt-get install -y --no-install-recommends libhwloc15 libpciaccess0 libigsc1 intel-metrics-library && \ mkdir xpu-m && \ cd xpu-m && \ - wget https://github.com/intel/xpumanager/releases/download/v2.0.0/libxpum2_2.0.0-238.24.04_amd64.deb && \ - wget https://github.com/intel/xpumanager/releases/download/v2.0.0/xpu-smi_2.0.0-238.24.04_amd64.deb && \ - dpkg-deb -x libxpum2_2.0.0-238.24.04_amd64.deb / && \ - dpkg-deb -x xpu-smi_2.0.0-238.24.04_amd64.deb / && \ - ldconfig && \ + wget https://github.com/intel/xpumanager/releases/download/v2.1.0/xpu-smi_2.1.0+26.33.6468cec-1.24.04_amd64.deb && \ + apt install -y ./xpu-smi_2.1.0+26.33.6468cec-1.24.04_amd64.deb --install-suggests && \ cd .. && \ rm -rf xpu-m && \ add-apt-repository --remove -y ppa:kobuk-team/intel-graphics && \ From c89d692cb5fa430ddca06a9e16682977a18b55dc Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Tue, 18 Aug 2026 16:37:37 +1000 Subject: [PATCH 14/15] [Build] Propagate vLLM version to Rust binaries (#52593) Signed-off-by: Bugen Zhao Co-authored-by: Nick Hill Co-authored-by: Andreas Karatzas --- .buildkite/scripts/ci-bake-rocm.sh | 8 ++--- .buildkite/test-amd.yaml | 2 +- .buildkite/test_areas/rust_frontend.yaml | 2 +- docker/Dockerfile | 5 +++ docker/Dockerfile.cpu | 5 +++ docker/Dockerfile.rocm | 35 ++++++++++++------- docker/Dockerfile.rocm_gfx1250 | 35 ++++++++++++------- docker/Dockerfile.xpu | 5 +++ requirements/build/rust.txt | 1 + rust/Cargo.lock | 7 ++++ rust/Cargo.toml | 2 ++ rust/src/bench/Cargo.toml | 1 + rust/src/bench/src/main.rs | 2 +- rust/src/build-info/Cargo.toml | 8 +++++ rust/src/build-info/src/lib.rs | 10 ++++++ rust/src/cmd/Cargo.toml | 1 + rust/src/cmd/src/cli.rs | 3 +- rust/src/server/Cargo.toml | 1 + rust/src/server/src/routes/tests.rs | 2 +- rust/src/server/src/routes/version.rs | 2 +- setup.py | 7 +++- tests/entrypoints/openai/test_uds.py | 7 ++-- .../serve/instrumentator/test_basic.py | 7 ++-- tools/build_rust.py | 16 +++++++++ 24 files changed, 135 insertions(+), 39 deletions(-) create mode 100644 rust/src/build-info/Cargo.toml create mode 100644 rust/src/build-info/src/lib.rs diff --git a/.buildkite/scripts/ci-bake-rocm.sh b/.buildkite/scripts/ci-bake-rocm.sh index 1e143113141f..e3681e214a42 100644 --- a/.buildkite/scripts/ci-bake-rocm.sh +++ b/.buildkite/scripts/ci-bake-rocm.sh @@ -23,8 +23,8 @@ DEFAULT_CI_BASE_METADATA_VERSION="3" # 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 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" +DEFAULT_ROCM_RUST_CONTENT_FILES=".dockerignore .git_archival.txt pyproject.toml 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 fetch_vllm_0 fetch_vllm vllm-version 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 # ("csrc-rocm-branch-" and "rust-rocm-branch-", both 17 characters). ROCM_CACHE_BRANCH_TAG_MAX_LEN=111 @@ -523,8 +523,8 @@ prepare_ci_build_context() { return 1 fi - # setuptools-scm understands Git's stable archive format, so full-source - # wheels retain their exact version without copying mutable Git history. + # setuptools-scm understands Git's stable archive format, so wheels and + # Rust artifacts retain their exact version without copying Git history. if ! is_ci_base_target; then write_ci_git_archival_metadata "${source_root}" "${context_root}" \ || return $? diff --git a/.buildkite/test-amd.yaml b/.buildkite/test-amd.yaml index e9726f73ff5c..41dbd629bc9b 100644 --- a/.buildkite/test-amd.yaml +++ b/.buildkite/test-amd.yaml @@ -2356,7 +2356,7 @@ steps: - export VLLM_USE_RUST_FRONTEND=1 - export VLLM_WORKER_MULTIPROC_METHOD=spawn - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py - - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not server_load" - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" - pytest -v -s entrypoints/serve/tokenize/test_tokenization.py -k "not tokenizer_info" diff --git a/.buildkite/test_areas/rust_frontend.yaml b/.buildkite/test_areas/rust_frontend.yaml index e9e66c46c0c0..6486df1c00a1 100644 --- a/.buildkite/test_areas/rust_frontend.yaml +++ b/.buildkite/test_areas/rust_frontend.yaml @@ -58,7 +58,7 @@ steps: - export VLLM_WORKER_MULTIPROC_METHOD=spawn - PYTHONPATH=/vllm-workspace pytest -v -s entrypoints/serve/dev/rpc/test_collective_rpc.py # server_load can be flaky under the Rust frontend; keep it excluded for now. - - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not show_version and not server_load" + - pytest -v -s entrypoints/serve/instrumentator/test_basic.py -k "not server_load" # test_generate_logprobs expects Python-style top_logprobs truncation (dedup sampled + cap at max(k, 1)). - pytest -v -s entrypoints/scale_out/token_in_token_out/test_serving_tokens.py -k "not stream and not lora and not test_generate_logprobs and not stop_string_workflow" - pytest -v -s entrypoints/serve/instrumentator/test_metrics.py -k "text and not show and not run_batch and not test_metrics_counts and not test_metrics_exist" diff --git a/docker/Dockerfile b/docker/Dockerfile index 804a3a4bb3ef..b2c949d2d155 100644 --- a/docker/Dockerfile +++ b/docker/Dockerfile @@ -279,11 +279,16 @@ ENV CARGO_BUILD_JOBS=4 # separate local sccache daemon while sharing the same remote cache backend. ENV SCCACHE_SERVER_PORT=4227 +# Only Rust inputs are present, so suppress the artificial dirty state caused +# by other tracked files being absent from this build stage. +ENV SETUPTOOLS_SCM_PRETEND_METADATA={dirty=false} + # Build the release artifacts. Cache cargo registry/git, but not target/, # because stale target metadata can outlive source updates across BuildKit # cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=bind,source=.git,target=.git \ --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ if [ "$USE_SCCACHE" = "1" ]; then \ if [ -n "${SCCACHE_ENDPOINT}" ]; then export SCCACHE_ENDPOINT="${SCCACHE_ENDPOINT}"; fi; \ diff --git a/docker/Dockerfile.cpu b/docker/Dockerfile.cpu index 69abab3b774f..43fab4bccb16 100644 --- a/docker/Dockerfile.cpu +++ b/docker/Dockerfile.cpu @@ -166,11 +166,16 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 +# Only Rust inputs are present, so suppress the artificial dirty state caused +# by other tracked files being absent from this build stage. +ENV SETUPTOOLS_SCM_PRETEND_METADATA={dirty=false} + # Build the release artifacts. Cache cargo registry/git, but not target/, # because stale target metadata can outlive source updates across BuildKit # cache reuse. RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=bind,source=.git,target=.git \ --mount=type=secret,id=aws-credentials,target=/root/.aws/credentials,required=false \ if [ "$USE_SCCACHE" = "1" ]; then \ export RUSTC_WRAPPER=sccache; \ diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 0487c0c9165a..0b02b9ec2153 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -125,6 +125,18 @@ ONBUILD RUN git clone ${VLLM_REPO} \ && git fetch upstream ; fi FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm +# Resolve the device-independent source version once for every downstream +# artifact. setuptools-scm reads Git metadata for local/remote sources and the +# archive metadata generated for canonical CI contexts. +FROM fetch_vllm AS vllm-version +ARG COMMON_WORKDIR +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt \ + && python3 -c \ + "from setuptools_scm import get_version; print(get_version())" \ + > /vllm-version.txt + # ----------------------- # Rust/protoc toolchain inputs # @@ -238,6 +250,7 @@ COPY --from=rust-input /rust-input/vllm/rust rust COPY --from=rust-input /rust-input/vllm/rust-toolchain.toml rust-toolchain.toml COPY --from=rust-input /rust-input/vllm/tools/build_rust.py tools/build_rust.py COPY --from=rust-input /rust-input/vllm/build_rust.sh build_rust.sh +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while @@ -249,6 +262,7 @@ RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry, export RUSTC_WRAPPER=sccache \ && sccache --show-stats; \ fi \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && bash build_rust.sh \ && test -x vllm/vllm-rs \ && if [ "$USE_SCCACHE" = "1" ]; then \ @@ -334,6 +348,7 @@ RUN --mount=type=bind,source=pyproject.toml,target=${COMMON_WORKDIR}/vllm/pyproj FROM build_vllm_dependencies AS build_vllm COPY --from=fetch_vllm ${COMMON_WORKDIR}/vllm ${COMMON_WORKDIR}/vllm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -341,6 +356,7 @@ COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/vllm-rs ${COMMON_WORKDIR}/vll COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/vllm/vllm/ RUN cd "${COMMON_WORKDIR}/vllm" \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && export VLLM_USE_PRECOMPILED=1 \ && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ && export VLLM_DOCKER_BUILD_CONTEXT=1 \ @@ -640,6 +656,7 @@ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR ARG ROCM_TRITON_KERNELS_COMMIT +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt COPY --from=rocm-triton-kernels \ /opt/rocm-triton-kernels /opt/rocm-triton-kernels ENV TRITON_KERNELS_SRC_DIR=/opt/rocm-triton-kernels @@ -678,17 +695,12 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ cd vllm && bash tools/check_repo.sh; \ fi -# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) -# 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 \ +# Install the packaging helpers before modifying requirements/rocm.txt. +RUN --mount=type=cache,target=/root/.cache/uv \ 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}" \ - && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + && uv pip install --system setuptools_scm regex # Fail if git-based package dependencies are found in requirements files # (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) @@ -726,13 +738,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ && echo "Building vLLM with custom wheels from /install" \ && uv pip install --system --find-links /install -r requirements/rocm.txt -# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt # (setup.py auto-detects ccache/sccache in PATH) -RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ && export CCACHE_BASEDIR="$PWD" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist diff --git a/docker/Dockerfile.rocm_gfx1250 b/docker/Dockerfile.rocm_gfx1250 index 4c0008f6ecb7..6e88eb9fed9e 100644 --- a/docker/Dockerfile.rocm_gfx1250 +++ b/docker/Dockerfile.rocm_gfx1250 @@ -127,6 +127,17 @@ ONBUILD RUN git clone ${VLLM_REPO} \ && git fetch upstream ; fi FROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm +# Resolve the device-independent source version once for every downstream +# artifact. setuptools-scm reads Git metadata from local and remote sources. +FROM fetch_vllm AS vllm-version +ARG COMMON_WORKDIR +RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ + cd ${COMMON_WORKDIR}/vllm \ + && uv pip install --system -r requirements/build/rust.txt \ + && python3 -c \ + "from setuptools_scm import get_version; print(get_version())" \ + > /vllm-version.txt + # ----------------------- # Rust build stage # Builds the `vllm-rs` frontend in a dedicated stage so the wheel build stages @@ -152,6 +163,8 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd ${COMMON_WORKDIR}/vllm \ && uv pip install --system -r requirements/build/rust.txt +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt + # Build the release binary. Cargo's registry/git caches can be written by # concurrent BuildKit jobs on shared workers, so lock those cache mounts while # keeping the cache benefit. Do not cache target/, because stale target metadata @@ -159,6 +172,7 @@ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ RUN --mount=type=cache,id=vllm-rocm-cargo-registry,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,id=vllm-rocm-cargo-git,target=/root/.cargo/git,sharing=locked \ cd ${COMMON_WORKDIR}/vllm \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && bash build_rust.sh \ && test -x vllm/vllm-rs @@ -211,6 +225,7 @@ ARG COMMON_WORKDIR ENV VLLM_TARGET_DEVICE=rocm COPY --from=csrc-build ${COMMON_WORKDIR}/vllm/dist /precompiled-wheels +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -220,6 +235,7 @@ COPY --from=rust-build ${COMMON_WORKDIR}/vllm/vllm/_rust_*.so ${COMMON_WORKDIR}/ RUN --mount=type=cache,id=vllm-rocm-uv,target=/root/.cache/uv \ cd vllm \ && uv pip install --system -r requirements/rocm.txt \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && export VLLM_USE_PRECOMPILED=1 \ && export VLLM_PRECOMPILED_WHEEL_LOCATION="$(ls /precompiled-wheels/*.whl)" \ && export VLLM_DOCKER_BUILD_CONTEXT=1 \ @@ -426,6 +442,7 @@ RUN /bin/bash -lc 'set -euo pipefail; \ FROM fetch_vllm AS build_vllm_wheel_release ARG COMMON_WORKDIR +COPY --from=vllm-version /vllm-version.txt /tmp/vllm-version.txt # Drop the pre-built Rust artifacts into the source tree. setup.py detects # them and ships them as-is, skipping the local Rust build. @@ -461,15 +478,10 @@ RUN if [ "$GIT_REPO_CHECK" != "0" ]; then \ cd vllm && bash tools/check_repo.sh; \ fi -# Extract version from git BEFORE any modifications (pin_rocm_dependencies.py modifies requirements/rocm.txt) -# 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 \ +# Install the packaging helpers before modifying requirements/rocm.txt. +RUN --mount=type=cache,target=/root/.cache/uv \ 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}" \ - && echo "${VLLM_VERSION}" > /tmp/vllm_version.txt + && uv pip install --system setuptools_scm regex # Fail if git-based package dependencies are found in requirements files # (uv doesn't handle git+ URLs well, and packages should be distributed on PyPI) @@ -507,13 +519,12 @@ RUN --mount=type=cache,target=/root/.cache/uv \ && echo "Building vLLM with custom wheels from /install" \ && uv pip install --system --find-links /install -r requirements/rocm.txt -# Build wheel using pre-extracted version to avoid dirty state from modified requirements/rocm.txt +# Build wheel using the shared source version to avoid dirty state from modified requirements/rocm.txt # (setup.py auto-detects ccache/sccache in PATH) -RUN --mount=type=bind,source=.git,target=vllm/.git \ - --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ +RUN --mount=type=cache,id=vllm-rocm-ccache,target=/root/.cache/ccache \ cd vllm \ && export CCACHE_BASEDIR="$PWD" \ - && export SETUPTOOLS_SCM_PRETEND_VERSION=$(cat /tmp/vllm_version.txt) \ + && export SETUPTOOLS_SCM_PRETEND_VERSION="$(cat /tmp/vllm-version.txt)" \ && echo "Building wheel with version: ${SETUPTOOLS_SCM_PRETEND_VERSION}" \ && MAX_JOBS="${MAX_JOBS:-$(nproc)}" python3 setup.py bdist_wheel --dist-dir=dist diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index ea7159f8c453..456be41f6f2a 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -28,8 +28,13 @@ COPY build_rust.sh build_rust.sh # (rustc spawns enough concurrent processes to hit RLIMIT_NOFILE otherwise). ENV CARGO_BUILD_JOBS=4 +# Only Rust inputs are present, so suppress the artificial dirty state caused +# by other tracked files being absent from this build stage. +ENV SETUPTOOLS_SCM_PRETEND_METADATA={dirty=false} + RUN --mount=type=cache,target=/root/.cargo/registry,sharing=locked \ --mount=type=cache,target=/root/.cargo/git,sharing=locked \ + --mount=type=bind,source=.git,target=.git \ bash build_rust.sh FROM ubuntu:24.04 AS vllm-base diff --git a/requirements/build/rust.txt b/requirements/build/rust.txt index e2874dee0aba..f507deb8037b 100644 --- a/requirements/build/rust.txt +++ b/requirements/build/rust.txt @@ -1,4 +1,5 @@ # Dependencies for building Rust artifacts through setuptools-rust. setuptools>=77.0.3,<81.0.0 +setuptools-scm>=9.2.0 setuptools-rust>=1.9.0 wheel diff --git a/rust/Cargo.lock b/rust/Cargo.lock index 6262bf2e783d..a8608e34f2fd 100644 --- a/rust/Cargo.lock +++ b/rust/Cargo.lock @@ -5505,9 +5505,14 @@ dependencies = [ "tracing", "url", "uuid", + "vllm-build-info", "vllm-tracing", ] +[[package]] +name = "vllm-build-info" +version = "0.1.0" + [[package]] name = "vllm-chat" version = "0.1.0" @@ -5574,6 +5579,7 @@ dependencies = [ "tracing", "uuid", "vllm-bench", + "vllm-build-info", "vllm-chat", "vllm-engine-core-client", "vllm-managed-engine", @@ -5755,6 +5761,7 @@ dependencies = [ "url", "uuid", "validator", + "vllm-build-info", "vllm-chat", "vllm-engine-core-client", "vllm-llm", diff --git a/rust/Cargo.toml b/rust/Cargo.toml index 723e34806f78..5c4b47d5817e 100644 --- a/rust/Cargo.toml +++ b/rust/Cargo.toml @@ -1,6 +1,7 @@ [workspace] members = [ "src/bench", + "src/build-info", "src/chat", "src/cmd", "src/engine-core-client", @@ -135,6 +136,7 @@ url = "2.5.7" uuid = { version = "1.22.0", features = ["v4"] } validator = { version = "0.20.0", features = ["derive"] } vllm-bench = { path = "src/bench" } +vllm-build-info = { path = "src/build-info" } vllm-chat = { path = "src/chat" } vllm-engine-core-client = { path = "src/engine-core-client" } vllm-llm = { path = "src/llm" } diff --git a/rust/src/bench/Cargo.toml b/rust/src/bench/Cargo.toml index 960e7a62f7dd..24a8f12e910c 100644 --- a/rust/src/bench/Cargo.toml +++ b/rust/src/bench/Cargo.toml @@ -34,6 +34,7 @@ tokio-stream.workspace = true tracing.workspace = true url.workspace = true uuid.workspace = true +vllm-build-info.workspace = true vllm-tracing.workspace = true [lints] diff --git a/rust/src/bench/src/main.rs b/rust/src/bench/src/main.rs index 1764983e4487..3485e22bb1e5 100644 --- a/rust/src/bench/src/main.rs +++ b/rust/src/bench/src/main.rs @@ -12,7 +12,7 @@ use clap::Parser; #[command( name = "vllm-bench", about = "Benchmark online serving throughput", - version + version = vllm_build_info::VERSION )] struct Cli { #[command(flatten)] diff --git a/rust/src/build-info/Cargo.toml b/rust/src/build-info/Cargo.toml new file mode 100644 index 000000000000..41d3092af0f2 --- /dev/null +++ b/rust/src/build-info/Cargo.toml @@ -0,0 +1,8 @@ +[package] +name = "vllm-build-info" +version.workspace = true +edition.workspace = true +license.workspace = true + +[lints] +workspace = true diff --git a/rust/src/build-info/src/lib.rs b/rust/src/build-info/src/lib.rs new file mode 100644 index 000000000000..962fbe2f7d92 --- /dev/null +++ b/rust/src/build-info/src/lib.rs @@ -0,0 +1,10 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +/// The vLLM package version supplied by the build system. +/// +/// Direct Cargo builds fall back to the internal crate version. +pub const VERSION: &str = match option_env!("VLLM_RS_BUILD_VERSION") { + Some(version) => version, + None => env!("CARGO_PKG_VERSION"), +}; diff --git a/rust/src/cmd/Cargo.toml b/rust/src/cmd/Cargo.toml index a326a0f9992a..bf92e05ee186 100644 --- a/rust/src/cmd/Cargo.toml +++ b/rust/src/cmd/Cargo.toml @@ -28,6 +28,7 @@ tokio-util.workspace = true tracing.workspace = true uuid.workspace = true vllm-bench.workspace = true +vllm-build-info.workspace = true vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-managed-engine.workspace = true diff --git a/rust/src/cmd/src/cli.rs b/rust/src/cmd/src/cli.rs index 833b93ee43e3..e851fad54996 100644 --- a/rust/src/cmd/src/cli.rs +++ b/rust/src/cmd/src/cli.rs @@ -39,7 +39,8 @@ use crate::cli::unsupported::UnsupportedArgs; #[derive(Debug, Parser)] #[command( name = "vllm-rs", - about = "Rust frontend and managed-engine CLI for vLLM." + about = "Rust frontend and managed-engine CLI for vLLM.", + version = vllm_build_info::VERSION )] pub struct Cli { #[command(subcommand)] diff --git a/rust/src/server/Cargo.toml b/rust/src/server/Cargo.toml index 8b8f57a0dac5..f6e6f08b530f 100644 --- a/rust/src/server/Cargo.toml +++ b/rust/src/server/Cargo.toml @@ -45,6 +45,7 @@ tracing-subscriber.workspace = true url.workspace = true uuid.workspace = true validator.workspace = true +vllm-build-info.workspace = true vllm-chat.workspace = true vllm-engine-core-client.workspace = true vllm-llm.workspace = true diff --git a/rust/src/server/src/routes/tests.rs b/rust/src/server/src/routes/tests.rs index ac65ae269afd..260851b4efe5 100644 --- a/rust/src/server/src/routes/tests.rs +++ b/rust/src/server/src/routes/tests.rs @@ -1746,7 +1746,7 @@ async fn version_returns_engine_vllm_version() { json, json!({ "version": "test-vllm-version", - "rust_frontend_version": env!("CARGO_PKG_VERSION"), + "rust_frontend_version": vllm_build_info::VERSION, }) ); } diff --git a/rust/src/server/src/routes/version.rs b/rust/src/server/src/routes/version.rs index e4f6f3f1ea5e..2398224c5c98 100644 --- a/rust/src/server/src/routes/version.rs +++ b/rust/src/server/src/routes/version.rs @@ -21,6 +21,6 @@ pub async fn version(State(state): State>) -> Json None: ): cmdclass["build_rust"] = precompiled_build_rust +# Resolve the Python version first because get_vllm_version() may set +# SETUPTOOLS_SCM_PRETEND_VERSION, which the Rust version should inherit. +vllm_version = get_vllm_version() +rust_build.prepare_build_environment() + # Rust artifacts, built via setuptools-rust and installed into the package # directory alongside the Python modules. rust_extensions = rust_build.rust_extensions( @@ -1501,7 +1506,7 @@ def add_vllm_package_data(filename: str) -> None: setup( # static metadata should rather go in pyproject.toml - version=get_vllm_version(), + version=vllm_version, ext_modules=ext_modules, rust_extensions=rust_extensions, install_requires=get_requirements(), diff --git a/tests/entrypoints/openai/test_uds.py b/tests/entrypoints/openai/test_uds.py index f79e40ee4132..7c8557141087 100644 --- a/tests/entrypoints/openai/test_uds.py +++ b/tests/entrypoints/openai/test_uds.py @@ -6,6 +6,7 @@ import httpx import pytest +from vllm import envs from vllm.version import __version__ as VLLM_VERSION from ...utils import RemoteOpenAIServer @@ -40,5 +41,7 @@ async def test_show_version(server: RemoteOpenAIServer): response = client.get(server.url_for("version")) response.raise_for_status() - # Tolerate additive fields (e.g. the Rust frontend reports its own version). - assert response.json()["version"] == VLLM_VERSION + payload = response.json() + assert payload["version"] == VLLM_VERSION + if envs.VLLM_USE_RUST_FRONTEND: + assert payload["rust_frontend_version"] == VLLM_VERSION diff --git a/tests/entrypoints/serve/instrumentator/test_basic.py b/tests/entrypoints/serve/instrumentator/test_basic.py index 5b00d2e578e5..73a97c4fa264 100644 --- a/tests/entrypoints/serve/instrumentator/test_basic.py +++ b/tests/entrypoints/serve/instrumentator/test_basic.py @@ -12,6 +12,7 @@ from fastapi import Request from tests.utils import RemoteOpenAIServer +from vllm import envs from vllm.v1.engine.exceptions import EngineDeadError from vllm.version import __version__ as VLLM_VERSION @@ -83,8 +84,10 @@ async def test_show_version(server: RemoteOpenAIServer): response = requests.get(server.url_for("version")) response.raise_for_status() - # Tolerate additive fields (e.g. the Rust frontend reports its own version). - assert response.json()["version"] == VLLM_VERSION + payload = response.json() + assert payload["version"] == VLLM_VERSION + if envs.VLLM_USE_RUST_FRONTEND: + assert payload["rust_frontend_version"] == VLLM_VERSION @pytest.mark.asyncio diff --git a/tools/build_rust.py b/tools/build_rust.py index b5951bfe576b..a000ec8169fd 100644 --- a/tools/build_rust.py +++ b/tools/build_rust.py @@ -11,8 +11,23 @@ from setuptools import setup from setuptools_rust import Binding, RustExtension +from setuptools_scm import get_version ROOT_DIR = Path(__file__).resolve().parents[1] +VLLM_RS_BUILD_VERSION = "VLLM_RS_BUILD_VERSION" + + +def prepare_build_environment() -> str | None: + """Set the device-independent vLLM source version for Rust artifacts.""" + version = os.getenv(VLLM_RS_BUILD_VERSION) or None + if version is None: + try: + version = get_version(root=ROOT_DIR) + except LookupError: + return None + + os.environ[VLLM_RS_BUILD_VERSION] = version + return version def rust_extensions(*, optional: bool = False) -> list[RustExtension]: @@ -61,6 +76,7 @@ def build_binary(build_rust_args: list[str]) -> None: def main() -> None: + prepare_build_environment() build_binary(sys.argv[1:]) From d785eb51ce24210ec6da82140729d30822332032 Mon Sep 17 00:00:00 2001 From: flb_ Date: Tue, 18 Aug 2026 14:46:44 +0800 Subject: [PATCH 15/15] [Test] Add pause/resume E2E tests (#52144) Signed-off-by: flb_ --- tests/entrypoints/serve/dev/rlhf/conftest.py | 414 ++++++++++++++++++ .../dev/rlhf/state_transitions/__init__.py | 0 .../state_transitions/test_pause_resume.py | 168 +++++++ 3 files changed, 582 insertions(+) create mode 100644 tests/entrypoints/serve/dev/rlhf/conftest.py create mode 100644 tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py create mode 100644 tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py diff --git a/tests/entrypoints/serve/dev/rlhf/conftest.py b/tests/entrypoints/serve/dev/rlhf/conftest.py new file mode 100644 index 000000000000..b2050997905a --- /dev/null +++ b/tests/entrypoints/serve/dev/rlhf/conftest.py @@ -0,0 +1,414 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Shared fixtures and helpers for the RL lifecycle test suite. + +All test modules under this directory import from here to avoid duplication. + +RFC: https://github.com/vllm-project/vllm/issues/45585 +PR: https://github.com/vllm-project/vllm/pull/45586 +""" + +import contextlib +import json +import os +import subprocess +import sys +import threading +import time +from collections.abc import Callable +from contextlib import contextmanager +from dataclasses import dataclass, field +from typing import Any + +import requests + +# --------------------------------------------------------------------------- +# Model / server defaults +# --------------------------------------------------------------------------- + + +MODEL_NAME = os.environ.get("VLLM_TEST_MODEL", "Qwen/Qwen3-0.6B") + + +_BASE_ARGS = [ + "--dtype", + "bfloat16", + "--max-model-len", + "2048", + "--max-num-seqs", + "32", + "--gpu-memory-utilization", + "0.75", + "--enable-sleep-mode", + "--enforce-eager", +] + + +# Lightweight args for state-machine / protocol tests that don't need real +# weights (avoids spending time downloading a 1B checkpoint in T0 tests). +_DUMMY_ARGS = [ + "--dtype", + "bfloat16", + "--max-model-len", + "128", + "--max-num-seqs", + "8", + "--gpu-memory-utilization", + "0.5", + "--enable-sleep-mode", + "--enforce-eager", + "--load-format", + "dummy", +] + + +# --------------------------------------------------------------------------- +# Server harness +# --------------------------------------------------------------------------- + + +@contextmanager +def server( + extra_args=None, + port: int = 8770, + timeout: float = 180.0, + dummy_weights: bool = False, +): + """Launch a vLLM server with the dev router; yield its base URL. + + Args: + extra_args: Additional CLI flags appended after the base args. + port: HTTP port to bind (caller is responsible for uniqueness). + timeout: Seconds to wait for /health before giving up. + dummy_weights: If True, use --load-format dummy (fast, no real weights). + """ + env = {**os.environ, "VLLM_SERVER_DEV_MODE": "1"} + base = _DUMMY_ARGS if dummy_weights else _BASE_ARGS + cmd = [ + sys.executable, + "-m", + "vllm.entrypoints.openai.api_server", + "--model", + MODEL_NAME, + "--port", + str(port), + "--served-model-name", + "m", + *(base + (extra_args or [])), + ] + proc = subprocess.Popen( + cmd, env=env, stdout=subprocess.DEVNULL, stderr=subprocess.PIPE + ) + url = f"http://localhost:{port}" + try: + deadline = time.time() + timeout + while time.time() < deadline: + if proc.poll() is not None: + err = ( + proc.stderr.read(4000).decode(errors="replace") + if proc.stderr + else "" + ) + raise RuntimeError(f"vllm server exited during startup:\n{err}") + with contextlib.suppress(Exception): + if requests.get(f"{url}/health", timeout=3).status_code == 200: + break + time.sleep(1) + else: + proc.terminate() + raise RuntimeError("vllm server did not start in time") + yield url + finally: + proc.terminate() + with contextlib.suppress(subprocess.TimeoutExpired): + proc.wait(timeout=10) + if proc.poll() is None: + proc.kill() + + +# --------------------------------------------------------------------------- +# Polling helper (200-lie workaround) +# --------------------------------------------------------------------------- + + +def poll_until( + predicate: Callable[[], bool], + timeout: float = 10.0, + interval: float = 0.5, +) -> bool: + """Poll predicate() until it returns True or timeout expires. + + Workaround for the vLLM sleep/wake "200-lie" — the HTTP endpoints may + return 200 before the underlying operation is complete, so callers that + need to verify state *after* an operation can use this helper instead of + assuming the 200 means completion. + + Returns True if predicate became true within timeout, False otherwise. + """ + deadline = time.time() + timeout + while time.time() < deadline: + try: + if predicate(): + return True + except Exception: + pass + time.sleep(interval) + return False + + +# --------------------------------------------------------------------------- +# HTTP helpers — generation +# --------------------------------------------------------------------------- + + +def gen(url, prompt="The capital of France is", max_tokens=8, timeout=30): + """Fire a /v1/completions request; return JSON or None on any error.""" + try: + r = requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0, + }, + timeout=timeout, + ) + return r.json() + except Exception: + return None + + +def gen_with_logprobs( + url, prompt="The capital of France is", max_tokens=8, logprobs=5, timeout=30 +): + """Fire a /v1/completions request with logprobs; return JSON or None.""" + try: + r = requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": prompt, + "max_tokens": max_tokens, + "temperature": 0, + "logprobs": logprobs, + }, + timeout=timeout, + ) + return r.json() + except Exception: + return None + + +def ok(resp) -> bool: + """True iff resp is a successful completion (has choices, no error key).""" + return ( + resp is not None + and "choices" in resp + and bool(resp["choices"]) + and "error" not in resp + ) + + +# --------------------------------------------------------------------------- +# HTTP helpers — stream generation +# --------------------------------------------------------------------------- + + +@dataclass +class StreamResult: + started: threading.Event = field(default_factory=threading.Event) + done: threading.Event = field(default_factory=threading.Event) + chunks: list[dict[str, Any]] = field(default_factory=list) + finish_reason: str | None = None + error: Exception | None = None + + +def stream_completion(url: str, result: StreamResult, max_tokens: int) -> None: + try: + with requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": "Count upward slowly: one, two, three,", + "max_tokens": max_tokens, + "temperature": 0, + "ignore_eos": True, + "stream": True, + }, + stream=True, + timeout=(5, 60), + ) as response: + response.raise_for_status() + for line in response.iter_lines(decode_unicode=True): + if not line or line == "data: [DONE]": + continue + assert line.startswith("data: ") + chunk = json.loads(line.removeprefix("data: ")) + result.chunks.append(chunk) + choice = chunk["choices"][0] + if choice.get("text"): + result.started.set() + if choice.get("finish_reason") is not None: + result.finish_reason = choice["finish_reason"] + except Exception as error: + result.error = error + finally: + result.done.set() + + +def start_stream(url: str, max_tokens: int) -> tuple[StreamResult, threading.Thread]: + result = StreamResult() + thread = threading.Thread( + target=stream_completion, + args=(url, result, max_tokens), + ) + thread.start() + started = result.started.wait(timeout=10) + if not started or result.done.is_set(): + pause(url, mode="abort") + resume(url) + thread.join(timeout=10) + assert started, "request did not start generating" + assert not result.done.is_set(), "request completed before it could be paused" + return result, thread + + +# --------------------------------------------------------------------------- +# HTTP helpers — pause / resume +# --------------------------------------------------------------------------- + + +def pause(url, mode="abort", clear_cache=True): + return requests.post( + f"{url}/pause", + params={"mode": mode, "clear_cache": clear_cache}, + timeout=15, + ).status_code + + +def resume(url): + return requests.post(f"{url}/resume", timeout=10).status_code + + +def completion_with_cache_details(url: str, prompt: str) -> dict[str, Any]: + response = requests.post( + f"{url}/v1/completions", + json={ + "model": "m", + "prompt": prompt, + "max_tokens": 8, + "temperature": 0, + "logprobs": 1, + }, + timeout=30, + ) + response.raise_for_status() + return response.json() + + +def golden_output(response: dict[str, Any]) -> dict[str, Any]: + choice = response["choices"][0] + usage = response["usage"] + return { + "text": choice["text"], + "finish_reason": choice["finish_reason"], + "tokens": choice["logprobs"]["tokens"], + "prompt_tokens": usage["prompt_tokens"], + "completion_tokens": usage["completion_tokens"], + } + + +def cached_tokens(response: dict[str, Any]) -> int: + return response["usage"]["prompt_tokens_details"]["cached_tokens"] + + +# --------------------------------------------------------------------------- +# HTTP helpers — sleep / wake +# --------------------------------------------------------------------------- + + +def sleep(url, level=1, mode="abort"): + return requests.post( + f"{url}/sleep", params={"level": level, "mode": mode}, timeout=15 + ).status_code + + +def wake(url, tags=None): + params = {"tags": tags} if tags else {} + return requests.post(f"{url}/wake_up", params=params, timeout=20).status_code + + +def is_sleeping(url) -> bool: + return requests.get(f"{url}/is_sleeping", timeout=5).json()["is_sleeping"] + + +def is_paused(url) -> bool: + return requests.get(f"{url}/is_paused", timeout=5).json()["is_paused"] + + +def health(url) -> int: + try: + return requests.get(f"{url}/health", timeout=5).status_code + except Exception: + return 0 + + +# --------------------------------------------------------------------------- +# HTTP helpers — weight transfer +# --------------------------------------------------------------------------- + + +def start_weight_update(url, is_checkpoint_format=True): + return requests.post( + f"{url}/start_weight_update", + json={"is_checkpoint_format": is_checkpoint_format}, + timeout=10, + ) + + +def finish_weight_update(url): + return requests.post(f"{url}/finish_weight_update", timeout=10) + + +def get_world_size(url, include_dp=True): + return requests.get( + f"{url}/get_world_size", + params={"include_dp": include_dp}, + timeout=5, + ) + + +# --------------------------------------------------------------------------- +# GPU / metrics helpers +# --------------------------------------------------------------------------- + + +def gpu_free_bytes(device: int = 0) -> int: + """Read GPU free bytes via subprocess to avoid import-time torch init.""" + out = subprocess.check_output( + [ + sys.executable, + "-c", + f"import torch; f,_=torch.accelerator.get_memory_info({device}); print(f)", + ], + timeout=10, + ) + return int(out.strip()) + + +def sleep_metrics(url): + """Return (awake, weights_offloaded, discard_all) from /metrics.""" + try: + from prometheus_client.parser import text_string_to_metric_families + except ImportError: + return None, None, None + + r = requests.get(f"{url}/metrics", timeout=5) + vals: dict = {} + for family in text_string_to_metric_families(r.text): + if family.name == "vllm:engine_sleep_state": + for s in family.samples: + vals[s.labels.get("sleep_state", "")] = s.value + return vals.get("awake"), vals.get("weights_offloaded"), vals.get("discard_all") diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py b/tests/entrypoints/serve/dev/rlhf/state_transitions/__init__.py new file mode 100644 index 000000000000..e69de29bb2d1 diff --git a/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py b/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py new file mode 100644 index 000000000000..3de9e0054963 --- /dev/null +++ b/tests/entrypoints/serve/dev/rlhf/state_transitions/test_pause_resume.py @@ -0,0 +1,168 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""End-to-end tests for the vLLM RL pause/resume lifecycle.""" + +import os +import threading +from typing import Any +from unittest.mock import patch + +import pytest +import requests + +from tests.entrypoints.serve.dev.rlhf.conftest import ( + cached_tokens, + completion_with_cache_details, + gen, + golden_output, + is_paused, + ok, + pause, + resume, + server, + start_stream, +) + + +@pytest.fixture(scope="module", params=[False, True], ids=["MRV1", "MRV2"]) +def use_v2(request): + return request.param + + +@pytest.fixture(scope="module") +def server_url(use_v2): + env_vars = { + "VLLM_USE_V2_MODEL_RUNNER": "1" if use_v2 else "0", + } + + with ( + patch.dict(os.environ, env_vars), + server( + extra_args=[ + "--enable-prefix-caching", + "--enable-prompt-tokens-details", + ] + ) as url, + ): + yield url + + +@pytest.fixture(autouse=True) +def restore_unpaused_state(server_url): + assert resume(server_url) == 200 + yield + assert resume(server_url) == 200 + + +class TestPauseResume: + def test_state_and_idempotency_across_cycles(self, server_url): + assert not is_paused(server_url) + + assert resume(server_url) == 200 + assert resume(server_url) == 200 + assert not is_paused(server_url) + + for _ in range(2): + for mode in ("abort", "wait", "keep"): + assert pause(server_url, mode=mode) == 200 + assert pause(server_url, mode=mode) == 200 + assert is_paused(server_url) + + assert resume(server_url) == 200 + assert resume(server_url) == 200 + assert not is_paused(server_url) + + def test_invalid_mode_preserves_state(self, server_url): + for paused in (False, True): + if paused: + assert pause(server_url) == 200 + assert is_paused(server_url) is paused + + response = requests.post( + f"{server_url}/pause", + params={"mode": "invalid"}, + timeout=10, + ) + assert response.status_code == 400 + assert response.json()["error"]["param"] == "query.mode" + assert is_paused(server_url) is paused + + assert resume(server_url) == 200 + + @pytest.mark.parametrize( + ("mode", "max_tokens", "inflight_finish_reason"), + [ + pytest.param("abort", 256, "abort", id="abort"), + pytest.param("wait", 256, "length", id="wait"), + pytest.param("keep", 256, "length", id="keep"), + ], + ) + def test_mode_request_lifecycle( + self, + server_url, + mode, + max_tokens, + inflight_finish_reason, + ): + inflight, inflight_thread = start_stream(server_url, max_tokens) + new_result: dict[str, Any] = {} + new_done = threading.Event() + + def _new_request(): + new_result["response"] = gen(server_url, max_tokens=4, timeout=60) + new_done.set() + + new_thread = threading.Thread(target=_new_request) + try: + assert pause(server_url, mode=mode) == 200 + assert is_paused(server_url) + + if mode in ("abort", "wait"): + assert inflight.done.is_set() + else: + chunks_after_pause = len(inflight.chunks) + assert not inflight.done.wait(timeout=5) + assert len(inflight.chunks) == chunks_after_pause, ( + "in-flight request continued generating in keep mode" + ) + + new_thread.start() + assert not new_done.wait(timeout=0.3), ( + "new request completed while generation was paused" + ) + finally: + assert resume(server_url) == 200 + inflight_thread.join(timeout=30) + if new_thread.ident is not None: + new_thread.join(timeout=30) + + assert not inflight_thread.is_alive() + assert inflight.error is None + assert inflight.finish_reason == inflight_finish_reason + assert not new_thread.is_alive() + assert ok(new_result.get("response")) + + def test_clear_cache_preserves_output_and_controls_prefix_cache(self, server_url): + prompt = ( + "Paris is the capital of France. Berlin is the capital of Germany. " + ) * 20 + + assert pause(server_url, clear_cache=True) == 200 + assert resume(server_url) == 200 + baseline = completion_with_cache_details(server_url, prompt) + warmed = completion_with_cache_details(server_url, prompt) + + assert cached_tokens(baseline) == 0 + assert cached_tokens(warmed) > 0 + + assert pause(server_url, clear_cache=False) == 200 + assert resume(server_url) == 200 + preserved = completion_with_cache_details(server_url, prompt) + assert golden_output(preserved) == golden_output(baseline) + assert cached_tokens(preserved) > 0 + + assert pause(server_url, clear_cache=True) == 200 + assert resume(server_url) == 200 + cleared = completion_with_cache_details(server_url, prompt) + assert golden_output(cleared) == golden_output(baseline) + assert cached_tokens(cleared) == 0