[Bugfix][SM70/SM75] Gate the SM70 graph tunings on the worker's own pre-Ampere device

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

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",

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,

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 "