diff --git a/tests/v1/cudagraph/test_sm70_graph_gates_pre_ampere.py b/tests/v1/cudagraph/test_sm70_graph_gates_pre_ampere.py new file mode 100644 index 0000000000..0f9635d22a --- /dev/null +++ b/tests/v1/cudagraph/test_sm70_graph_gates_pre_ampere.py @@ -0,0 +1,59 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""The SM70 graph gates are pre-Ampere gates on the worker's own device. + +Volta and Turing take the same graph tunings. The gate must ask the device +this process selected, not index 0 of the visibility list, which on a mixed +rig belongs to a different worker. +""" + +import pytest + +from vllm.platforms.interface import DeviceCapability +from vllm.v1.worker.gpu import cudagraph_utils as cg + +VOLTA = DeviceCapability(7, 0) +TURING = DeviceCapability(7, 5) +AMPERE = DeviceCapability(8, 0) + + +def _fake_devices(monkeypatch, capabilities, current: int) -> list[int]: + asked: list[int] = [] + + def get_device_capability(device_id: int = 0) -> DeviceCapability: + asked.append(device_id) + return capabilities[device_id] + + monkeypatch.setattr(cg.current_platform, "is_cuda", lambda: True) + monkeypatch.setattr( + cg.current_platform, "get_device_capability", get_device_capability + ) + monkeypatch.setattr(cg.torch.accelerator, "current_device_index", lambda: current) + return asked + + +@pytest.mark.parametrize( + ("capability", "expected"), + [(VOLTA, True), (TURING, True), (AMPERE, False)], + ids=["volta", "turing", "ampere"], +) +def test_pre_ampere_by_capability(monkeypatch, capability, expected): + _fake_devices(monkeypatch, [capability], current=0) + assert cg._worker_device_is_pre_ampere() is expected + + +def test_gate_asks_the_workers_own_device(monkeypatch): + # Index 0 is Ampere, this worker sits on the Turing card at index 1. + asked = _fake_devices(monkeypatch, [AMPERE, TURING], current=1) + assert cg._worker_device_is_pre_ampere() is True + assert asked == [1] + + +def test_gate_is_off_without_cuda(monkeypatch): + monkeypatch.setattr(cg.current_platform, "is_cuda", lambda: False) + monkeypatch.setattr( + cg.current_platform, + "get_device_capability", + lambda device_id=0: pytest.fail("no device query without CUDA"), + ) + assert cg._worker_device_is_pre_ampere() is False diff --git a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py index 16fd480cc6..fe98bd8a63 100644 --- a/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py +++ b/tests/v1/cudagraph/test_sm70_mtp_split_cudagraphs.py @@ -8,6 +8,7 @@ import torch from vllm.config import CompilationConfig, CUDAGraphMode +from vllm.platforms.interface import DeviceCapability from vllm.v1.worker.gpu.cudagraph_utils import CudaGraphManager @@ -31,6 +32,7 @@ def _make_config(max_num_seqs: int, verifier_sizes: list[int]): ) +@pytest.mark.parametrize("capability", [(7, 0), (7, 5)], ids=["volta", "turing"]) @pytest.mark.parametrize( ("max_num_seqs", "request_sizes", "verifier_sizes"), [ @@ -43,15 +45,21 @@ def test_split_managers_keep_exact_full_graph_shapes( max_num_seqs: int, request_sizes: list[int], verifier_sizes: list[int], + capability: tuple[int, int], ): monkeypatch.setenv("VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", "1") monkeypatch.setattr( "vllm.v1.worker.gpu.cudagraph_utils.current_platform.is_cuda", lambda: True, ) + # The gate asks the worker's own device; answer for it only. monkeypatch.setattr( - "vllm.v1.worker.gpu.cudagraph_utils.current_platform.is_device_capability", - lambda capability: capability == (7, 0), + "vllm.v1.worker.gpu.cudagraph_utils.current_platform.get_device_capability", + lambda device_id=0: DeviceCapability(*capability), + ) + monkeypatch.setattr( + "vllm.v1.worker.gpu.cudagraph_utils.torch.accelerator.current_device_index", + lambda: 0, ) monkeypatch.setattr( "vllm.v1.worker.gpu.cudagraph_utils.current_platform.get_global_graph_pool", diff --git a/tests/v1/worker/test_sm70_long_attention_graphs.py b/tests/v1/worker/test_sm70_long_attention_graphs.py index d4999facb8..4a8e5b02fe 100644 --- a/tests/v1/worker/test_sm70_long_attention_graphs.py +++ b/tests/v1/worker/test_sm70_long_attention_graphs.py @@ -9,6 +9,7 @@ import torch from vllm.config.compilation import CompilationConfig, CUDAGraphMode +from vllm.platforms.interface import DeviceCapability from vllm.v1.attention.ops.sm70_e4m3_long import BUILTIN_MAX_CONTEXT from vllm.v1.worker.gpu import cudagraph_utils as cg from vllm.v1.worker.gpu.cudagraph_utils import ( @@ -94,24 +95,32 @@ def original(*args, **kwargs): @pytest.mark.parametrize( - "enabled,method,sm70,target,sequence_parallel,expect_tail", + "enabled,method,capability,target,sequence_parallel,expect_tail", [ - (True, "dflash", True, True, False, True), - (False, "dflash", True, True, False, False), - (True, "mtp", True, True, False, False), - (True, "dflash", False, True, False, False), - (True, "dflash", True, False, False, False), - (True, "dflash", True, True, True, False), + (True, "dflash", (7, 0), True, False, True), + # Turing takes the same tail graphs as Volta. + (True, "dflash", (7, 5), True, False, True), + (False, "dflash", (7, 0), True, False, False), + (True, "mtp", (7, 0), True, False, False), + (True, "dflash", (8, 0), True, False, False), + (True, "dflash", (7, 0), False, False, False), + (True, "dflash", (7, 0), True, True, False), ], ) def test_tail_capture_and_dispatch_from_real_initialization( - monkeypatch, enabled, method, sm70, target, sequence_parallel, expect_tail + monkeypatch, enabled, method, capability, target, sequence_parallel, expect_tail ): monkeypatch.setenv("VLLM_SM70_DFLASH2_TAIL_CUDAGRAPHS", str(int(enabled))) monkeypatch.setenv("VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS", "0") monkeypatch.delenv("VLLM_SM70_E4M3_LONG_ATTENTION_MANIFEST", raising=False) monkeypatch.setattr(cg.current_platform, "is_cuda", lambda: True) - monkeypatch.setattr(cg.current_platform, "is_device_capability", lambda cap: sm70) + # The gate asks the worker's own device; answer for it only. + monkeypatch.setattr( + cg.current_platform, + "get_device_capability", + lambda device_id=0: DeviceCapability(*capability), + ) + monkeypatch.setattr(cg.torch.accelerator, "current_device_index", lambda: 0) monkeypatch.setattr(cg.current_platform, "get_global_graph_pool", lambda: None) monkeypatch.setattr( cg, diff --git a/vllm/v1/worker/gpu/cudagraph_utils.py b/vllm/v1/worker/gpu/cudagraph_utils.py index 9a03f604ff..11b63b2f9c 100644 --- a/vllm/v1/worker/gpu/cudagraph_utils.py +++ b/vllm/v1/worker/gpu/cudagraph_utils.py @@ -56,14 +56,32 @@ def get_explicit_cudagraph_memory_reserve(cudagraph_mode: CUDAGraphMode) -> int: return reserve_bytes +def _worker_device_is_pre_ampere() -> bool: + """Whether this worker's own device is Volta or Turing. + + The SM70 graph tunings are pre-Ampere tunings: Turing runs the same + kernels, the same fp16 contract and the same compile graph as Volta. + Asking index 0 of the visibility list would answer for a different + worker on a mixed rig, so ask the device this process has selected. + """ + if not current_platform.is_cuda(): + return False + capability = current_platform.get_device_capability( + device_id=torch.accelerator.current_device_index() + ) + return capability is not None and (capability.major, capability.minor) in ( + (7, 0), + (7, 5), + ) + + def _use_split_sm70_mtp_cudagraphs(vllm_config: VllmConfig) -> bool: speculative_config = vllm_config.speculative_config return bool( envs.VLLM_SM70_MTP_SPLIT_DRAFT_CUDAGRAPHS and speculative_config is not None and speculative_config.method == "mtp" - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) + and _worker_device_is_pre_ampere() ) @@ -171,8 +189,7 @@ def __init__( and speculative_config is not None and speculative_config.method == "dflash" and decode_query_len == 8 - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) + and _worker_device_is_pre_ampere() and not self.compilation_config.pass_config.enable_sp ) if self._sm70_dflash2_tail_graphs: @@ -361,8 +378,7 @@ def capture( ) if ( envs.VLLM_SM70_FLASH_V100_0DOT3_COMPILE_GRAPH - and current_platform.is_cuda() - and current_platform.is_device_capability((7, 0)) + and _worker_device_is_pre_ampere() ): logger.info_once( "Running SM70 Flash-V100 compile full-graph "