[Bug Fix] Allow pinned memory for WSL2 (#41496)

Signed-off-by: Jimmy Lee <hirejimmylee@gmail.com>

benchmarks/benchmark_pin_memory.py

@@ -0,0 +1,358 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark and regression-test pinned (page-locked) CPU memory for vLLM.
+
+Verifies that enabling pinned memory does not regress throughput or latency
+compared to unpinned memory.  Each condition runs in an isolated ``spawn``
+subprocess so both start from a cold CUDA context, giving an unbiased
+comparison.
+
+Usage
+-----
+Run all tests with the default model::
+
+    python benchmarks/benchmark_pin_memory.py -v
+
+Override the model and optional max-model-len::
+
+    python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B -v
+    python benchmarks/benchmark_pin_memory.py --model unsloth/Qwen3-1.7B \
+        --max-model-len 8192 -v
+
+Run only throughput or latency tests::
+
+    python benchmarks/benchmark_pin_memory.py -v -k test_throughput
+    python benchmarks/benchmark_pin_memory.py -v -k test_latency
+
+Run only the v1 or v2 runner variant::
+
+    python benchmarks/benchmark_pin_memory.py -v -k v1
+    python benchmarks/benchmark_pin_memory.py -v -k v2
+
+Note: on WSL2, v1 runner tests are skipped because pin memory is not available
+for the v1 runner without cpu_offload_gb.  Run on other platforms to exercise v1.
+"""
+
+import argparse
+import json
+import multiprocessing
+import sys
+import tempfile
+
+import pytest
+
+# Allow up to 2% degradation.  Both benchmark runs start from an identical
+# cold CUDA context (separate spawn subprocesses), so the measured difference
+# reflects the genuine pin_memory overhead rather than cold/warm ordering bias.
+_THROUGHPUT_TOLERANCE = 0.98
+_THROUGHPUT_NUM_REQUESTS = 200
+_THROUGHPUT_INPUT_LEN = 128
+_THROUGHPUT_OUTPUT_LEN = 512
+_THROUGHPUT_MAX_NUM_SEQS = 128
+
+# Latency benchmark constants — match latency.py defaults.
+_LATENCY_TOLERANCE = 1.02  # Allow up to 2% latency regression.
+_LATENCY_BATCH_SIZE = 64
+_LATENCY_INPUT_LEN = 32
+_LATENCY_OUTPUT_LEN = 128
+_LATENCY_WARMUP_ITERS = 5
+_LATENCY_BENCH_ITERS = 15
+
+_DEFAULT_MODEL = "unsloth/Qwen3-1.7B"
+_DEFAULT_MAX_MODEL_LEN = 16384
+
+
+def _benchmark_args() -> argparse.Namespace:
+    parser = argparse.ArgumentParser(add_help=False)
+    parser.add_argument("--model", default=_DEFAULT_MODEL)
+    parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
+    args, _ = parser.parse_known_args()
+    return args
+
+
+@pytest.fixture
+def model() -> str:
+    return _benchmark_args().model
+
+
+@pytest.fixture
+def max_model_len() -> int:
+    return _benchmark_args().max_model_len
+
+
+def _skip_if_pin_memory_not_available(engine_args_kwargs: dict) -> None:
+    """Skip the current pytest test if pin_memory is unavailable for this config."""
+    import vllm.utils.platform_utils as pu
+    from vllm.config import set_current_vllm_config
+    from vllm.engine.arg_utils import EngineArgs
+
+    vllm_config = EngineArgs(**engine_args_kwargs).create_engine_config()
+    with set_current_vllm_config(vllm_config):
+        pu.is_pin_memory_available.cache_clear()
+        if not pu.is_pin_memory_available():
+            import os
+
+            runner = "v2" if os.environ.get("VLLM_USE_V2_MODEL_RUNNER") == "1" else "v1"
+            model = engine_args_kwargs.get("model", "unknown")
+            print(
+                f"\033[33mSKIP: pin_memory not available for "
+                f"{runner} runner, model={model}\033[0m"
+            )
+            pytest.skip("pin_memory not available for this configuration")
+
+
+def _throughput_worker(
+    pin: bool,
+    engine_args_kwargs: dict,
+    q: "multiprocessing.Queue[float]",
+    v2_mode: bool = False,
+) -> None:
+    """Run throughput benchmark in a fresh spawn subprocess.
+
+    Delegates to vllm/benchmarks/throughput.py main() using the random dataset,
+    so the methodology matches the official benchmark.  Results are written to a
+    temp JSON file and forwarded through the queue as tokens/s.
+
+    v2_mode: when True, monkeypatches is_uva_available() to always return True
+    so the v2 model runner's UVA buffers remain functional even when pin=False.
+    This isolates the non-UVA pin_memory paths in v2.
+    """
+    import vllm.utils.platform_utils as pu
+    from vllm.platforms import current_platform
+
+    pu.is_pin_memory_available.cache_clear()
+    pu.is_uva_available.cache_clear()
+    type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
+    if v2_mode:
+        pu.is_uva_available = lambda: True
+
+    from vllm.benchmarks.throughput import add_cli_args
+    from vllm.benchmarks.throughput import main as throughput_main
+
+    parser = argparse.ArgumentParser()
+    add_cli_args(parser)
+    args = parser.parse_args([])
+
+    for key, val in engine_args_kwargs.items():
+        setattr(args, key, val)
+    args.max_num_seqs = _THROUGHPUT_MAX_NUM_SEQS
+    args.dataset_name = "random"
+    args.input_len = _THROUGHPUT_INPUT_LEN
+    args.output_len = _THROUGHPUT_OUTPUT_LEN
+    # Nullify defaults that conflict with explicit input/output_len.
+    args.random_input_len = None
+    args.random_output_len = None
+    args.random_prefix_len = None
+    args.num_prompts = _THROUGHPUT_NUM_REQUESTS
+    args.seed = 0
+    args.disable_detokenize = True
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        tmp_path = f.name
+    args.output_json = tmp_path
+
+    throughput_main(args)
+
+    with open(tmp_path) as f:
+        results = json.load(f)
+    q.put(results["tokens_per_second"])
+
+
+def _run_throughput_benchmark(
+    pin: bool,
+    engine_args_kwargs: dict,
+    v2_mode: bool = False,
+) -> float:
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    p = ctx.Process(
+        target=_throughput_worker,
+        args=(pin, engine_args_kwargs, q, v2_mode),
+    )
+    p.start()
+    p.join()
+    if p.exitcode != 0:
+        raise RuntimeError(
+            f"Throughput benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
+        )
+    return q.get()
+
+
+def _latency_worker(
+    pin: bool,
+    engine_args_kwargs: dict,
+    q: "multiprocessing.Queue[dict]",
+    v2_mode: bool = False,
+) -> None:
+    """Run latency benchmark in a fresh spawn subprocess.
+
+    Follows latency.py methodology: fixed batch of dummy token IDs, warmup
+    iterations to reach steady state, then timed iterations reduced to avg
+    and percentiles.  Results are written to a temp JSON file by latency_main
+    and forwarded through the queue.
+    """
+    import vllm.utils.platform_utils as pu
+    from vllm.platforms import current_platform
+
+    pu.is_pin_memory_available.cache_clear()
+    pu.is_uva_available.cache_clear()
+    type(current_platform).is_pin_memory_available = classmethod(lambda cls: pin)
+    if v2_mode:
+        pu.is_uva_available = lambda: True
+
+    from vllm.benchmarks.latency import add_cli_args
+    from vllm.benchmarks.latency import main as latency_main
+
+    parser = argparse.ArgumentParser()
+    add_cli_args(parser)
+    args = parser.parse_args([])
+
+    for key, val in engine_args_kwargs.items():
+        setattr(args, key, val)
+    args.input_len = _LATENCY_INPUT_LEN
+    args.output_len = _LATENCY_OUTPUT_LEN
+    args.batch_size = _LATENCY_BATCH_SIZE
+    args.num_iters_warmup = _LATENCY_WARMUP_ITERS
+    args.num_iters = _LATENCY_BENCH_ITERS
+    args.profile = False
+    args.disable_detokenize = True
+
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".json", delete=False) as f:
+        tmp_path = f.name
+    args.output_json = tmp_path
+
+    latency_main(args)
+
+    with open(tmp_path) as f:
+        results = json.load(f)
+    q.put(results)
+
+
+def _run_latency_benchmark(
+    pin: bool,
+    engine_args_kwargs: dict,
+    v2_mode: bool = False,
+) -> dict:
+    ctx = multiprocessing.get_context("spawn")
+    q = ctx.Queue()
+    p = ctx.Process(
+        target=_latency_worker,
+        args=(pin, engine_args_kwargs, q, v2_mode),
+    )
+    p.start()
+    p.join()
+    if p.exitcode != 0:
+        raise RuntimeError(
+            f"Latency benchmark subprocess (pin={pin}) exited with code {p.exitcode}"
+        )
+    return q.get()
+
+
+@pytest.mark.parametrize(
+    "test_v2_runner",
+    [
+        pytest.param(False, id="v1"),
+        pytest.param(True, id="v2"),
+    ],
+)
+class TestPinnedMemory:
+    """Verify pinned memory yields >= throughput vs unpinned via real vLLM inference."""
+
+    def test_throughput(self, monkeypatch, test_v2_runner, model, max_model_len):
+        """Benchmark throughput with pin_memory forced on then off.
+
+        Delegates to vllm/benchmarks/throughput.py main() with the random
+        dataset.  Each condition runs in an isolated spawn subprocess so both
+        start from a cold CUDA context, giving an unbiased comparison.
+        """
+        monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
+
+        engine_args_kwargs = dict(
+            model=model,
+            gpu_memory_utilization=0.88,
+            max_model_len=max_model_len,
+            enable_prefix_caching=False,
+        )
+
+        _skip_if_pin_memory_not_available(engine_args_kwargs)
+
+        unpinned_tps = _run_throughput_benchmark(
+            False, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+        pinned_tps = _run_throughput_benchmark(
+            True, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+
+        pct_diff = (pinned_tps - unpinned_tps) / unpinned_tps * 100
+        runner = "v2" if test_v2_runner else "v1"
+        print(
+            f"\n=== Throughput results ({runner} runner, {model}) ==="
+            f"\npin_memory=True:  {pinned_tps:.1f} tok/s"
+            f"\npin_memory=False: {unpinned_tps:.1f} tok/s"
+            f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
+        )
+
+        assert pinned_tps >= unpinned_tps * _THROUGHPUT_TOLERANCE, (
+            f"Pinned throughput ({pinned_tps:.1f} tok/s) fell more than "
+            f"{(1.0 - _THROUGHPUT_TOLERANCE) * 100:.1f}% below "
+            f"unpinned ({unpinned_tps:.1f} tok/s)."
+        )
+
+    def test_latency(self, monkeypatch, test_v2_runner, model, max_model_len):
+        """Benchmark per-batch latency with pin_memory forced on then off.
+
+        Follows vllm/benchmarks/latency.py: fixed dummy-token batch, warmup
+        iterations to reach steady state, then timed iterations reduced to avg
+        and percentiles.  Subprocesses run serially so each gets a cold CUDA
+        context without GPU memory pressure from the other run.
+        """
+        monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+        monkeypatch.setenv("VLLM_USE_V2_MODEL_RUNNER", "1" if test_v2_runner else "0")
+
+        engine_args_kwargs = dict(
+            model=model,
+            gpu_memory_utilization=0.88,
+            max_model_len=max_model_len,
+            enable_prefix_caching=False,
+        )
+
+        _skip_if_pin_memory_not_available(engine_args_kwargs)
+
+        unpinned = _run_latency_benchmark(
+            False, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+        pinned = _run_latency_benchmark(
+            True, engine_args_kwargs, v2_mode=test_v2_runner
+        )
+
+        pct_diff = (
+            (pinned["avg_latency"] - unpinned["avg_latency"])
+            / unpinned["avg_latency"]
+            * 100
+        )
+        runner = "v2" if test_v2_runner else "v1"
+        print(
+            f"\n=== Latency results ({runner} runner, {model}) ==="
+            f"\npin_memory=True:  avg={pinned['avg_latency']:.3f}s"
+            f"  p50={pinned['percentiles']['50']:.3f}s"
+            f"  p99={pinned['percentiles']['99']:.3f}s"
+            f"\npin_memory=False: avg={unpinned['avg_latency']:.3f}s"
+            f"  p50={unpinned['percentiles']['50']:.3f}s"
+            f"  p99={unpinned['percentiles']['99']:.3f}s"
+            f"\nDifference: {pct_diff:+.1f}% (pinned vs unpinned)"
+        )
+
+        assert pinned["avg_latency"] <= unpinned["avg_latency"] * _LATENCY_TOLERANCE, (
+            f"Pinned avg latency ({pinned['avg_latency']:.3f}s) exceeded "
+            f"unpinned ({unpinned['avg_latency']:.3f}s) by more than "
+            f"{(_LATENCY_TOLERANCE - 1.0) * 100:.1f}%."
+        )
+
+
+if __name__ == "__main__":
+    _parser = argparse.ArgumentParser(add_help=False)
+    _parser.add_argument("--model", default=_DEFAULT_MODEL)
+    _parser.add_argument("--max-model-len", type=int, default=_DEFAULT_MAX_MODEL_LEN)
+    _, _remaining = _parser.parse_known_args()
+    sys.exit(pytest.main([__file__] + _remaining))

vllm/envs.py

@@ -260,6 +260,7 @@ if TYPE_CHECKING:
     VLLM_DEBUG_MFU_METRICS: bool = False
     VLLM_WEIGHT_OFFLOADING_DISABLE_PIN_MEMORY: bool = False
     VLLM_WEIGHT_OFFLOADING_DISABLE_UVA: bool = False
+    VLLM_WSL2_ENABLE_PIN_MEMORY: bool = False
     VLLM_DISABLE_LOG_LOGO: bool = False
     VLLM_LORA_DISABLE_PDL: bool = False
     VLLM_ENABLE_CUDA_COMPATIBILITY: bool = False
@@ -1839,6 +1840,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_WEIGHT_OFFLOADING_DISABLE_UVA": lambda: bool(
         int(os.getenv("VLLM_WEIGHT_OFFLOADING_DISABLE_UVA", "0"))
     ),
+    # On WSL2 with a compatible kernel (>= 4.19.121), pinned memory is
+    # supported but disabled by default due to a small performance regression.
+    # Set to 1 when pinned memory or UVA is required (e.g. CPU offloading
+    # or v2 model runner).
+    "VLLM_WSL2_ENABLE_PIN_MEMORY": lambda: bool(
+        int(os.getenv("VLLM_WSL2_ENABLE_PIN_MEMORY", "0"))
+    ),
     # Disable logging of vLLM logo at server startup time.
     "VLLM_DISABLE_LOG_LOGO": lambda: bool(int(os.getenv("VLLM_DISABLE_LOG_LOGO", "0"))),
     # Disable PDL for LoRA, as enabling PDL with LoRA on SM100 causes

vllm/platforms/cuda.py

@@ -7,6 +7,7 @@ pynvml. However, it should not initialize cuda context.
 from __future__ import annotations
 
 import os
+import platform
 from collections.abc import Callable
 from datetime import timedelta
 from functools import cache, lru_cache, wraps
@@ -26,7 +27,7 @@ from vllm.utils.import_utils import import_pynvml
 from vllm.utils.torch_utils import is_quantized_kv_cache
 from vllm.v1.attention.backends.registry import AttentionBackendEnum
 
-from .interface import DeviceCapability, Platform, PlatformEnum
+from .interface import DeviceCapability, Platform, PlatformEnum, in_wsl
 
 if TYPE_CHECKING:
     from vllm.config import VllmConfig
@@ -159,6 +160,21 @@ def with_nvml_context(fn: Callable[_P, _R]) -> Callable[_P, _R]:
     return wrapper
 
 
+@cache
+def _get_wsl_kernel_version() -> tuple[int, ...] | None:
+    """Return the WSL2 kernel version as a tuple, or None on parse failure.
+
+    platform.uname().release on WSL2 looks like
+    "5.15.167.4-microsoft-standard-WSL2"; we take the numeric prefix.
+    """
+    try:
+        release = platform.uname().release
+        parts = release.split("-")[0].split(".")
+        return tuple(int(x) for x in parts[:3])
+    except Exception:
+        return None
+
+
 class CudaPlatformBase(Platform):
     _enum = PlatformEnum.CUDA
     device_name: str = "cuda"
@@ -224,6 +240,27 @@ class CudaPlatformBase(Platform):
     def log_warnings(cls):
         pass
 
+    @classmethod
+    def is_pin_memory_available(cls) -> bool:
+        if in_wsl():
+            # WSL1 has no CUDA support, so being on the CUDA platform under
+            # WSL implies WSL2. Gate on kernel >= 4.19.121, the first WSL2
+            # kernel with limited pinned memory support for CUDA.
+            version = _get_wsl_kernel_version()
+            if version is None or version < (4, 19, 121):
+                logger.warning(
+                    "Using 'pin_memory=False' as WSL is detected and the "
+                    "WSL2 kernel version is below 4.19.121. This may slow "
+                    "down performance. Please run `wsl --update`."
+                )
+                return False
+            # On compatible WSL2 kernels, pinned memory is supported but
+            # disabled by default. Enable it via VLLM_WSL2_ENABLE_PIN_MEMORY=1.
+            import vllm.envs as envs
+
+            return envs.VLLM_WSL2_ENABLE_PIN_MEMORY
+        return True
+
     @classmethod
     def check_and_update_config(cls, vllm_config: VllmConfig) -> None:
         parallel_config = vllm_config.parallel_config
@@ -246,6 +283,27 @@ class CudaPlatformBase(Platform):
             )
             scheduler_config.disable_chunked_mm_input = True
 
+        if (
+            in_wsl()
+            and vllm_config.offload_config.uva.cpu_offload_gb > 0
+            and bool(vllm_config.compilation_config.cudagraph_mode)
+        ):
+            logger.warning(
+                "--cpu-offload-gb is enabled with CUDA graphs on WSL2. "
+                "This combination requires pinned (page-locked) memory "
+                "allocations. WARNING: Windows (WDDM) enforces a hard "
+                "system-wide cap of roughly 50%% of physical RAM on pinned "
+                "memory shared across ALL processes by default (limit can "
+                "changed via %%USERPROFILE%%\\.wslconfig). "
+                "Excessive use of page-locked memory can prevent Windows "
+                "from reclaiming memory under load, which can cause the "
+                "entire host OS to become unresponsive and may require a "
+                "hard reboot to recover. Proceed at your own risk. "
+                "To raise the WSL2 VM memory ceiling, increase the `memory` "
+                "setting in %%USERPROFILE%%\\.wslconfig and run "
+                "`wsl --shutdown`."
+            )
+
     @classmethod
     def get_current_memory_usage(
         cls, device: torch.types.Device | None = None

vllm/platforms/interface.py

@@ -2,6 +2,7 @@
 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
 import contextlib
 import enum
+import functools
 import os
 import platform
 import sys
@@ -30,6 +31,7 @@ else:
 logger = init_logger(__name__)
 
 
+@functools.cache
 def in_wsl() -> bool:
     # Reference: https://github.com/microsoft/WSL/issues/4071
     return "microsoft" in " ".join(platform.uname()).lower()
@@ -752,11 +754,13 @@ class Platform:
     def is_pin_memory_available(cls) -> bool:
         """Checks whether pin memory is available on the current platform."""
         if in_wsl():
-            # Pinning memory in WSL is not supported.
             # https://docs.nvidia.com/cuda/wsl-user-guide/index.html#known-limitations-for-linux-cuda-applications
+            # Pinned memory support under WSL depends on the vendor and driver
+            # version. Conservative default: return False. Platform subclasses
+            # that can verify support (e.g. CudaPlatformBase) override this.
             logger.warning(
                 "Using 'pin_memory=False' as WSL is detected. "
-                "This may slow down the performance."
+                "This may slow down performance."
             )
             return False
         return True