[#20586][fix] DeepSeek-V4.1 vision: skip cuDNN SDPA on SM107 with cuDNN < 9.25

tensorrt_llm/_torch/models/deepseek_v41_vision.py

@@ -4,18 +4,55 @@
 
 from __future__ import annotations
 
+import contextlib
+import functools
 from collections.abc import Callable
 from typing import TYPE_CHECKING
 
 import torch
 import torch.nn.functional as F
 from torch import nn
+from torch.nn.attention import SDPBackend, sdpa_kernel
 
 if TYPE_CHECKING:
     from ..configs.deepseek_v41 import DeepseekV41VisionConfig
 
 __all__ = ["DeepseekV41VisionModel", "DeepseekV41Aligner"]
 
+# PyTorch dispatches SDPA to cuDNN first on SM 10.x and does not fall back when
+# plan building fails. cuDNN before 9.25 has no SM 10.7 SDPA forward engine that
+# builds: its runtime-compiled candidates all fail NVRTC compilation.
+_NON_CUDNN_SDPA_BACKENDS = [
+    SDPBackend.FLASH_ATTENTION,
+    SDPBackend.EFFICIENT_ATTENTION,
+    SDPBackend.MATH,
+]
+_FIRST_CUDNN_WITH_SM107_SDPA = 92500
+
+
+@functools.cache
+def _excludes_cudnn_sdpa(device_index: int) -> bool:
+    cudnn_version = torch.backends.cudnn.version() or 0
+    return (
+        torch.cuda.get_device_capability(device_index) == (10, 7)
+        and cudnn_version < _FIRST_CUDNN_WITH_SM107_SDPA
+    )
+
+
+def _vision_sdpa_context(device: torch.device) -> contextlib.AbstractContextManager:
+    """Disable only cuDNN SDPA where it cannot build a plan; keep the caller's other backends."""
+    if device.type != "cuda" or not _excludes_cudnn_sdpa(
+        device.index if device.index is not None else torch.cuda.current_device()
+    ):
+        return contextlib.nullcontext()
+    enabled = {
+        SDPBackend.FLASH_ATTENTION: torch.backends.cuda.flash_sdp_enabled(),
+        SDPBackend.EFFICIENT_ATTENTION: torch.backends.cuda.mem_efficient_sdp_enabled(),
+        SDPBackend.MATH: torch.backends.cuda.math_sdp_enabled(),
+    }
+    backends = [backend for backend in _NON_CUDNN_SDPA_BACKENDS if enabled[backend]]
+    return sdpa_kernel(backends) if backends else contextlib.nullcontext()
+
 
 def _get_vision_cos_sin(
     n_h: int, n_w: int, dim: int, theta: float, device: torch.device
@@ -82,11 +119,12 @@ def forward(self, x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torc
         )
         query = _apply_vision_rotary(query, cos, sin)
         key = _apply_vision_rotary(key, cos, sin)
-        attended = F.scaled_dot_product_attention(
-            query.transpose(0, 1).unsqueeze(0),
-            key.transpose(0, 1).unsqueeze(0),
-            value.transpose(0, 1).unsqueeze(0),
-        )
+        with _vision_sdpa_context(query.device):
+            attended = F.scaled_dot_product_attention(
+                query.transpose(0, 1).unsqueeze(0),
+                key.transpose(0, 1).unsqueeze(0),
+                value.transpose(0, 1).unsqueeze(0),
+            )
         return self.wo(attended.squeeze(0).transpose(0, 1).reshape(num_patches, -1))
 
 

tests/integration/test_lists/test-db/l0_gr100.yml

@@ -14,3 +14,4 @@ l0_gr100:
       backend: pytorch
   tests:
   - accuracy/test_llm_api_pytorch.py::TestDeepSeekV3Lite::test_bfloat16_flashinfer[enable_chunked_prefill=True]
+  - unittest/_torch/modeling/test_deepseek_v41_vision_sm107.py

tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py

@@ -14,6 +14,7 @@
 from transformers import PretrainedConfig
 
 from tensorrt_llm._torch.configs.deepseek_v41 import DeepseekV41VisionConfig
+from tensorrt_llm._torch.models import deepseek_v41_vision
 from tensorrt_llm._torch.models.deepseek_v41_processor import DeepseekV41InputProcessor
 from tensorrt_llm._torch.models.deepseek_v41_vision import (
     DeepseekV41Aligner,
@@ -203,6 +204,59 @@ def test_complete_vision_and_aligner_match_reference(
     assert model(patches.bfloat16(), 4, 5).dtype == torch.bfloat16
 
 
+def _sdpa_flags() -> dict[str, bool]:
+    return {
+        "cudnn": torch.backends.cuda.cudnn_sdp_enabled(),
+        "flash": torch.backends.cuda.flash_sdp_enabled(),
+        "efficient": torch.backends.cuda.mem_efficient_sdp_enabled(),
+        "math": torch.backends.cuda.math_sdp_enabled(),
+    }
+
+
+@pytest.mark.parametrize(
+    ("device", "capability", "cudnn_version", "excludes_cudnn"),
+    [
+        ("cuda:0", (10, 7), 92400, True),
+        ("cuda:0", (10, 7), 92500, False),
+        ("cuda:0", (10, 0), 92400, False),
+        ("cuda:0", (10, 3), 92400, False),
+        ("cpu", (10, 7), 92400, False),
+    ],
+)
+def test_vision_sdpa_skips_cudnn_only_on_sm107_with_old_cudnn(
+    monkeypatch, device, capability, cudnn_version, excludes_cudnn
+) -> None:
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda index=None: capability)
+    monkeypatch.setattr(torch.backends.cudnn, "version", lambda: cudnn_version)
+    deepseek_v41_vision._excludes_cudnn_sdpa.cache_clear()
+    try:
+        before = _sdpa_flags()
+        with deepseek_v41_vision._vision_sdpa_context(torch.device(device)):
+            inside = _sdpa_flags()
+        assert _sdpa_flags() == before
+        if excludes_cudnn:
+            assert inside == {**before, "cudnn": False}
+        else:
+            assert inside == before
+    finally:
+        deepseek_v41_vision._excludes_cudnn_sdpa.cache_clear()
+
+
+def test_vision_sdpa_keeps_backends_the_caller_disabled(monkeypatch) -> None:
+    monkeypatch.setattr(torch.cuda, "get_device_capability", lambda index=None: (10, 7))
+    monkeypatch.setattr(torch.backends.cudnn, "version", lambda: 92400)
+    deepseek_v41_vision._excludes_cudnn_sdpa.cache_clear()
+    try:
+        with torch.nn.attention.sdpa_kernel(
+            [torch.nn.attention.SDPBackend.CUDNN_ATTENTION, torch.nn.attention.SDPBackend.MATH]
+        ):
+            with deepseek_v41_vision._vision_sdpa_context(torch.device("cuda:0")):
+                inside = _sdpa_flags()
+        assert inside == {"cudnn": False, "flash": False, "efficient": False, "math": True}
+    finally:
+        deepseek_v41_vision._excludes_cudnn_sdpa.cache_clear()
+
+
 def test_wrapper_encodes_image_spans_and_tracks_active_positions(
     vision_config: DeepseekV41VisionConfig,
 ) -> None:

tests/unittest/_torch/modeling/test_deepseek_v41_vision_sm107.py

@@ -0,0 +1,31 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+import pytest
+import torch
+from torch.nn.attention import SDPBackend, sdpa_kernel
+
+from tensorrt_llm._torch.configs.deepseek_v41 import DeepseekV41VisionConfig
+from tensorrt_llm._torch.models.deepseek_v41_vision import DeepseekV41VisionModel
+
+
+@pytest.mark.skipif(
+    not torch.cuda.is_available() or torch.cuda.get_device_capability() != (10, 7),
+    reason="SM 10.7 only",
+)
+def test_vision_attention_runs_on_sm107() -> None:
+    torch.manual_seed(0)
+    config = DeepseekV41VisionConfig(
+        num_hidden_layers=1,
+        hidden_size=1024,
+        num_attention_heads=16,
+        intermediate_size=64,
+        patch_size=2,
+        downsample_ratio=3,
+    )
+    model = DeepseekV41VisionModel(config).cuda().bfloat16().eval()
+    patches = torch.randn(32 * 32, 3, 2, 2, device="cuda", dtype=torch.bfloat16)
+    with torch.inference_mode():
+        actual = model(patches, 32, 32)
+        with sdpa_kernel([SDPBackend.MATH]):
+            expected = model(patches, 32, 32)
+    torch.testing.assert_close(actual, expected, rtol=2e-2, atol=2e-2)