From 80828b1796cc51ddbf325351cf08e0d20e5f9aa5 Mon Sep 17 00:00:00 2001 From: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> Date: Mon, 19 Oct 2026 01:01:24 -0700 Subject: [PATCH] [#20586][fix] DeepSeek-V4.1 vision: skip cuDNN SDPA on SM107 with cuDNN < 9.25 PyTorch prefers cuDNN SDPA on SM 10.x and does not fall back when plan building fails. cuDNN < 9.25 cannot build any SDPA forward plan on SM 10.7, so the vision tower fails during KV-cache profiling. Exclude only the cuDNN backend around the vision SDPA call in that configuration. Signed-off-by: Zihua Wu <13583761+lucifer1004@users.noreply.github.com> --- .../_torch/models/deepseek_v41_vision.py | 48 +++++++++++++++-- .../test_lists/test-db/l0_gr100.yml | 1 + .../modeling/test_deepseek_v41_multimodal.py | 54 +++++++++++++++++++ .../test_deepseek_v41_vision_sm107.py | 31 +++++++++++ 4 files changed, 129 insertions(+), 5 deletions(-) create mode 100644 tests/unittest/_torch/modeling/test_deepseek_v41_vision_sm107.py diff --git a/tensorrt_llm/_torch/models/deepseek_v41_vision.py b/tensorrt_llm/_torch/models/deepseek_v41_vision.py index 86795ae7bf77..cc3f034b8136 100644 --- a/tensorrt_llm/_torch/models/deepseek_v41_vision.py +++ b/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)) diff --git a/tests/integration/test_lists/test-db/l0_gr100.yml b/tests/integration/test_lists/test-db/l0_gr100.yml index 45cd6e6a35a8..dfdc557a799c 100644 --- a/tests/integration/test_lists/test-db/l0_gr100.yml +++ b/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 diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py b/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py index c80425c5deb9..0f3db9eae835 100644 --- a/tests/unittest/_torch/modeling/test_deepseek_v41_multimodal.py +++ b/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: diff --git a/tests/unittest/_torch/modeling/test_deepseek_v41_vision_sm107.py b/tests/unittest/_torch/modeling/test_deepseek_v41_vision_sm107.py new file mode 100644 index 000000000000..631cb0763560 --- /dev/null +++ b/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)