diff --git a/flash-attention-v100/include/fused_mha.h b/flash-attention-v100/include/fused_mha.h index 3a721d1d59..4004701800 100644 --- a/flash-attention-v100/include/fused_mha.h +++ b/flash-attention-v100/include/fused_mha.h @@ -59,6 +59,8 @@ at::Tensor flash_attention_grouped_verify_paged( int64_t flash_attention_grouped_verify_max_query_tokens(); +int64_t flash_attention_grouped_sparse_page4_abi_version(); + at::Tensor flash_attention_grouped_sparse_page4( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, diff --git a/flash-attention-v100/kernel/flash_decode_paged.cu b/flash-attention-v100/kernel/flash_decode_paged.cu index 2f83d2e25f..7384538f75 100644 --- a/flash-attention-v100/kernel/flash_decode_paged.cu +++ b/flash-attention-v100/kernel/flash_decode_paged.cu @@ -4034,6 +4034,12 @@ int64_t flash_attention_grouped_verify_max_query_tokens() { return kGroupedVerifyMaxSupportedQ; } +int64_t flash_attention_grouped_sparse_page4_abi_version() { + // Version 1 accepted FP16 K/V through the nine-argument forward binding. + // Version 2 adds kv_cache_dtype and calibrated K/V scales. + return 2; +} + at::Tensor flash_attention_grouped_verify_paged( const at::Tensor& q, const at::Tensor& k_cache, const at::Tensor& v_cache, std::optional& out_, const at::Tensor& block_table, diff --git a/flash-attention-v100/kernel/fused_mha_api.cpp b/flash-attention-v100/kernel/fused_mha_api.cpp index 77f595d4ad..76859871ac 100644 --- a/flash-attention-v100/kernel/fused_mha_api.cpp +++ b/flash-attention-v100/kernel/fused_mha_api.cpp @@ -29,6 +29,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { "Maximum query length supported by grouped DFlash2 verification"); m.def("grouped_sparse_page4_fwd", &flash_attention_grouped_sparse_page4, "Grouped exact QSA page4 attention over paged KV cache (Volta)"); + m.def("grouped_sparse_page4_abi_version", + &flash_attention_grouped_sparse_page4_abi_version, + "Grouped sparse page4 forward ABI version"); m.def("grouped_sparse_page4_plan_fwd", &flash_attention_grouped_sparse_page4_plan, "Build grouped exact QSA page4 tables over paged KV cache (Volta)"); diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 9de7aa40e5..8dcff8d3c5 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -308,19 +308,8 @@ def test_sm70_mtp_defaults_require_env_opt_in(monkeypatch): } assert args.enable_prefix_caching is True assert args.mamba_cache_mode == "align" - assert args.max_num_seqs == 4 - assert args.compilation_config.cudagraph_capture_sizes == [ - 1, - 2, - 4, - 5, - 8, - 9, - 10, - 15, - 18, - 20, - ] + assert args.max_num_seqs is None + assert args.compilation_config.cudagraph_capture_sizes is None def test_sm70_mtp_split_cudagraphs_are_opt_in(monkeypatch): @@ -332,7 +321,7 @@ def test_sm70_mtp_split_cudagraphs_are_opt_in(monkeypatch): }, ) - assert args.compilation_config.cudagraph_capture_sizes == [5, 10, 20] + assert args.compilation_config.cudagraph_capture_sizes is None def test_sm70_mtp_split_cudagraphs_cover_production_batches(monkeypatch): @@ -348,6 +337,7 @@ def test_sm70_mtp_split_cudagraphs_cover_production_batches(monkeypatch): assert args.compilation_config.cudagraph_capture_sizes == [ 5, 10, + 15, 20, 30, 40, @@ -395,7 +385,7 @@ def test_sm70_explicit_mtp_still_gets_safe_defaults(monkeypatch): } assert args.enable_prefix_caching is True assert args.mamba_cache_mode == "align" - assert args.max_num_seqs == 4 + assert args.max_num_seqs is None def test_sm70_explicit_dflash_preserves_probabilistic_default(monkeypatch): diff --git a/tests/kernels/attention/test_sm70_qsa_grouped_page4.py b/tests/kernels/attention/test_sm70_qsa_grouped_page4.py index 2880b1b913..9b3de24e7f 100644 --- a/tests/kernels/attention/test_sm70_qsa_grouped_page4.py +++ b/tests/kernels/attention/test_sm70_qsa_grouped_page4.py @@ -21,10 +21,21 @@ def _require_grouped_page4(): return extension +def _grouped_page4_abi_version(extension) -> int: + capability = getattr(extension, "grouped_sparse_page4_abi_version", None) + if callable(capability): + return int(capability()) + doc = getattr(extension.grouped_sparse_page4_fwd, "__doc__", "") or "" + return 2 if "arg11:" in doc else 1 if "arg8:" in doc else 0 + + @pytest.mark.parametrize("kv_cache_dtype", ["auto", "fp8_e4m3"]) @torch.inference_mode() def test_sm70_qsa_grouped_page4_calibrated_kv(kv_cache_dtype: str) -> None: extension = _require_grouped_page4() + abi_version = _grouped_page4_abi_version(extension) + if kv_cache_dtype == "fp8_e4m3" and abi_version < 2: + pytest.skip("installed grouped page4 ABI does not support quantized K/V") torch.manual_seed(7) query = torch.randn((8, 6, 256), dtype=torch.float16, device="cuda") * 0.2 key = torch.randn((1, 4, 1, 256), dtype=torch.float16, device="cuda") * 0.35 @@ -48,7 +59,7 @@ def test_sm70_qsa_grouped_page4_calibrated_kv(kv_cache_dtype: str) -> None: seq_lens = torch.tensor([4], dtype=torch.int32, device="cuda") output = torch.empty_like(query) lse = torch.empty((8, 6), dtype=torch.float32, device="cuda") - extension.grouped_sparse_page4_fwd( + args = ( query, key_cache, value_cache, @@ -58,10 +69,16 @@ def test_sm70_qsa_grouped_page4_calibrated_kv(kv_cache_dtype: str) -> None: seq_lens, lse, 256**-0.5, - kv_cache_dtype, - k_scale, - v_scale, ) + if abi_version >= 2: + extension.grouped_sparse_page4_fwd( + *args, + kv_cache_dtype, + k_scale, + v_scale, + ) + else: + extension.grouped_sparse_page4_fwd(*args) reference_key = reference_key.view(4, 256) reference_value = reference_value.view(4, 256) diff --git a/tests/kernels/moe/test_sm70_unquantized_moe_config.py b/tests/kernels/moe/test_sm70_unquantized_moe_config.py index 533e2fdf21..1f05df95d8 100644 --- a/tests/kernels/moe/test_sm70_unquantized_moe_config.py +++ b/tests/kernels/moe/test_sm70_unquantized_moe_config.py @@ -8,6 +8,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import ( _get_sm70_mtp_moe_decode_config, force_sm70_mtp_moe_legacy_config, + fused_moe_kernel, ) @@ -20,6 +21,10 @@ def test_mtp_sm70_decode_config_keeps_legacy_tile_at_m1(): assert _get_sm70_mtp_moe_decode_config(1, 256, 128, 2048, 8) is None +def test_fused_moe_does_not_specialize_on_routing_dependent_em_alignment(): + assert "EM" in fused_moe_kernel.do_not_specialize_on_alignment + + @pytest.mark.parametrize("m", range(2, 17)) def test_mtp_sm70_decode_config_uses_exact_local_tile(m): config = _get_sm70_mtp_moe_decode_config(m, 256, 128, 2048, 8) diff --git a/tests/models/qwen4_exp/test_qsa_ops.py b/tests/models/qwen4_exp/test_qsa_ops.py index b1e8651d8f..9a95dd9e05 100644 --- a/tests/models/qwen4_exp/test_qsa_ops.py +++ b/tests/models/qwen4_exp/test_qsa_ops.py @@ -162,6 +162,66 @@ def test_qsa_xqa_page4_route_uses_configured_boundary(monkeypatch): assert not qsa_ops._use_sm70_qsa_xqa_page4(query, *args) +def test_qsa_grouped_page4_modern_abi_forwards_quantized_kv_metadata(): + calls = [] + + def forward(*args): + calls.append(args) + + extension = SimpleNamespace( + grouped_sparse_page4_abi_version=lambda: 2, + grouped_sparse_page4_plan_fwd=lambda *args: None, + grouped_sparse_page4_fwd=forward, + ) + tensors = [torch.empty(0) for _ in range(8)] + + assert qsa_ops._qsa_grouped_page4_supported(extension, "auto") + assert qsa_ops._qsa_grouped_page4_supported(extension, "fp8_e4m3") + qsa_ops._qsa_grouped_page4_forward( + extension, + *tensors, + 0.0625, + "fp8_e4m3", + 0.125, + 0.25, + ) + + assert len(calls) == 1 + assert len(calls[0]) == 12 + assert calls[0][-3:] == ("fp8_e4m3", 0.125, 0.25) + + +def test_qsa_grouped_page4_legacy_abi_is_fp16_only(): + calls = [] + + def forward(*args): + calls.append(args) + + forward.__doc__ = "grouped_sparse_page4_fwd(" + ", ".join( + f"arg{index}: object" for index in range(9) + ) + extension = SimpleNamespace( + grouped_sparse_page4_plan_fwd=lambda *args: None, + grouped_sparse_page4_fwd=forward, + ) + tensors = [torch.empty(0) for _ in range(8)] + + assert qsa_ops._qsa_grouped_page4_abi_version(extension) == 1 + assert qsa_ops._qsa_grouped_page4_supported(extension, "auto") + assert not qsa_ops._qsa_grouped_page4_supported(extension, "fp8_e4m3") + qsa_ops._qsa_grouped_page4_forward( + extension, + *tensors, + 0.0625, + "auto", + 1.0, + 1.0, + ) + + assert len(calls) == 1 + assert len(calls[0]) == 9 + + def test_qsa_e4m3_page4_routes_large_mixed_batch_below_prefill_boundary( monkeypatch, ): @@ -194,13 +254,21 @@ def test_qsa_e4m3_page4_routes_large_mixed_batch_below_prefill_boundary( ) -def test_qsa_e4m3_xqa_page4_splits_non_grouped_large_batch(monkeypatch): - rows = 49 +@pytest.mark.parametrize( + ("rows", "kv_cache_dtype"), + [(49, "fp8_e4m3"), (65, "auto")], +) +def test_qsa_xqa_page4_splits_non_grouped_large_batch( + monkeypatch, + rows, + kv_cache_dtype, +): query = torch.empty(rows, 6, 256, dtype=torch.float16) flash_cuda = SimpleNamespace( decode_paged_xqa_fwd=object(), - grouped_sparse_page4_plan_fwd=object(), - grouped_sparse_page4_fwd=object(), + grouped_sparse_page4_abi_version=lambda: 2, + grouped_sparse_page4_plan_fwd=lambda *args: None, + grouped_sparse_page4_fwd=lambda *args: None, ) flash_interface = ModuleType("flash_attn_v100.flash_attn_interface") cast(Any, flash_interface).flash_attn_v100_cuda = flash_cuda @@ -292,15 +360,23 @@ def fake_xqa_batch( query_positions, sequence_lengths, out, - "fp8_e4m3", + kv_cache_dtype, 0.05, 0.05, ) assert result is out + grouped_rows = rows // 8 * 8 assert calls == [ - ("grouped", 48, 48, 48, 48, 48), - ("xqa", 1, 1, 1, 1, 1), + ( + "grouped", + grouped_rows, + grouped_rows, + grouped_rows, + grouped_rows, + grouped_rows, + ), + ("xqa", *(rows - grouped_rows,) * 5), ] diff --git a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py index e351c60fa6..02e23680e2 100644 --- a/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py +++ b/tests/quantization/test_sm70_modelopt_mixed_nvfp4.py @@ -21,6 +21,7 @@ _prepare_compact_slot_groups, _prepare_single_token_slots, _single_token_weighted_reduce, + _use_compact_grouped, _use_qwen38_indexed_prefill, _use_qwen38_qpn_m1_decode, _validate_weight_layout, @@ -28,6 +29,15 @@ ) +@pytest.mark.parametrize( + ("top_k", "compact_tokens", "dense_tokens"), + [(8, 10, 11), (10, 8, 9)], +) +def test_nvfp4_compact_work_limit_is_routed_rows(top_k, compact_tokens, dense_tokens): + assert _use_compact_grouped(compact_tokens, top_k) + assert not _use_compact_grouped(dense_tokens, top_k) + + def _mixed_config() -> ModelOptMixedPrecisionConfig: fp8 = ModelOptFp8Config("FP8", True, None, []) nvfp4 = ModelOptNvFp4Config( @@ -249,7 +259,7 @@ def test_nvfp4_sm70_moe_owns_routing_without_generic_modular_wrapper(): not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires an exact SM70 CUDA device", ) -@pytest.mark.parametrize("total_slots", (8, 72, 80, 100)) +@pytest.mark.parametrize("total_slots", (8, 72, 80)) def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slots): sorted_expert_ids = ( torch.arange(total_slots, dtype=torch.int32, device="cuda") // 3 @@ -266,6 +276,18 @@ def test_nvfp4_compact_groups_keep_duplicate_expert_slots_independent(total_slot assert torch.equal(active_expert_ids.cpu(), sorted_expert_ids.cpu()) +@pytest.mark.parametrize("total_slots", (81, 100)) +def test_nvfp4_compact_groups_reject_work_above_80_rows(total_slots): + sorted_expert_ids = torch.empty(total_slots, dtype=torch.int32) + compact_offsets = torch.empty(total_slots + 1, dtype=torch.int32) + active_expert_ids = torch.empty(total_slots, dtype=torch.int32) + + with pytest.raises(ValueError, match="active-expert slots"): + _prepare_compact_slot_groups( + sorted_expert_ids, compact_offsets, active_expert_ids + ) + + @pytest.mark.skipif( not torch.cuda.is_available() or torch.cuda.get_device_capability() != (7, 0), reason="requires an exact SM70 CUDA device", diff --git a/tests/quantization/test_sm70_warmup.py b/tests/quantization/test_sm70_warmup.py index 9e7bdb3424..b59b67bd0f 100644 --- a/tests/quantization/test_sm70_warmup.py +++ b/tests/quantization/test_sm70_warmup.py @@ -278,7 +278,7 @@ def _nvfp4_moe_layer() -> nn.Module: layer.sm70_nvfp4_w2_n_dim = 2048 layer.sm70_nvfp4_group_size = 16 layer.sm70_nvfp4_graph_safe_max_tokens = 18 - layer.sm70_nvfp4_compact_grouped_max_tokens = 10 + layer.sm70_nvfp4_compact_grouped_max_slots = 80 layer.w13_tm_weight = nn.Parameter( torch.empty((1, 1), dtype=torch.uint8), requires_grad=False ) @@ -366,7 +366,7 @@ def test_nvfp4_moe_warmup_includes_opted_in_cuda_graph_shapes(monkeypatch): ) == [*range(1, 11), 18, 20, 40, 60, 80] -def test_nvfp4_moe_warmup_uses_slot_compact_through_b10(monkeypatch): +def test_nvfp4_moe_warmup_uses_slot_compact_through_80_rows(monkeypatch): layer = _nvfp4_moe_layer() calls = [] monkeypatch.setattr( @@ -400,7 +400,7 @@ def test_nvfp4_moe_warmup_uses_slot_compact_through_b10(monkeypatch): assert calls[2][2].tolist() == list(range(81)) -def test_nvfp4_moe_warmup_uses_full_expert_groups_above_compact_b10(monkeypatch): +def test_nvfp4_moe_warmup_uses_full_expert_groups_above_80_rows(monkeypatch): layer = _nvfp4_moe_layer() calls = [] monkeypatch.setattr( diff --git a/tests/v1/attention/test_sm70_flash_v100_policy.py b/tests/v1/attention/test_sm70_flash_v100_policy.py index 58e322653e..46e1f49fd2 100644 --- a/tests/v1/attention/test_sm70_flash_v100_policy.py +++ b/tests/v1/attention/test_sm70_flash_v100_policy.py @@ -1581,11 +1581,11 @@ def test_sm70_nomtp_cudagraph_capture_sizes_cover_concurrency( [ (1, [5]), (2, [5, 10]), - (4, [5, 10, 20]), - (6, [5, 10, 20, 30]), - (12, [5, 10, 20, 30, 40, 60]), - (16, [5, 10, 20, 30, 40, 60, 80]), - (32, [5, 10, 20, 30, 40, 60, 80]), + (4, [5, 10, 15, 20]), + (6, [5, 10, 15, 20, 30]), + (12, [5, 10, 15, 20, 30, 40, 60]), + (16, [5, 10, 15, 20, 30, 40, 60, 80]), + (32, [5, 10, 15, 20, 30, 40, 60, 80]), ], ) def test_sm70_mtp_cudagraph_capture_sizes_cover_production_concurrency( @@ -1597,6 +1597,24 @@ def test_sm70_mtp_cudagraph_capture_sizes_cover_production_concurrency( assert _sm70_mtp_cudagraph_capture_sizes(max_num_seqs, 5) == expected +def test_sm70_speculative_cudagraph_shapes_are_tp_independent_and_bounded(): + from vllm.config.vllm import _sm70_speculative_cudagraph_capture_sizes + + assert _sm70_speculative_cudagraph_capture_sizes(4, 5) == [ + 1, + 2, + 4, + 5, + 8, + 9, + 10, + 15, + 18, + 20, + ] + assert _sm70_speculative_cudagraph_capture_sizes(256, 5)[-1] == 80 + + def test_flash_v100_decode_query_does_not_attach_smallq_metadata( monkeypatch, local_flash_v100_model, diff --git a/tests/v1/executor/test_executor.py b/tests/v1/executor/test_executor.py index 494e8aa67d..f7201cb31b 100644 --- a/tests/v1/executor/test_executor.py +++ b/tests/v1/executor/test_executor.py @@ -6,6 +6,7 @@ from collections.abc import Callable from concurrent.futures import Future from typing import Any +from unittest.mock import MagicMock import pytest @@ -14,6 +15,7 @@ from vllm.sampling_params import SamplingParams from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.llm_engine import LLMEngine +from vllm.v1.executor import uniproc_executor as uniproc_executor_module from vllm.v1.executor.abstract import Executor from vllm.v1.executor.multiproc_executor import MultiprocExecutor from vllm.v1.executor.uniproc_executor import ( @@ -43,6 +45,42 @@ def test_supports_async_scheduling_multiproc_executor(): assert MultiprocExecutor.supports_async_scheduling() is True +def test_uniproc_executor_starts_ple_worker_around_model_load(monkeypatch): + """TP1 must not silently wait for a PLE process that was never spawned.""" + driver_worker = MagicMock() + monkeypatch.setattr( + uniproc_executor_module, + "WorkerWrapperBase", + MagicMock(return_value=driver_worker), + ) + monkeypatch.setattr(uniproc_executor_module.envs, "VLLM_PLE_CPU_OFFLOAD", True) + monkeypatch.setattr( + uniproc_executor_module.envs, + "VLLM_ELASTIC_EP_SCALE_UP_LAUNCH", + False, + ) + monkeypatch.setattr( + uniproc_executor_module, + "set_worker_net_device", + lambda *args: None, + ) + monkeypatch.setattr(uniproc_executor_module, "current_platform", MagicMock()) + + executor = UniProcExecutor.__new__(UniProcExecutor) + executor.vllm_config = object() + monkeypatch.setattr(executor, "_distributed_args", lambda: ("local://", 0, 0)) + + executor._init_executor() + + assert [call[0] for call in driver_worker.method_calls] == [ + "init_worker", + "init_device", + "spawn_ple_offload", + "load_model", + "wait_ple_offload_ready", + ] + + class CustomMultiprocExecutor(MultiprocExecutor): def collective_rpc( self, diff --git a/tests/v1/worker/test_gpu_warmup_blocks.py b/tests/v1/worker/test_gpu_warmup_blocks.py index 4beda08b46..3e84c9f2c9 100644 --- a/tests/v1/worker/test_gpu_warmup_blocks.py +++ b/tests/v1/worker/test_gpu_warmup_blocks.py @@ -2,6 +2,9 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """MRV2 warmup must reserve the same speculative KV tail as the scheduler.""" +from types import SimpleNamespace +from typing import Any + import pytest import torch @@ -13,13 +16,83 @@ MambaSpec, UniformTypeKVCacheSpecs, ) -from vllm.v1.worker.gpu.warmup import _reserved_block_count +from vllm.v1.worker.gpu.warmup import ( + _kernel_prefill_warmup_token_counts, + _reserved_block_count, + warmup_kernels, +) BLOCK_SIZE = 16 MAX_MODEL_LEN = 1024 NUM_SPEC_TOKENS = 7 +def test_kernel_prefill_warmup_profiles_are_capability_advertised() -> None: + runner = SimpleNamespace( + scheduler_config=SimpleNamespace(max_num_batched_tokens=128), + max_model_len=64, + compilation_config=SimpleNamespace( + static_forward_context={ + "plain": object(), + "qsa_0": SimpleNamespace( + kernel_warmup_prefill_token_counts=(33, 257, True, -1) + ), + "qsa_1": SimpleNamespace(kernel_warmup_prefill_token_counts=(33,)), + } + ), + ) + + assert _kernel_prefill_warmup_token_counts(runner, 6) == (6, 33) + + +def test_kernel_prefill_warmup_runs_default_batch_and_extra_profile( + monkeypatch: pytest.MonkeyPatch, +) -> None: + connector_states: list[bool] = [] + attention_spec = _full_attention_spec() + runner = SimpleNamespace( + num_speculative_steps=4, + scheduler_config=SimpleNamespace( + max_num_seqs=4, + max_num_batched_tokens=128, + ), + max_model_len=64, + compilation_config=SimpleNamespace( + static_forward_context={ + "profile": SimpleNamespace(kernel_warmup_prefill_token_counts=(33,)) + } + ), + kv_cache_config=SimpleNamespace( + kv_cache_groups=[SimpleNamespace(kv_cache_spec=attention_spec)], + num_blocks=64, + ), + vllm_config=SimpleNamespace(num_lookahead_tokens=4), + model_state=SimpleNamespace(max_encoder_len=0), + is_pooling_model=False, + is_last_pp_rank=True, + model_config=SimpleNamespace(get_vocab_size=lambda: 64), + kv_connector=SimpleNamespace( + set_disabled=lambda disabled: connector_states.append(disabled) + ), + ) + executions: list[Any] = [] + samples: list[Any] = [] + monkeypatch.setattr(torch.accelerator, "synchronize", lambda: None) + + warmup_kernels(runner, executions.append, samples.append) + + assert [output.total_num_scheduled_tokens for output in executions] == [ + 24, + 20, + 0, + 33, + 5, + 0, + ] + assert connector_states == [True, False] + assert len(samples) == 4 + + def _speculative_config(method: str) -> SpeculativeConfig: config = object.__new__(SpeculativeConfig) object.__setattr__(config, "method", method) diff --git a/tests/v1/worker/test_ple_offload_worker.py b/tests/v1/worker/test_ple_offload_worker.py index 6dc55e8c3c..473de34ddf 100644 --- a/tests/v1/worker/test_ple_offload_worker.py +++ b/tests/v1/worker/test_ple_offload_worker.py @@ -1,20 +1,24 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import queue from contextlib import nullcontext from types import SimpleNamespace +from unittest.mock import Mock import msgspec import pytest import torch import vllm.envs as envs +import vllm.v1.ple_offload.connector as ple_offload_connector_module import vllm.v1.worker.gpu_worker as gpu_worker_module from vllm.config import VllmConfig, get_current_vllm_config_or_none from vllm.model_executor.layers import ple_offload_layer from vllm.model_executor.layers.ple_offload_layer import PleOffloadLayer from vllm.model_executor.models.utils import AutoWeightsLoader, WeightsMapper from vllm.v1.ple_offload import worker as ple_offload_worker +from vllm.v1.ple_offload.connector import PleOffloadConnector from vllm.v1.worker.gpu_worker import Worker @@ -143,6 +147,292 @@ def test_ple_offload_rejects_missing_materialized_parameters( ) +def test_ple_storage_estimate_deduplicates_shared_storage() -> None: + module = torch.nn.Module() + weight = torch.nn.Parameter(torch.zeros(16, dtype=torch.float32)) + module.register_parameter("weight", weight) + module.register_buffer("weight_view", weight.detach().view(4, 4)) + + estimated = ple_offload_worker._estimate_module_storage_bytes([module]) + + assert estimated == weight.untyped_storage().nbytes() + + +def test_ple_auto_numa_uses_gpu_local_allowed_cpus( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls: list[str | set[int]] = [] + + def get_mempolicy(mode, *_): + mode._obj.value = 4 # MPOL_LOCAL + return 0 + + fake_libnuma = SimpleNamespace( + numa_available=lambda: 0, + numa_set_localalloc=lambda: calls.append("localalloc"), + get_mempolicy=get_mempolicy, + ) + monkeypatch.setattr(envs, "VLLM_PLE_OFFLOAD_AUTO_NUMA", True) + monkeypatch.setattr( + ple_offload_worker, + "os", + SimpleNamespace( + sched_getaffinity=lambda _: {0, 1, 24, 25}, + sched_setaffinity=lambda _, cpus: calls.append(set(cpus)), + ), + ) + monkeypatch.setattr( + ple_offload_worker.psutil, + "Process", + lambda *_: SimpleNamespace(), + ) + monkeypatch.setattr( + "vllm.platforms.current_platform.get_all_device_numa_nodes", + lambda: [1, 1, 1, 1], + ) + monkeypatch.setattr( + "vllm.utils.cpu_resource_utils.get_allowed_cpu_list", + lambda: [ + SimpleNamespace(id=0, numa_node=0), + SimpleNamespace(id=1, numa_node=0), + SimpleNamespace(id=24, numa_node=1), + SimpleNamespace(id=25, numa_node=1), + ], + ) + monkeypatch.setattr("vllm.utils.numa_utils.get_libnuma", lambda: fake_libnuma) + config = SimpleNamespace( + parallel_config=SimpleNamespace(numa_bind_nodes=None), + ) + + node = ple_offload_worker._configure_ple_numa_locality(config) + + assert node == 1 + assert calls == ["localalloc", {24, 25}] + + +def test_ple_auto_numa_keeps_cpu_affinity_when_mempolicy_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + calls = [] + + def fail_localalloc() -> None: + raise PermissionError("set_mempolicy denied") + + fake_libnuma = SimpleNamespace( + numa_available=lambda: 0, + numa_set_localalloc=fail_localalloc, + ) + monkeypatch.setattr(envs, "VLLM_PLE_OFFLOAD_AUTO_NUMA", True) + monkeypatch.setattr( + ple_offload_worker, + "os", + SimpleNamespace( + sched_getaffinity=lambda _: {0, 1}, + sched_setaffinity=lambda _, cpus: calls.append(set(cpus)), + ), + ) + monkeypatch.setattr( + "vllm.platforms.current_platform.get_all_device_numa_nodes", + lambda: [0], + ) + monkeypatch.setattr( + "vllm.utils.cpu_resource_utils.get_allowed_cpu_list", + lambda: [ + SimpleNamespace(id=0, numa_node=0), + SimpleNamespace(id=1, numa_node=0), + ], + ) + monkeypatch.setattr("vllm.utils.numa_utils.get_libnuma", lambda: fake_libnuma) + config = SimpleNamespace( + parallel_config=SimpleNamespace(numa_bind_nodes=None), + ) + + assert ple_offload_worker._configure_ple_numa_locality(config) == 0 + assert calls == [{0, 1}] + + +def test_ple_prefault_touches_unique_storage_when_capacity_allows( + monkeypatch: pytest.MonkeyPatch, +) -> None: + module = torch.nn.Module() + weight = torch.nn.Parameter(torch.arange(8192, dtype=torch.float32)) + module.register_parameter("weight", weight) + module.register_buffer("weight_view", weight.detach().view(4096, 2)) + required = weight.untyped_storage().nbytes() + monkeypatch.setattr(envs, "VLLM_PLE_OFFLOAD_PREFAULT", True) + monkeypatch.setattr(ple_offload_worker.os, "sysconf", lambda _: 4096) + monkeypatch.setattr( + ple_offload_worker.psutil, + "virtual_memory", + lambda: SimpleNamespace( + total=64 * ple_offload_worker.GiB_bytes, + available=32 * ple_offload_worker.GiB_bytes, + ), + ) + monkeypatch.setattr( + ple_offload_worker.psutil, + "Process", + lambda _: SimpleNamespace( + memory_full_info=lambda: SimpleNamespace(rss=required, swap=0), + ), + ) + + assert ple_offload_worker._prefault_module_storage([module]) == required + + +def test_ple_prefault_skips_when_host_capacity_is_insufficient( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + module = torch.nn.Linear(4, 4, bias=False) + required = module.weight.untyped_storage().nbytes() + gib = ple_offload_worker.GiB_bytes + monkeypatch.setattr(envs, "VLLM_PLE_OFFLOAD_PREFAULT", True) + monkeypatch.setattr( + ple_offload_worker.psutil, + "virtual_memory", + lambda: SimpleNamespace(total=16 * gib, available=1), + ) + monkeypatch.setattr( + ple_offload_worker.psutil, + "Process", + lambda _: SimpleNamespace( + memory_full_info=lambda: SimpleNamespace(rss=0, swap=0) + ), + ) + + with caplog.at_level("WARNING", logger=ple_offload_worker.__name__): + touched = ple_offload_worker._prefault_module_storage([module]) + + assert touched == 0 + assert required > 0 + assert "Skipping PLE RAM prefault" in caplog.text + + +def test_ple_host_memory_pressure_warns_without_rejecting( + monkeypatch: pytest.MonkeyPatch, + caplog: pytest.LogCaptureFixture, +) -> None: + gib = ple_offload_worker.GiB_bytes + monkeypatch.setattr( + ple_offload_worker, + "_estimate_module_storage_bytes", + lambda _: 6 * gib, + ) + monkeypatch.setattr( + ple_offload_worker.psutil, + "virtual_memory", + lambda: SimpleNamespace(total=16 * gib, available=8 * gib), + ) + monkeypatch.setattr( + ple_offload_worker.psutil, + "swap_memory", + lambda: SimpleNamespace(used=3 * gib), + ) + + with caplog.at_level("WARNING", logger=ple_offload_worker.__name__): + required = ple_offload_worker._log_ple_host_memory_capacity([]) + + assert required == 6 * gib + assert "may page its 6.00 GiB table" in caplog.text + assert "input-dependent prefill and decode latency" in caplog.text + + +def test_ple_offload_uses_event_per_inflight_mrv2_batch() -> None: + """A later batch must not retarget an earlier batch's D2H event.""" + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector.tp_rank = 0 + connector.dp_rank = 0 + connector._uses_cuda_inputs = True + connector._request_queue = queue.Queue(maxsize=2) + connector._d2h_event_pool = queue.Queue(maxsize=2) + first_event = Mock() + second_event = Mock() + connector._d2h_event_pool.put_nowait(first_event) + connector._d2h_event_pool.put_nowait(second_event) + enqueue_cuda_inputs = Mock() + connector._enqueue_cuda_inputs = enqueue_cuda_inputs # type: ignore[method-assign] + + connector._launch(num_reqs=4, num_tokens=20) + connector._launch(num_reqs=1, num_tokens=5) + + first_pending = connector._request_queue.get_nowait() + second_pending = connector._request_queue.get_nowait() + assert first_pending is not None + assert second_pending is not None + assert first_pending.d2h_done_event is first_event + assert second_pending.d2h_done_event is second_event + assert first_pending.request.num_reqs == 4 + assert first_pending.request.num_tokens == 20 + assert second_pending.request.num_reqs == 1 + assert second_pending.request.num_tokens == 5 + assert enqueue_cuda_inputs.call_args_list[0].args[1] is first_event + assert enqueue_cuda_inputs.call_args_list[1].args[1] is second_event + assert connector._d2h_event_pool.empty() + + with pytest.raises(RuntimeError, match="configured concurrent batches"): + connector._launch(num_reqs=1, num_tokens=1) + + +def test_ple_offload_request_waits_for_its_bound_event( + monkeypatch: pytest.MonkeyPatch, +) -> None: + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector.device = SimpleNamespace(index=0) + connector._uses_cuda_inputs = True + connector._d2h_event_pool = queue.Queue(maxsize=1) + event = Mock() + socket = Mock() + request = ple_offload_worker.PleOffloadRequest( + dp_rank=0, + num_tokens=5, + num_reqs=1, + ) + pending = ple_offload_connector_module._PendingPleOffloadRequest(request, event) + monkeypatch.setattr( + ple_offload_connector_module.torch.accelerator, + "device_index", + lambda *_: nullcontext(), + ) + monkeypatch.setattr( + ple_offload_connector_module.torch.cuda.nvtx, + "range", + lambda *_: nullcontext(), + ) + + connector._process_request(pending, socket) + + event.synchronize.assert_called_once_with() + socket.send.assert_called_once_with(msgspec.msgpack.encode(request)) + assert connector._d2h_event_pool.get_nowait() is event + + +def test_ple_offload_preserves_mrv1_cpu_staging( + monkeypatch: pytest.MonkeyPatch, +) -> None: + connector = PleOffloadConnector.__new__(PleOffloadConnector) + connector._uses_cuda_inputs = False + connector._d2h_event_pool = None + connector._copy_cpu_inputs = Mock() # type: ignore[method-assign] + socket = Mock() + request = ple_offload_worker.PleOffloadRequest( + dp_rank=0, + num_tokens=3, + num_reqs=1, + ) + pending = ple_offload_connector_module._PendingPleOffloadRequest(request, None) + monkeypatch.setattr( + ple_offload_connector_module.torch.cuda.nvtx, + "range", + lambda *_: nullcontext(), + ) + + connector._process_request(pending, socket) + + connector._copy_cpu_inputs.assert_called_once_with(request) + socket.send.assert_called_once_with(msgspec.msgpack.encode(request)) + + def test_ple_offload_wait_only_waits_for_done( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/tests/v1/worker/test_qwen4_exp_v2.py b/tests/v1/worker/test_qwen4_exp_v2.py index 954779157f..468b23454b 100644 --- a/tests/v1/worker/test_qwen4_exp_v2.py +++ b/tests/v1/worker/test_qwen4_exp_v2.py @@ -3,10 +3,12 @@ from contextlib import nullcontext from types import SimpleNamespace +from unittest import mock import pytest import torch +from vllm.models.qwen4_exp.nvidia.ops import qsa as qsa_ops from vllm.v1.kv_cache_interface import ( CircularBufferSpec, FullAttentionSpec, @@ -20,6 +22,89 @@ from vllm.v1.worker.gpu.spec_decode.eagle import speculator as eagle_speculator +def test_qsa_request_count_is_not_a_triton_compile_key() -> None: + for kernel in ( + qsa_ops._qsa_mqa_paged_kernel, + qsa_ops._qsa_sparse_paged_gqa_splitk_kernel, + ): + assert "num_requests" in kernel.do_not_specialize + + +@pytest.mark.parametrize( + ("capture_sizes", "decode_query_len", "max_num_reqs", "expected"), + [ + ([1, 2, 4, 5, 8, 9, 10, 15, 18, 20], 5, 4, (1, 2, 3, 4)), + ([5, 10, 15, 20, 25], 5, 3, (1, 2, 3)), + ([1, 2, 4, 8], 5, 4, (1,)), + (None, 5, 4, (1,)), + ([5], 0, 4, ()), + ([5], 5, 0, ()), + ], +) +def test_mtp_decode_warmup_uses_graph_request_shapes( + capture_sizes: list[int] | None, + decode_query_len: int, + max_num_reqs: int, + expected: tuple[int, ...], +) -> None: + assert ( + eagle_speculator._mtp_decode_warmup_request_sizes( + capture_sizes, + decode_query_len, + max_num_reqs, + ) + == expected + ) + + +def test_mtp_moe_warmup_executes_each_captured_concurrency( + monkeypatch: pytest.MonkeyPatch, +) -> None: + speculator = eagle_speculator.EagleSpeculator.__new__( + eagle_speculator.EagleSpeculator + ) + speculator.method = "mtp" + speculator.device = torch.device("cuda") + speculator.num_speculative_steps = 4 + speculator.max_num_reqs = 4 + speculator.max_num_tokens = 64 + speculator._sm70_mtp_moe_warmed = False + speculator.draft_model_config = SimpleNamespace( + hf_text_config=SimpleNamespace( + num_experts_per_tok=10, + moe_intermediate_size=640, + ), + get_num_experts=lambda: 512, + get_hidden_size=lambda: 2560, + ) + speculator.vllm_config = SimpleNamespace( + parallel_config=SimpleNamespace(tensor_parallel_size=4), + compilation_config=SimpleNamespace( + cudagraph_capture_sizes=[1, 2, 4, 5, 8, 9, 10, 15, 18, 20] + ), + ) + dummy_run = mock.Mock() + monkeypatch.setattr( + eagle_speculator.current_platform, + "is_device_capability", + lambda *_args: True, + ) + monkeypatch.setattr(eagle_speculator.envs, "VLLM_SM70_MTP_MOE_TUNED_CONFIG", True) + monkeypatch.setattr(torch.accelerator, "synchronize", mock.Mock()) + + assert speculator.warmup_sm70_mtp_moe_kernels(dummy_run) == ( + "mtp_draft_moe_prefill_m16", + "mtp_draft_moe_decode_reqs_1_2_3_4", + ) + assert dummy_run.call_args_list == [ + mock.call(16), + mock.call(5, uniform_decode=True), + mock.call(10, uniform_decode=True), + mock.call(15, uniform_decode=True), + mock.call(20, uniform_decode=True), + ] + + def test_qwen4_exp_mtp_v2_unpacks_logits_and_feedback_hidden_states( monkeypatch: pytest.MonkeyPatch, ) -> None: diff --git a/vllm/config/vllm.py b/vllm/config/vllm.py index acdc15a1d7..7532ef36c0 100644 --- a/vllm/config/vllm.py +++ b/vllm/config/vllm.py @@ -81,7 +81,8 @@ } ) _SM70_NOMTP_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 16) -_SM70_MTP_CUDAGRAPH_REQUEST_SIZES = (1, 2, 4, 6, 8, 12, 16) +_SM70_MTP_CUDAGRAPH_REQUEST_SIZES = (1, 2, 3, 4, 6, 8, 12, 16) +_SM70_SPECULATIVE_AUX_CUDAGRAPH_CAPTURE_SIZES = (1, 2, 4, 8, 9, 18) _SM70_DFLASH2_VERIFIER_DEFAULTS = { # This is the target projection's memory-neutral FP8 layout, not the @@ -183,6 +184,20 @@ def _sm70_mtp_cudagraph_capture_sizes( return [decode_query_len * size for size in sorted(request_sizes)] +def _sm70_speculative_cudagraph_capture_sizes( + max_num_seqs: int, + decode_query_len: int, +) -> list[int]: + """Return bounded auxiliary and verifier shapes without a TP contract.""" + verifier_sizes = _sm70_mtp_cudagraph_capture_sizes( + max_num_seqs, + decode_query_len, + ) + return sorted( + set(_SM70_SPECULATIVE_AUX_CUDAGRAPH_CAPTURE_SIZES) | set(verifier_sizes) + ) + + class OptimizationLevel(IntEnum): """Optimization level enum.""" @@ -1682,31 +1697,15 @@ def __post_init__(self): ) else: cudagraph_capture_sizes = ( - [1, 2, 4, 8, 9, 18] - if self.parallel_config.tensor_parallel_size >= 4 - else [1, 2, 4, 8, 9] - ) - max_graph_reqs = ( - 4 - if self.parallel_config.tensor_parallel_size >= 4 - else 1 - ) - max_graph_reqs = min( - max(int(self.scheduler_config.max_num_seqs), 1), - max_graph_reqs, - ) - cudagraph_capture_sizes = sorted( - set(cudagraph_capture_sizes) - | { - decode_query_len * num_reqs - for num_reqs in range(1, max_graph_reqs + 1) - } + _sm70_speculative_cudagraph_capture_sizes( + self.scheduler_config.max_num_seqs, + decode_query_len, + ) ) logger.info_once( - "Using SM70 speculative verifier cudagraph shapes " - "%sx1..%s for Flash-V100 compile graph.", - decode_query_len, - max_graph_reqs, + "Using bounded SM70 speculative cudagraph token " + "shapes %s for Flash-V100 compile graph.", + tuple(cudagraph_capture_sizes), ) elif cudagraph_capture_sizes != [1, 2]: logger.info_once( diff --git a/vllm/engine/arg_utils.py b/vllm/engine/arg_utils.py index 16cf09c17c..6b73e33f00 100644 --- a/vllm/engine/arg_utils.py +++ b/vllm/engine/arg_utils.py @@ -97,6 +97,7 @@ OptimizationLevel, PerformanceMode, _sm70_mtp_cudagraph_capture_sizes, + _sm70_speculative_cudagraph_capture_sizes, ) from vllm.logger import init_logger, suppress_logging from vllm.platforms import CpuArchEnum, current_platform @@ -1852,10 +1853,6 @@ def _maybe_apply_sm70_mtp_defaults( self.speculative_config["use_local_argmax_reduction"] = True profile_updates.append("speculative_config.use_local_argmax_reduction=True") - if self.max_num_seqs is None: - self.max_num_seqs = 4 if self.tensor_parallel_size >= 4 else 1 - profile_updates.append(f"max_num_seqs={self.max_num_seqs}") - if has_native_mtp and self.compilation_config.fast_moe_cold_start is None: # Native MTP layers are encoded explicitly by MoERunner and do not # consume the target-model forward-context MoE order. @@ -1885,40 +1882,38 @@ def _maybe_apply_sm70_mtp_defaults( num_speculative_tokens, ddtree_budget ) decode_query_len = num_speculative_state_tokens + 1 - max_num_seqs = max(int(self.max_num_seqs or 1), 1) - if envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS and spec_method == "mtp": + if self.max_num_seqs is None: + # Scheduler defaults are resolved later in VllmConfig. Leave + # graph sizing to that stage instead of turning a graph policy + # into a service-capacity limit. + profile_updates.append("mtp_cudagraph_shapes=deferred_to_scheduler") + elif envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS and spec_method == "mtp": cudagraph_capture_sizes = _sm70_mtp_cudagraph_capture_sizes( - max_num_seqs, + self.max_num_seqs, decode_query_len, ) profile_updates.append( f"mtp_split_verifier_cudagraph_shapes={cudagraph_capture_sizes}" ) else: - cudagraph_capture_sizes = ( - [1, 2, 4, 8, 9, 18] - if self.tensor_parallel_size >= 4 - else [1, 2, 4, 8, 9] + cudagraph_capture_sizes = _sm70_speculative_cudagraph_capture_sizes( + self.max_num_seqs, + decode_query_len, + ) + profile_updates.append( + f"mtp_cudagraph_shapes={cudagraph_capture_sizes}" ) - cudagraph_capture_sizes = sorted( - set(cudagraph_capture_sizes) - | { - decode_query_len * num_reqs - for num_reqs in range(1, max_num_seqs + 1) - } + if self.max_num_seqs is not None: + self.compilation_config.cudagraph_capture_sizes = ( + cudagraph_capture_sizes + ) + self.compilation_config.max_cudagraph_capture_size = max( + cudagraph_capture_sizes ) profile_updates.append( - "mtp_verifier_cudagraph_shapes=" - f"{decode_query_len}x1..{max_num_seqs}" + "cudagraph_capture_sizes=" + f"{self.compilation_config.cudagraph_capture_sizes}" ) - self.compilation_config.cudagraph_capture_sizes = cudagraph_capture_sizes - self.compilation_config.max_cudagraph_capture_size = max( - cudagraph_capture_sizes - ) - profile_updates.append( - "cudagraph_capture_sizes=" - f"{self.compilation_config.cudagraph_capture_sizes}" - ) if profile_updates: logger.info_once( diff --git a/vllm/envs.py b/vllm/envs.py index a47ed2e4d1..f7ca6f836f 100755 --- a/vllm/envs.py +++ b/vllm/envs.py @@ -729,6 +729,8 @@ VLLM_COMPILE_CACHE_SAVE_FORMAT: Literal["binary", "unpacked"] = "binary" VLLM_USE_V2_MODEL_RUNNER: bool | None = None VLLM_PLE_CPU_OFFLOAD: bool = False + VLLM_PLE_OFFLOAD_AUTO_NUMA: bool = True + VLLM_PLE_OFFLOAD_PREFAULT: bool = True VLLM_PLE_OFFLOAD_READY_TIMEOUT: float = 600.0 VLLM_LOG_MODEL_INSPECTION: bool = False VLLM_DEBUG_MFU_METRICS: bool = False @@ -4272,6 +4274,19 @@ def _resolve_rust_frontend_path() -> str | None: "VLLM_PLE_CPU_OFFLOAD": lambda: ( os.getenv("VLLM_PLE_CPU_OFFLOAD", "False").lower() in ("true", "1") ), + # Keep the latency-critical PLE lookup process on the NUMA node local to + # its first visible GPU. This changes CPU placement only; allocations use + # a local-first policy with fallback so large tables are not forced into a + # single NUMA node and swapped out. + "VLLM_PLE_OFFLOAD_AUTO_NUMA": lambda: ( + os.getenv("VLLM_PLE_OFFLOAD_AUTO_NUMA", "True").lower() in ("true", "1") + ), + # Fault PLE table pages back into RAM after GPU workers finish loading. + # Concurrent checkpoint loading can otherwise leave anonymous table pages + # in swap while reclaimable checkpoint page cache occupies host memory. + "VLLM_PLE_OFFLOAD_PREFAULT": lambda: ( + os.getenv("VLLM_PLE_OFFLOAD_PREFAULT", "True").lower() in ("true", "1") + ), # Timeout for PLE weight loading and TP worker registration. "VLLM_PLE_OFFLOAD_READY_TIMEOUT": lambda: float( os.getenv("VLLM_PLE_OFFLOAD_READY_TIMEOUT", "600") diff --git a/vllm/model_executor/layers/fused_moe/fused_moe.py b/vllm/model_executor/layers/fused_moe/fused_moe.py index e8eea2c04b..8b9b74625c 100644 --- a/vllm/model_executor/layers/fused_moe/fused_moe.py +++ b/vllm/model_executor/layers/fused_moe/fused_moe.py @@ -305,7 +305,7 @@ def fused_moe_kernel_gptq_awq( tl.store(c_ptrs, accumulator, mask=c_mask) -@triton.jit +@triton.jit(do_not_specialize_on_alignment=["EM"]) def fused_moe_kernel( # Pointers to matrices a_ptr, diff --git a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py index 3e4d36f326..d4afc2f852 100644 --- a/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py +++ b/vllm/model_executor/layers/quantization/nvfp4_sm70_moe.py @@ -48,8 +48,7 @@ } _SUPPORTED_TP_SIZES: Final = (1, 2, 4) _GRAPH_SAFE_MAX_TOKENS: Final = 18 -_COMPACT_GROUPED_MAX_TOKENS: Final = 10 -_MAX_SUPPORTED_TOP_K: Final = max(contract[3] for contract in _SUPPORTED_CONTRACTS) +_COMPACT_GROUPED_MAX_SLOTS: Final = 80 _QWEN38_QPN_M1_W13_SPLIT_K: Final = 8 _QWEN38_QPN_M1_W2_SPLIT_K: Final = 1 _QWEN38_INDEXED_PREFILL_MIN_TOKENS: Final = 128 @@ -223,8 +222,7 @@ def _prepare_compact_slot_groups( active_expert_ids: torch.Tensor, ) -> None: total_slots = sorted_expert_ids.numel() - max_slots = _COMPACT_GROUPED_MAX_TOKENS * _MAX_SUPPORTED_TOP_K - if not (0 < total_slots <= max_slots): + if not (0 < total_slots <= _COMPACT_GROUPED_MAX_SLOTS): raise ValueError(f"Unsupported SM70 NVFP4 active-expert slots: {total_slots}") block = triton.next_power_of_2(total_slots + 1) # TurboMind's compact grouped dispatch forces one row per group. Keep each @@ -241,6 +239,12 @@ def _prepare_compact_slot_groups( ) +def _use_compact_grouped(num_tokens: int, top_k: int) -> bool: + """Bound compact dispatch by its routed-row workload, not batch size.""" + total_slots = num_tokens * top_k + return 0 < total_slots <= _COMPACT_GROUPED_MAX_SLOTS + + def validate_nvfp4_sm70_moe_contract(moe: FusedMoEConfig) -> None: """Reject every topology outside the validated SM70 NVFP4 contract.""" if moe.tp_size not in _SUPPORTED_TP_SIZES: @@ -565,7 +569,7 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: layer.sm70_nvfp4_qwen38_fused_swiglu_prefill = fused_swiglu_prefill layer.sm70_nvfp4_qwen38_fast_prefill = fast_prefill layer.sm70_nvfp4_graph_safe_max_tokens = _GRAPH_SAFE_MAX_TOKENS - layer.sm70_nvfp4_compact_grouped_max_tokens = _COMPACT_GROUPED_MAX_TOKENS + layer.sm70_nvfp4_compact_grouped_max_slots = _COMPACT_GROUPED_MAX_SLOTS self._allocate_graph_safe_decode_buffers(layer) del layer.w13_weight @@ -579,13 +583,13 @@ def process_weights_after_loading(self, layer: RoutedExperts) -> None: logger.info_once( "SM70 ModelOpt NVFP4 TurboMind MoE path enabled " "(hidden=%d, local_intermediate=%d, local_experts=%d, top_k=%d, " - "graph_safe_decode=B1-B%d, compact_grouped_decode=B1-B%d).", + "graph_safe_decode=B1-B%d, compact_grouped_decode<=%d routed rows).", hidden, intermediate, num_experts, layer.sm70_nvfp4_top_k, _GRAPH_SAFE_MAX_TOKENS, - _COMPACT_GROUPED_MAX_TOKENS, + _COMPACT_GROUPED_MAX_SLOTS, ) if fused_swiglu_prefill: logger.info_once( @@ -931,7 +935,7 @@ def apply( buffers["expert_offsets64"], non_blocking=True ) - if not direct_single_token and num_tokens <= _COMPACT_GROUPED_MAX_TOKENS: + if not direct_single_token and _use_compact_grouped(num_tokens, top_k): _prepare_compact_slot_groups( buffers["permuted_experts_id"], buffers["compact_offsets"], diff --git a/vllm/model_executor/warmup/awq_sm70_warmup.py b/vllm/model_executor/warmup/awq_sm70_warmup.py index 0cd9d59a37..6e04c86f88 100644 --- a/vllm/model_executor/warmup/awq_sm70_warmup.py +++ b/vllm/model_executor/warmup/awq_sm70_warmup.py @@ -17,6 +17,7 @@ from vllm.model_executor.layers.quantization import sm70_turbomind as sm70_tm from vllm.model_executor.layers.quantization.nvfp4_sm70_moe import ( _prepare_compact_slot_groups, + _use_compact_grouped, ) if TYPE_CHECKING: @@ -744,12 +745,18 @@ def _warmup_nvfp4_moe_decode_layers( device = layer.w13_tm_weight.device top_k = int(layer.moe_config.experts_per_token) max_experts = int(layer.sm70_nvfp4_num_experts) - compact_max_tokens = int( - getattr(layer, "sm70_nvfp4_compact_grouped_max_tokens", 8) + compact_max_slots = int( + getattr( + layer, + "sm70_nvfp4_compact_grouped_max_slots", + 8 * top_k, + ) ) for num_tokens in token_counts: total_slots = num_tokens * top_k - if num_tokens <= compact_max_tokens: + if _use_compact_grouped(num_tokens, top_k) and ( + total_slots <= compact_max_slots + ): stage_experts = total_slots sorted_expert_ids = layer._nvfp4_sm70_dense_expert_ids[:total_slots] expert_offsets = torch.empty( @@ -840,7 +847,14 @@ def _get_nvfp4_moe_token_counts( for layer in moe_layers ) compact_max_tokens = max( - int(getattr(layer, "sm70_nvfp4_compact_grouped_max_tokens", 8)) + int( + getattr( + layer, + "sm70_nvfp4_compact_grouped_max_slots", + 8 * int(layer.moe_config.experts_per_token), + ) + ) + // max(int(layer.moe_config.experts_per_token), 1) for layer in moe_layers ) max_top_k = max(int(layer.moe_config.experts_per_token) for layer in moe_layers) diff --git a/vllm/models/qwen4_exp/nvidia/ops/qsa.py b/vllm/models/qwen4_exp/nvidia/ops/qsa.py index f00236eb0a..c020dc0f24 100644 --- a/vllm/models/qwen4_exp/nvidia/ops/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/ops/qsa.py @@ -7,6 +7,7 @@ import math import os +import regex as re import torch from vllm.logger import init_logger @@ -32,7 +33,10 @@ ) _SM70_QSA_XQA_PAGE4 = os.getenv("VLLM_SM70_QSA_XQA_PAGE4", "1") == "1" _SM70_QSA_XQA_PAGE4_MIN_ROWS = int( - os.getenv("VLLM_SM70_QSA_XQA_PAGE4_MIN_ROWS", "4096") + # Operator crossover on SM70 is around 48 rows for the fixed QSA width. + # Use a conservative 64-row workload gate rather than coupling the route + # to a particular server's max_num_batched_tokens setting. + os.getenv("VLLM_SM70_QSA_XQA_PAGE4_MIN_ROWS", "64") ) _SM70_QSA_XQA_PAGE4_PARTITION = 1024 _SM70_QSA_XQA_PAGE4_PAGES = 513 @@ -56,9 +60,60 @@ torch.Tensor, ], ] = {} +_SM70_QSA_GROUPED_PAGE4_ABI_CACHE: tuple[object, int] | None = None + + +def _qsa_grouped_page4_abi_version(flash_attn_v100_cuda) -> int: + """Return the grouped-page4 ABI without probing it on the hot path. + + New Flash-V100 builds expose an explicit version. Wheels predating that + capability query are recognized conservatively from pybind's generated + signature: ABI v1 has arguments 0..8, while ABI v2 has arguments 0..11. + An unknown binding is treated as unsupported instead of risking a server + crash on the first large prefill. + """ + global _SM70_QSA_GROUPED_PAGE4_ABI_CACHE + cached = _SM70_QSA_GROUPED_PAGE4_ABI_CACHE + if cached is not None and cached[0] is flash_attn_v100_cuda: + return cached[1] + + version = 0 + capability = getattr(flash_attn_v100_cuda, "grouped_sparse_page4_abi_version", None) + if callable(capability): + try: + version = int(capability()) + except (RuntimeError, TypeError, ValueError): + version = 0 + else: + binding = getattr(flash_attn_v100_cuda, "grouped_sparse_page4_fwd", None) + doc = getattr(binding, "__doc__", "") or "" + argument_ids = [int(match) for match in re.findall(r"\barg(\d+):", doc)] + if argument_ids: + highest_argument = max(argument_ids) + if highest_argument >= 11: + version = 2 + elif highest_argument >= 8: + version = 1 + _SM70_QSA_GROUPED_PAGE4_ABI_CACHE = (flash_attn_v100_cuda, version) + return version -@triton.jit + +def _qsa_grouped_page4_supported( + flash_attn_v100_cuda, + kv_cache_dtype: str, +) -> bool: + forward = getattr(flash_attn_v100_cuda, "grouped_sparse_page4_fwd", None) + planner = getattr(flash_attn_v100_cuda, "grouped_sparse_page4_plan_fwd", None) + if not callable(forward) or not callable(planner): + return False + abi_version = _qsa_grouped_page4_abi_version(flash_attn_v100_cuda) + return abi_version >= 2 or ( + abi_version == 1 and kv_cache_dtype in ("auto", "float16") + ) + + +@triton.jit(do_not_specialize=["num_requests"]) def _qsa_mqa_paged_kernel( q_ptr, k_cache_ptr, @@ -457,7 +512,7 @@ def _qsa_xqa_page4_table_kernel( ) -@triton.jit +@triton.jit(do_not_specialize=["num_requests"]) def _qsa_sparse_paged_gqa_splitk_kernel( q_ptr, k_cache_ptr, @@ -1564,6 +1619,48 @@ def _qsa_xqa_page4_physical_kv( return physical_k_cache, physical_v_cache +def _qsa_grouped_page4_forward( + flash_attn_v100_cuda, + q: torch.Tensor, + k_cache: torch.Tensor, + v_cache: torch.Tensor, + out: torch.Tensor, + grouped_pages: torch.Tensor, + token_masks: torch.Tensor, + grouped_sequence_lengths: torch.Tensor, + lse: torch.Tensor, + softmax_scale: float, + kv_cache_dtype: str, + k_scale: float, + v_scale: float, +) -> None: + forward_args = ( + q, + k_cache, + v_cache, + out, + grouped_pages, + token_masks, + grouped_sequence_lengths, + lse, + softmax_scale, + ) + abi_version = _qsa_grouped_page4_abi_version(flash_attn_v100_cuda) + if abi_version >= 2: + flash_attn_v100_cuda.grouped_sparse_page4_fwd( + *forward_args, + kv_cache_dtype, + k_scale, + v_scale, + ) + return + + # ABI v1 only supports the original FP16 K/V contract. The route + # eligibility check rejects quantized K/V before the planner runs. + assert abi_version == 1 and kv_cache_dtype in ("auto", "float16") + flash_attn_v100_cuda.grouped_sparse_page4_fwd(*forward_args) + + def _qsa_sparse_paged_attention_sm70_grouped_page4( q: torch.Tensor, k_cache: torch.Tensor, @@ -1597,7 +1694,8 @@ def _qsa_sparse_paged_attention_sm70_grouped_page4( k_cache.shape[0], ) physical_k_cache, physical_v_cache = _qsa_xqa_page4_physical_kv(q, k_cache, v_cache) - flash_attn_v100_cuda.grouped_sparse_page4_fwd( + _qsa_grouped_page4_forward( + flash_attn_v100_cuda, q, physical_k_cache, physical_v_cache, @@ -1712,61 +1810,70 @@ def _qsa_sparse_paged_attention_sm70_xqa_page4( ) return None - grouped_bindings_available = hasattr( - flash_attn_v100_cuda, "grouped_sparse_page4_plan_fwd" - ) and hasattr(flash_attn_v100_cuda, "grouped_sparse_page4_fwd") - grouped_enabled = _SM70_QSA_GROUPED_PAGE4 and grouped_bindings_available - if grouped_enabled and q.shape[0] % _SM70_QSA_GROUPED_PAGE4_QUERIES == 0: - return _qsa_sparse_paged_attention_sm70_grouped_page4( - q, + grouped_enabled = _SM70_QSA_GROUPED_PAGE4 and _qsa_grouped_page4_supported( + flash_attn_v100_cuda, kv_cache_dtype + ) + if grouped_enabled: + grouped_rows = ( + q.shape[0] // _SM70_QSA_GROUPED_PAGE4_QUERIES + ) * _SM70_QSA_GROUPED_PAGE4_QUERIES + if grouped_rows: + _qsa_sparse_paged_attention_sm70_grouped_page4( + q[:grouped_rows], + k_cache, + v_cache, + logical_indices[:grouped_rows], + block_table, + token_to_req[:grouped_rows], + query_positions[:grouped_rows], + sequence_lengths, + out[:grouped_rows], + kv_cache_dtype, + k_scale, + v_scale, + flash_attn_v100_cuda, + ) + if grouped_rows == q.shape[0]: + return out + + logger.info_once( + "Splitting a non-grouped page4 batch across grouped/XQA routes " + "(rows=%d, grouped_rows=%d, kv_cache_dtype=%s).", + q.shape[0], + grouped_rows, + kv_cache_dtype, + ) + _qsa_sparse_paged_attention_sm70_xqa_page4_batch( + q[grouped_rows:], k_cache, v_cache, - logical_indices, + logical_indices[grouped_rows:], block_table, - token_to_req, - query_positions, + token_to_req[grouped_rows:], + query_positions[grouped_rows:], sequence_lengths, - out, + out[grouped_rows:], kv_cache_dtype, k_scale, v_scale, flash_attn_v100_cuda, ) + return out if kv_cache_dtype == "fp8_e4m3" and q.shape[0] > 16: # The generic E4M3 XQA kernel accepts at most 16 query rows. Scheduler # iterations can mix a large prefill (or catch-up chunk) with decode - # rows, so split the bulk across the grouped route and keep only the - # remainder on supported XQA batches. This avoids both an invalid - # B>16 XQA launch and the larger Triton split-K fallback workspace. - grouped_rows = 0 - if grouped_enabled: - grouped_rows = ( - q.shape[0] // _SM70_QSA_GROUPED_PAGE4_QUERIES - ) * _SM70_QSA_GROUPED_PAGE4_QUERIES - if grouped_rows: - _qsa_sparse_paged_attention_sm70_grouped_page4( - q[:grouped_rows], - k_cache, - v_cache, - logical_indices[:grouped_rows], - block_table, - token_to_req[:grouped_rows], - query_positions[:grouped_rows], - sequence_lengths, - out[:grouped_rows], - kv_cache_dtype, - k_scale, - v_scale, - flash_attn_v100_cuda, - ) + # rows. If an older Flash-V100 build lacks the quantized grouped ABI, + # retain correctness by slicing the work into supported XQA batches. + # This avoids both an invalid B>16 launch and the larger Triton split-K + # fallback workspace. logger.info_once( - "Splitting a non-grouped E4M3 page4 batch across supported " - "grouped/XQA routes (rows=%d, grouped_rows=%d).", + "Splitting an E4M3 page4 batch across supported XQA launches " + "because this Flash-V100 build lacks the quantized grouped ABI " + "(rows=%d).", q.shape[0], - grouped_rows, ) - for row_start in range(grouped_rows, q.shape[0], 16): + for row_start in range(0, q.shape[0], 16): row_end = min(row_start + 16, q.shape[0]) _qsa_sparse_paged_attention_sm70_xqa_page4_batch( q[row_start:row_end], diff --git a/vllm/models/qwen4_exp/nvidia/qsa.py b/vllm/models/qwen4_exp/nvidia/qsa.py index f35d4bfb0e..c0d4147e1f 100644 --- a/vllm/models/qwen4_exp/nvidia/qsa.py +++ b/vllm/models/qwen4_exp/nvidia/qsa.py @@ -208,6 +208,10 @@ class Qwen4ExpQSAAttention(Qwen3NextAttention, AttentionLayerBase): """Merged Qwen full-attention owner with a QSA index side branch.""" supports_dcp = False + # The paged indexer and sparse attention switch launch profiles after 32 + # query rows. Advertise the first row count in the wider profile so the + # generic MRV2 warmup can compile it before serving traffic. + kernel_warmup_prefill_token_counts = (33,) def __init__( self, diff --git a/vllm/v1/executor/uniproc_executor.py b/vllm/v1/executor/uniproc_executor.py index 430df22594..4b3b3ad46a 100644 --- a/vllm/v1/executor/uniproc_executor.py +++ b/vllm/v1/executor/uniproc_executor.py @@ -63,10 +63,15 @@ def _init_executor(self) -> None: self.driver_worker.init_worker(all_kwargs=[kwargs]) self.driver_worker.init_device() + if envs.VLLM_PLE_CPU_OFFLOAD: + self.driver_worker.spawn_ple_offload() + if envs.VLLM_ELASTIC_EP_SCALE_UP_LAUNCH: self.driver_worker.elastic_ep_execute("load_model") else: self.driver_worker.load_model() + if envs.VLLM_PLE_CPU_OFFLOAD: + self.driver_worker.wait_ple_offload_ready() current_platform.update_block_size_for_backend(self.vllm_config) def _distributed_args(self) -> tuple[str, int, int]: diff --git a/vllm/v1/ple_offload/connector.py b/vllm/v1/ple_offload/connector.py index 64b4474538..f8da6e2898 100644 --- a/vllm/v1/ple_offload/connector.py +++ b/vllm/v1/ple_offload/connector.py @@ -5,6 +5,7 @@ import os import queue import threading +from dataclasses import dataclass from multiprocessing.reduction import ForkingPickler from typing import Any @@ -29,6 +30,14 @@ logger = init_logger(__name__) +@dataclass(frozen=True) +class _PendingPleOffloadRequest: + """Bind request metadata to its MRV2 D2H completion event.""" + + request: PleOffloadRequest + d2h_done_event: torch.cuda.Event | None + + def _cuda_check(result: Any, operation: str) -> Any: """Check the ``(CUresult, ...)`` tuple returned by cuda-python calls.""" error = result[0] if isinstance(result, tuple) else result @@ -92,18 +101,20 @@ def __init__( self._validate_input_sources() self._pinned_input_buffers: list[torch.Tensor] = [] - # PLE rejects DBO, and each forward consumes its output before the - # next launch, so one pending request is sufficient. - self._request_queue: queue.Queue[PleOffloadRequest | None] = queue.Queue( - maxsize=1 + # MRV2 may queue more than one batch under asynchronous scheduling. + # MRV1 remains strictly serialized and retains its existing depth-one + # CPU staging contract. + request_queue_size = ( + vllm_config.max_concurrent_batches if self._uses_cuda_inputs else 1 + ) + self._request_queue: queue.Queue[_PendingPleOffloadRequest | None] = ( + queue.Queue(maxsize=request_queue_size) ) self._request_thread: threading.Thread | None = None self._request_thread_ready = threading.Event() self._zmq_ctx: zmq.Context | None = None self._registration_socket: zmq.Socket | None = None - self._d2h_stream: torch.cuda.Stream | None = None - self._input_ready_event: torch.cuda.Event | None = None - self._d2h_done_event: torch.cuda.Event | None = None + self._d2h_event_pool: queue.Queue[torch.cuda.Event] | None = None try: self._zmq_ctx = zmq.Context() @@ -117,9 +128,11 @@ def __init__( with torch.accelerator.device_index(self.device.index): self._pin_input_buffers() if self._uses_cuda_inputs: - self._d2h_stream = torch.cuda.Stream(device=self.device) - self._input_ready_event = torch.cuda.Event() - self._d2h_done_event = torch.cuda.Event() + self._d2h_event_pool = queue.Queue( + maxsize=vllm_config.max_concurrent_batches + ) + for _ in range(vllm_config.max_concurrent_batches): + self._d2h_event_pool.put_nowait(torch.cuda.Event()) self._start_request_thread(ipc_addr) except Exception: self.close() @@ -251,7 +264,7 @@ def _start_request_thread(self, ipc_addr: str) -> None: raise RuntimeError("Timed out starting the PLE request thread") def _request_loop(self, ipc_addr: str) -> None: - """Stage fixed runner inputs, then notify the CPU worker.""" + """Wait for staged inputs, then notify the CPU worker.""" socket: zmq.Socket | None = None try: if self._zmq_ctx is None: @@ -272,15 +285,33 @@ def _request_loop(self, ipc_addr: str) -> None: if socket is not None: socket.close(linger=0) - def _process_request(self, request: PleOffloadRequest, socket: zmq.Socket) -> None: - """Stage one batch from fixed sources and publish its request.""" + def _process_request( + self, + pending: _PendingPleOffloadRequest, + socket: zmq.Socket, + ) -> None: + """Wait for one staged batch and publish its request.""" + request = pending.request + event = pending.d2h_done_event if self._uses_cuda_inputs: - self._copy_cuda_inputs(request) + event_pool = self._d2h_event_pool + if event_pool is None or event is None: + raise RuntimeError("MRV2 PLE request is missing its D2H event") + with ( + torch.accelerator.device_index(self.device.index), + torch.cuda.nvtx.range("ple_offload.wait_d2h"), + ): + event.synchronize() else: + if event is not None: + raise RuntimeError("MRV1 PLE request unexpectedly has a D2H event") self._copy_cpu_inputs(request) with torch.cuda.nvtx.range("ple_offload.send_request"): socket.send(msgspec.msgpack.encode(request)) + if event is not None: + assert self._d2h_event_pool is not None + self._d2h_event_pool.put_nowait(event) def _copy_cpu_inputs(self, request: PleOffloadRequest) -> None: """Stage MRV1's existing CPU mirrors in the notifier thread.""" @@ -332,38 +363,38 @@ def _validate_input_sources(self) -> None: ): raise ValueError(f"PLE {name} source is incompatible") - def _copy_cuda_inputs(self, request: PleOffloadRequest) -> None: - """Stage MRV2 inputs on the background D2H stream.""" - if ( - self._d2h_stream is None - or self._input_ready_event is None - or self._d2h_done_event is None - ): - raise RuntimeError("PLE D2H resources are not initialized") + def _enqueue_cuda_inputs( + self, + request: PleOffloadRequest, + d2h_done_event: torch.cuda.Event, + ) -> None: + """Stage MRV2 inputs on the model stream and record completion. + The request owns ``d2h_done_event`` until its shared inputs are visible + to the CPU worker. Reusing one background-stream event lets a later + batch retarget the event before the previous request has consumed its + inputs, which can deadlock CUDA Graph replay when the batch size changes. + """ with torch.accelerator.device_index(self.device.index): - with torch.cuda.stream(self._d2h_stream): - self._d2h_stream.wait_event(self._input_ready_event) - with torch.cuda.nvtx.range("ple_offload.copy_input_ids"): - self._input_ids_buf[: request.num_tokens].copy_( - self._input_ids_source[: request.num_tokens], - non_blocking=True, - ) - with torch.cuda.nvtx.range("ple_offload.copy_query_start_loc"): - self._query_start_loc_buf[: request.num_reqs + 1].copy_( - self._query_start_loc_source[: request.num_reqs + 1], + stream = torch.cuda.current_stream(self.device) + with torch.cuda.nvtx.range("ple_offload.copy_input_ids"): + self._input_ids_buf[: request.num_tokens].copy_( + self._input_ids_source[: request.num_tokens], + non_blocking=True, + ) + with torch.cuda.nvtx.range("ple_offload.copy_query_start_loc"): + self._query_start_loc_buf[: request.num_reqs + 1].copy_( + self._query_start_loc_source[: request.num_reqs + 1], + non_blocking=True, + ) + if self._ngram_context_buf is not None: + assert self._ngram_context_source is not None + with torch.cuda.nvtx.range("ple_offload.copy_ngram_context"): + self._ngram_context_buf[: request.num_reqs].copy_( + self._ngram_context_source[: request.num_reqs], non_blocking=True, ) - if self._ngram_context_buf is not None: - assert self._ngram_context_source is not None - with torch.cuda.nvtx.range("ple_offload.copy_ngram_context"): - self._ngram_context_buf[: request.num_reqs].copy_( - self._ngram_context_source[: request.num_reqs], - non_blocking=True, - ) - self._d2h_done_event.record(self._d2h_stream) - with torch.cuda.nvtx.range("ple_offload.wait_d2h"): - self._d2h_done_event.synchronize() + d2h_done_event.record(stream) def _launch( self, @@ -376,17 +407,25 @@ def _launch( if self.tp_rank != 0: return - if self._uses_cuda_inputs: - assert self._input_ready_event is not None - # The background copy stream waits for runner input production - # without making the model stream wait for D2H completion. - self._input_ready_event.record(torch.cuda.current_stream(self.device)) request = PleOffloadRequest( dp_rank=self.dp_rank, num_tokens=num_tokens, num_reqs=num_reqs, ) - self._request_queue.put_nowait(request) + d2h_done_event = None + if self._uses_cuda_inputs: + if self._d2h_event_pool is None: + raise RuntimeError("PLE D2H event pool is not initialized") + try: + d2h_done_event = self._d2h_event_pool.get_nowait() + except queue.Empty as exc: + raise RuntimeError( + "PLE has more MRV2 requests than configured concurrent batches" + ) from exc + self._enqueue_cuda_inputs(request, d2h_done_event) + self._request_queue.put_nowait( + _PendingPleOffloadRequest(request, d2h_done_event) + ) def prepare_forward( self, @@ -435,9 +474,7 @@ def close(self) -> None: if self._pinned_input_buffers: with torch.accelerator.device_index(self.device.index): self._unpin_input_buffers() - self._d2h_done_event = None - self._input_ready_event = None - self._d2h_stream = None + self._d2h_event_pool = None if self._registration_socket is not None: self._registration_socket.close(linger=0) self._registration_socket = None diff --git a/vllm/v1/ple_offload/worker.py b/vllm/v1/ple_offload/worker.py index 06b66c7eb8..8b203593a6 100644 --- a/vllm/v1/ple_offload/worker.py +++ b/vllm/v1/ple_offload/worker.py @@ -20,10 +20,13 @@ """ import contextlib +import ctypes import multiprocessing.process +import os import signal import tempfile import threading +import time from collections.abc import Iterable from dataclasses import dataclass from multiprocessing.connection import Connection @@ -31,6 +34,7 @@ from typing import Any, cast import msgspec +import psutil import torch import torch.distributed as dist import zmq @@ -55,6 +59,7 @@ process_weights_after_loading, ) from vllm.model_executor.model_loader.weight_utils import initialize_dummy_weights +from vllm.utils.mem_constants import GiB_bytes from vllm.utils.system_utils import decorate_logs, get_mp_context from vllm.utils.torch_utils import set_default_torch_dtype from vllm.v1.ple_offload.protocol import ( @@ -66,6 +71,218 @@ logger = init_logger(__name__) +def _configure_ple_numa_locality(vllm_config: VllmConfig) -> int | None: + """Use GPU-local CPUs without imposing a single-node memory hard limit. + + PLE tables can occupy tens of GiB. Inheriting ``numactl --membind`` from + the spawning GPU worker can therefore swap a large part of the table even + when another local NUMA node has ample RAM. Use local-first allocation, + which retains NUMA locality but permits normal fallback under pressure, + and bind future CPU lookup threads to the first visible GPU's node. + """ + if not envs.VLLM_PLE_OFFLOAD_AUTO_NUMA or not hasattr(os, "sched_setaffinity"): + return None + + try: + from vllm.platforms import current_platform + from vllm.utils.cpu_resource_utils import get_allowed_cpu_list + + parallel_config = vllm_config.parallel_config + numa_nodes = parallel_config.numa_bind_nodes + if numa_nodes is None: + numa_nodes = current_platform.get_all_device_numa_nodes() + if not numa_nodes: + logger.warning( + "PLE automatic NUMA placement skipped: GPU topology is unavailable." + ) + return None + + target_node = int(numa_nodes[0]) + allowed_cpus = set(os.sched_getaffinity(0)) + local_cpus = { + cpu.id + for cpu in get_allowed_cpu_list() + if cpu.numa_node == target_node and cpu.id in allowed_cpus + } + if not local_cpus: + logger.warning( + "PLE automatic NUMA placement skipped: no allowed CPU belongs " + "to GPU-local NUMA node %d.", + target_node, + ) + return None + + except Exception as error: + # Placement is an optimization. Keep the existing functional path on + # platforms or containers that cannot expose/change NUMA affinity. + logger.warning("PLE automatic NUMA placement failed: %s", error) + return None + + # A PLE child spawned from a NUMA-bound GPU worker inherits its strict + # MPOL_BIND policy. Relax only this child to local-first allocation so a + # table larger than one node can spill into other host RAM, not swap. Keep + # this best-effort step independent from CPU affinity: containers often + # permit sched_setaffinity while denying set_mempolicy. + memory_policy = "inherited" + try: + from vllm.utils.numa_utils import get_libnuma + + libnuma = get_libnuma() + if libnuma is not None and libnuma.numa_available() >= 0: + libnuma.numa_set_localalloc() + # numa_set_localalloc() has a void C signature and some libnuma + # builds only print when set_mempolicy is denied. Read the policy + # back before claiming that the strict inherited bind was relaxed. + policy_mode = ctypes.c_int(-1) + status = libnuma.get_mempolicy( + ctypes.byref(policy_mode), + None, + ctypes.c_ulong(0), + None, + ctypes.c_ulong(0), + ) + if status != 0 or policy_mode.value != 4: # MPOL_LOCAL + raise OSError( + "numa_set_localalloc did not install MPOL_LOCAL " + f"(status={status}, mode={policy_mode.value})" + ) + memory_policy = "local-first with host-RAM fallback" + except Exception as error: + logger.warning( + "PLE could not relax its inherited NUMA memory policy: %s", error + ) + + try: + os.sched_setaffinity(0, local_cpus) + except Exception as error: + logger.warning("PLE automatic CPU affinity failed: %s", error) + return None + + logger.info( + "PLE CPU lookup affinity: GPU-local NUMA node %d, CPUs %s; " + "memory policy is %s.", + target_node, + ",".join(str(cpu) for cpu in sorted(local_cpus)), + memory_policy, + ) + return target_node + + +def _iter_unique_module_storage_views( + modules: Iterable[torch.nn.Module], +) -> Iterable[torch.Tensor]: + """Yield one byte view spanning each unique materialized CPU storage.""" + seen: set[tuple[str, int, int]] = set() + for module in modules: + tensors = (*module.parameters(), *module.buffers()) + for tensor in tensors: + if tensor.device.type == "meta": + continue + storage = tensor.untyped_storage() + storage_bytes = storage.nbytes() + key = (str(tensor.device), storage.data_ptr(), storage_bytes) + if key in seen: + continue + seen.add(key) + yield torch.empty(0, dtype=torch.uint8, device=tensor.device).set_( + storage, + 0, + (storage_bytes,), + (1,), + ) + + +def _estimate_module_storage_bytes(modules: Iterable[torch.nn.Module]) -> int: + """Return unique materialized parameter and buffer storage bytes.""" + return sum(view.numel() for view in _iter_unique_module_storage_views(modules)) + + +def _prefault_module_storage(modules: Iterable[torch.nn.Module]) -> int: + """Bring offloaded table pages back to RAM when host capacity permits.""" + if not envs.VLLM_PLE_OFFLOAD_PREFAULT: + return 0 + + storage_views = list(_iter_unique_module_storage_views(modules)) + required_bytes = sum(view.numel() for view in storage_views) + if required_bytes == 0: + return 0 + + process = psutil.Process(os.getpid()) + process_info = process.memory_full_info() + resident_bytes = process_info.rss + swap_bytes_before = int(getattr(process_info, "swap", 0)) + memory = psutil.virtual_memory() + headroom_bytes = max(4 * GiB_bytes, memory.total // 10) + missing_bytes = min( + required_bytes, + max(required_bytes - resident_bytes, swap_bytes_before, 0), + ) + if missing_bytes + headroom_bytes > memory.available: + logger.warning( + "Skipping PLE RAM prefault: %.2f GiB may be non-resident, but " + "only %.2f GiB host RAM is available with %.2f GiB reserved " + "as runtime headroom.", + missing_bytes / GiB_bytes, + memory.available / GiB_bytes, + headroom_bytes / GiB_bytes, + ) + return 0 + + page_bytes = int(os.sysconf("SC_PAGE_SIZE")) + chunk_bytes = 256 * 1024 * 1024 + checksum = 0 + started = time.perf_counter() + for view in storage_views: + for start in range(0, view.numel(), chunk_bytes): + page_heads = view[start : start + chunk_bytes : page_bytes] + checksum ^= int(page_heads.sum(dtype=torch.int64).item()) + elapsed = time.perf_counter() - started + full_info = process.memory_full_info() + swap_bytes = int(getattr(full_info, "swap", 0)) + logger.info( + "PLE RAM prefault touched %.2f GiB in %.2f s " + "(process RSS %.2f GiB, swap %.2f -> %.2f GiB).", + required_bytes / GiB_bytes, + elapsed, + full_info.rss / GiB_bytes, + swap_bytes_before / GiB_bytes, + swap_bytes / GiB_bytes, + ) + del checksum + return required_bytes + + +def _log_ple_host_memory_capacity(layers: Iterable[PleOffloadLayer]) -> int: + """Report whether PLE weights can remain resident without swap pressure.""" + required_bytes = _estimate_module_storage_bytes(layers) + memory = psutil.virtual_memory() + swap = psutil.swap_memory() + # Leave enough memory for the API process, model workers, page tables, and + # transient request buffers. This is a warning threshold, not an artificial + # model or checkpoint admission gate. + headroom_bytes = max(4 * GiB_bytes, memory.total // 10) + logger.info( + "PLE host-memory requirement: %.2f GiB weights/buffers, %.2f GiB " + "currently available of %.2f GiB total (%.2f GiB swap in use).", + required_bytes / GiB_bytes, + memory.available / GiB_bytes, + memory.total / GiB_bytes, + swap.used / GiB_bytes, + ) + if required_bytes + headroom_bytes > memory.available: + logger.warning( + "PLE CPU offload may page its %.2f GiB table: only %.2f GiB host " + "RAM is currently available and %.2f GiB is reserved as runtime " + "headroom. Paging causes severe, input-dependent prefill and " + "decode latency. Free host memory, add RAM, or disable PLE CPU " + "offload when the table fits on the accelerators.", + required_bytes / GiB_bytes, + memory.available / GiB_bytes, + headroom_bytes / GiB_bytes, + ) + return required_bytes + + @dataclass class PleOffloadOutputTarget: """GPU output destination and semaphore for one TP worker.""" @@ -245,6 +462,7 @@ def proc_main( ) -> None: """Load PLE weights, accept registrations, and run the request loop.""" decorate_logs("PleOffloadWorker") + _configure_ple_numa_locality(vllm_config) ready_reader, ready_writer = ready_pipe ready_reader.close() shutdown_event = threading.Event() @@ -295,6 +513,7 @@ def handle_signal(_signum: int, _frame: object) -> None: # READY means that the process can immediately serve requests. Wait # for every DP/TP worker to register before notifying the parent. runner.accept_registrations(pull_socket, num_workers) + _prefault_module_storage(runner._layers.values()) ready_writer.send( { "status": PleOffloadWorker.READY_STR, @@ -387,6 +606,7 @@ def _load_weights(self) -> None: len(offload_layers), sorted(offload_layers), ) + _log_ple_host_memory_capacity(offload_layers.values()) offload_prefixes = tuple(f"{name}." for name in offload_layers) # Step 3: filter checkpoint tensors before model.load_weights(). The diff --git a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py index ca6db6cad3..138efba94e 100644 --- a/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/eagle/speculator.py @@ -41,6 +41,27 @@ logger = init_logger(__name__) +def _mtp_decode_warmup_request_sizes( + capture_sizes: list[int] | None, + decode_query_len: int, + max_num_reqs: int, +) -> tuple[int, ...]: + """Return production request counts represented by verifier graph shapes.""" + if decode_query_len <= 0 or max_num_reqs <= 0: + return () + + # A single-request decode is always a valid eager fallback even when the + # user disables graphs or supplies only prefill-oriented capture sizes. + request_sizes = {1} + for num_tokens in capture_sizes or (): + if num_tokens <= 0 or num_tokens % decode_query_len: + continue + num_reqs = num_tokens // decode_query_len + if 1 <= num_reqs <= max_num_reqs: + request_sizes.add(num_reqs) + return tuple(sorted(request_sizes)) + + class EagleSpeculator: def __init__(self, vllm_config: VllmConfig, device: torch.device): self.vllm_config = vllm_config @@ -237,11 +258,28 @@ def warmup_sm70_mtp_moe_kernels( dummy_run(16) warmed.append("mtp_draft_moe_prefill_m16") - # One real decode-shaped round executes the M5 verifier window - # followed by all three M1 continuation passes. - dummy_run(1, uniform_decode=True) + decode_query_len = 1 + self.num_speculative_steps + request_sizes = _mtp_decode_warmup_request_sizes( + self.vllm_config.compilation_config.cudagraph_capture_sizes, + decode_query_len, + self.max_num_reqs, + ) + executed_request_sizes = [] + for num_reqs in request_sizes: + num_tokens = decode_query_len * num_reqs + if num_tokens > self.max_num_tokens: + continue + # Exercise the same verifier/draft row counts that production + # graphs advertise. This avoids a first-request JIT without + # tying the tuned MoE path to a fixed concurrency. + dummy_run(num_tokens, uniform_decode=True) + executed_request_sizes.append(num_reqs) torch.accelerator.synchronize() - warmed.append("mtp_draft_moe_decode_m5_m1") + if executed_request_sizes: + warmed.append( + "mtp_draft_moe_decode_reqs_" + + "_".join(str(size) for size in executed_request_sizes) + ) except Exception as err: # pragma: no cover - best-effort warmup logger.warning_once("SM70 V2 MTP MoE warmup skipped: %s", err) return () diff --git a/vllm/v1/worker/gpu/warmup.py b/vllm/v1/worker/gpu/warmup.py index 9bcafe5210..9ad49e90a2 100644 --- a/vllm/v1/worker/gpu/warmup.py +++ b/vllm/v1/worker/gpu/warmup.py @@ -26,6 +26,28 @@ from vllm.v1.worker.gpu.model_runner import GPUModelRunner +def _kernel_prefill_warmup_token_counts( + model_runner: GPUModelRunner, + default_prompt_len: int, +) -> tuple[int, ...]: + """Collect per-request prefill sizes advertised by active kernel owners.""" + max_tokens = min( + model_runner.scheduler_config.max_num_batched_tokens, + model_runner.max_model_len, + ) + token_counts = {default_prompt_len} + static_context = model_runner.compilation_config.static_forward_context + for layer in static_context.values(): + for token_count in getattr(layer, "kernel_warmup_prefill_token_counts", ()): + if ( + isinstance(token_count, int) + and not isinstance(token_count, bool) + and default_prompt_len < token_count <= max_tokens + ): + token_counts.add(token_count) + return tuple(sorted(token_counts)) + + def _reserved_block_count( num_tokens: int, kv_cache_spec: KVCacheSpec, @@ -85,33 +107,16 @@ def warmup_kernels( # Use 1 + num_spec_steps + 1 tokens so the prefill batch's per-request # query length exceeds decode_query_len (= 1 + num_spec_steps), preventing # it from being misclassified as a uniform decode batch. - prompt_len = 2 + num_spec_steps - prompt_token_ids = list(range(prompt_len)) - # After prefill, decode generates 1 verified + num_spec_steps draft tokens. - decode_len = prompt_len + 1 + num_spec_steps + default_prompt_len = 2 + num_spec_steps + prompt_lengths = _kernel_prefill_warmup_token_counts( + model_runner, default_prompt_len + ) kv_cache_groups = model_runner.kv_cache_config.kv_cache_groups num_kv_cache_groups = len(kv_cache_groups) - # Compute per-request block counts for each KV cache group. block_count = _warmup_block_counter(model_runner) kv_cache_specs = [group.kv_cache_spec for group in kv_cache_groups] - prefill_block_counts = [block_count(prompt_len, spec) for spec in kv_cache_specs] - decode_block_counts = [block_count(decode_len, spec) for spec in kv_cache_specs] - decode_block_deltas = [ - d - p for d, p in zip(decode_block_counts, prefill_block_counts) - ] - max_blocks_per_req = sum(decode_block_counts) - - num_reqs = min( - model_runner.scheduler_config.max_num_seqs, - model_runner.scheduler_config.max_num_batched_tokens - // max(prompt_len, 1 + num_spec_steps), - # Reserve block 0 (null block) and ensure we have enough blocks. - max(1, (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req), - ) - - req_ids = [f"_warmup_{i}_" for i in range(num_reqs)] # SamplingParams exercising all sampling features. if model_runner.is_pooling_model: @@ -121,82 +126,121 @@ def warmup_kernels( sampling_params = SamplingParams.for_sampler_warmup() pooling_params = None - # Assign distinct block IDs per request per group. 0 null block, start from 1. - next_block_id = 1 - - def _alloc_blocks(num_blocks: int) -> list[int]: - nonlocal next_block_id - return list(range(next_block_id, next_block_id := next_block_id + num_blocks)) - - # Step 1: Prefill all requests with 2 + num_spec_steps prompt tokens each. - new_reqs = [ - NewRequestData.from_request( - Request(req_ids[i], prompt_token_ids, sampling_params, pooling_params), - block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), - prefill_token_ids=prompt_token_ids, - ) - for i in range(num_reqs) - ] - - prefill_output = SchedulerOutput.make_empty() - prefill_output.scheduled_new_reqs = new_reqs - prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} - prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs - prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups - - # Disable KV connector for warmup run. + # Disable KV connector for all warmup runs. Kernel-advertised profiles run + # with one request so adding an operator does not silently multiply startup + # work by max_num_seqs. model_runner.kv_connector.set_disabled(True) - worker_execute_model(prefill_output) - - if not model_runner.is_pooling_model: - # Warm up sampler and perform a decode step for non-pooling models. - - grammar_output = None - if model_runner.is_last_pp_rank: - # Build a GrammarOutput to exercise the structured output bitmask - # kernel during the prefill step. - vocab_size = model_runner.model_config.get_vocab_size() - bitmask_width = (vocab_size + 31) // 32 - grammar_bitmask = np.full( - (len(req_ids), bitmask_width), fill_value=-1, dtype=np.int32 - ) - grammar_output = GrammarOutput( - structured_output_request_ids=req_ids, grammar_bitmask=grammar_bitmask + try: + for profile_idx, prompt_len in enumerate(prompt_lengths): + prompt_token_ids = list(range(prompt_len)) + decode_len = prompt_len + 1 + num_spec_steps + prefill_block_counts = [ + block_count(prompt_len, spec) for spec in kv_cache_specs + ] + decode_block_counts = [ + block_count(decode_len, spec) for spec in kv_cache_specs + ] + decode_block_deltas = [ + d - p for d, p in zip(decode_block_counts, prefill_block_counts) + ] + max_blocks_per_req = sum(decode_block_counts) + num_reqs = min( + model_runner.scheduler_config.max_num_seqs, + model_runner.scheduler_config.max_num_batched_tokens + // max(prompt_len, 1 + num_spec_steps), + # Reserve block 0 (null block) and ensure enough blocks. + max( + 1, + (model_runner.kv_cache_config.num_blocks - 1) // max_blocks_per_req, + ), ) - - worker_sample_tokens(grammar_output) - - # Step 2: Decode all requests with 1 + num_spec_steps tokens each. - cached_req_data = CachedRequestData.make_empty() - cached_req_data.req_ids = list(req_ids) - cached_req_data.num_computed_tokens = [prompt_len] * num_reqs - cached_req_data.num_output_tokens = [1] * num_reqs - new_block = any(decode_block_deltas) - cached_req_data.new_block_ids = [ - tuple(_alloc_blocks(n) for n in decode_block_deltas) if new_block else None - for _ in range(num_reqs) - ] - - decode_output = SchedulerOutput.make_empty() - decode_output.scheduled_cached_reqs = cached_req_data - decode_output.num_scheduled_tokens = { - req_id: 1 + num_spec_steps for req_id in req_ids - } - if num_spec_steps > 0: - decode_output.scheduled_spec_decode_tokens = { - req_id: [0] * num_spec_steps for req_id in req_ids - } - decode_output.total_num_scheduled_tokens = sum( - decode_output.num_scheduled_tokens.values() - ) - decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups - - worker_execute_model(decode_output) - worker_sample_tokens(None) - - # Clean up - process finish_req_ids. - cleanup_output = SchedulerOutput.make_empty() - cleanup_output.finished_req_ids = set(req_ids) - worker_execute_model(cleanup_output) - model_runner.kv_connector.set_disabled(False) + if profile_idx: + num_reqs = min(num_reqs, 1) + if num_reqs <= 0: + continue + + req_ids = [f"_warmup_{profile_idx}_{i}_" for i in range(num_reqs)] + next_block_id = 1 + + def _alloc_blocks(num_blocks: int) -> list[int]: + nonlocal next_block_id + return list( + range(next_block_id, next_block_id := next_block_id + num_blocks) + ) + + new_reqs = [ + NewRequestData.from_request( + Request( + req_ids[i], + prompt_token_ids, + sampling_params, + pooling_params, + ), + block_ids=tuple(_alloc_blocks(n) for n in prefill_block_counts), + prefill_token_ids=prompt_token_ids, + ) + for i in range(num_reqs) + ] + + prefill_output = SchedulerOutput.make_empty() + prefill_output.scheduled_new_reqs = new_reqs + prefill_output.num_scheduled_tokens = {rid: prompt_len for rid in req_ids} + prefill_output.total_num_scheduled_tokens = prompt_len * num_reqs + prefill_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + worker_execute_model(prefill_output) + + if not model_runner.is_pooling_model: + grammar_output = None + if profile_idx == 0 and model_runner.is_last_pp_rank: + # Exercise the structured-output bitmask once; extra + # operator profiles only need the model path. + vocab_size = model_runner.model_config.get_vocab_size() + bitmask_width = (vocab_size + 31) // 32 + grammar_bitmask = np.full( + (len(req_ids), bitmask_width), + fill_value=-1, + dtype=np.int32, + ) + grammar_output = GrammarOutput( + structured_output_request_ids=req_ids, + grammar_bitmask=grammar_bitmask, + ) + worker_sample_tokens(grammar_output) + + cached_req_data = CachedRequestData.make_empty() + cached_req_data.req_ids = list(req_ids) + cached_req_data.num_computed_tokens = [prompt_len] * num_reqs + cached_req_data.num_output_tokens = [1] * num_reqs + new_block = any(decode_block_deltas) + cached_req_data.new_block_ids = [ + ( + tuple(_alloc_blocks(n) for n in decode_block_deltas) + if new_block + else None + ) + for _ in range(num_reqs) + ] + + decode_output = SchedulerOutput.make_empty() + decode_output.scheduled_cached_reqs = cached_req_data + decode_output.num_scheduled_tokens = { + req_id: 1 + num_spec_steps for req_id in req_ids + } + if num_spec_steps > 0: + decode_output.scheduled_spec_decode_tokens = { + req_id: [0] * num_spec_steps for req_id in req_ids + } + decode_output.total_num_scheduled_tokens = sum( + decode_output.num_scheduled_tokens.values() + ) + decode_output.num_common_prefix_blocks = [0] * num_kv_cache_groups + + worker_execute_model(decode_output) + worker_sample_tokens(None) + + cleanup_output = SchedulerOutput.make_empty() + cleanup_output.finished_req_ids = set(req_ids) + worker_execute_model(cleanup_output) + finally: + model_runner.kv_connector.set_disabled(False) torch.accelerator.synchronize()