[None][fix] cuDNN attention on Rubin (SM107): allow FP8 and fix LSE layout

tensorrt_llm/_torch/visual_gen/attention_backend/cudnn.py

@@ -39,7 +39,6 @@
 import cudnn
 import torch
 
-from tensorrt_llm._utils import get_sm_version
 from tensorrt_llm.logger import logger
 from tensorrt_llm.visual_gen.args import QuantAttentionConfig
 
@@ -208,6 +207,12 @@ class _CuDNNProblemShape:
 # ============================================================================
 
 
+def _device_sm_version(device: torch.device) -> int:
+    """SM version (e.g. 107) of ``device``, not of CUDA device 0."""
+    major, minor = torch.cuda.get_device_capability(device)
+    return major * 10 + minor
+
+
 class CuDNNAttention(AttentionBackend):
     """cuDNN SDPA backend for visual generation.
 
@@ -300,8 +305,13 @@ def _get_handle(cls, device: torch.device) -> Any:
     def check_hardware_compatibility(
         cls, device: torch.device, quant_dtype: str | None = None
     ) -> None:
-        if get_sm_version() not in (100, 103) and quant_dtype is not None:
-            raise RuntimeError("cuDNN quantized attention requires NVIDIA Blackwell-class GPU.")
+        if quant_dtype is None:
+            return
+        if _device_sm_version(device) not in (100, 103, 107):
+            raise RuntimeError(
+                "cuDNN quantized attention requires an NVIDIA Blackwell- or Rubin-class GPU "
+                "(SM 100, 103 or 107)."
+            )
 
     @classmethod
     def check_library_feature(cls, quant_dtype: str | None = None) -> None:
@@ -708,8 +718,8 @@ def _run(
 
         self._execute_graph(bundle, tensor_map, device)
 
-        # Stats is packed [B, H, S, 1]; other backends expose LSE as [B, S, H].
-        lse = None if stats is None else stats.squeeze(-1).transpose(1, 2).contiguous()
+        # Stats is packed [B, H, S, 1]; LSE consumers (Ring/Attention2D) expect [B, H, S].
+        lse = None if stats is None else stats.squeeze(-1)
         return output, lse
 
     def forward(
@@ -752,7 +762,7 @@ def forward_with_lse(
 
         Returns:
             output: ``[B, S_q, H, D_v]``
-            lse: ``[B, S_q, H]`` float32
+            lse: ``[B, H, S_q]`` float32
         """
         output, lse = self._run(
             q, k, v, is_causal=self._resolve_mask(attention_mask, key_padding_mask), with_lse=True

tests/unittest/_torch/visual_gen/multi_gpu/test_attn2d_attention.py

@@ -35,6 +35,7 @@
 
 from tensorrt_llm._torch.attention.backends.interface import PredefinedAttentionMask
 from tensorrt_llm._torch.visual_gen.attention_backend import Attention2DAttention
+from tensorrt_llm._torch.visual_gen.attention_backend.cudnn import CuDNNAttention
 from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import FlashAttn4Attention
 from tensorrt_llm._torch.visual_gen.attention_backend.flash_attn4 import _flash_attn_fwd as _fa4_fwd
 from tensorrt_llm._torch.visual_gen.attention_backend.interface import AttentionTensorLayout
@@ -285,6 +286,51 @@ def _logic_attn2d_vs_standard(rank, world_size):
     )
 
 
+def _logic_attn2d_cudnn_vs_standard(rank, world_size):
+    """cuDNN inner backend: its LSE layout must match what Attention2D combines."""
+    row_size, col_size = 2, 2
+    batch, num_heads, head_dim = 1, 4, 64
+    seq_per_rank = 64
+    seq_full = seq_per_rank * world_size
+    device = torch.device(f"cuda:{rank}")
+
+    row_pg, col_pg = _make_process_groups(rank, world_size, row_size, col_size)
+
+    torch.manual_seed(42)
+    q_full, k_full, v_full = (
+        torch.randn(batch, seq_full, num_heads, head_dim, device=device, dtype=torch.bfloat16)
+        for _ in range(3)
+    )
+    shard = slice(rank * seq_per_rank, (rank + 1) * seq_per_rank)
+
+    inner = CuDNNAttention(num_heads=num_heads, head_dim=head_dim, dtype=torch.bfloat16)
+    try:
+        attn = Attention2DAttention(inner, row_pg, col_pg)
+    except ImportError:
+        pytest.skip("flash_attn_combine JIT kernels not available")
+
+    output = attn(
+        q_full[:, shard].contiguous(),
+        k_full[:, shard].contiguous(),
+        v_full[:, shard].contiguous(),
+        batch_size=batch,
+    )
+
+    reference = F.scaled_dot_product_attention(
+        q_full.transpose(1, 2).float(),
+        k_full.transpose(1, 2).float(),
+        v_full.transpose(1, 2).float(),
+        scale=1.0 / math.sqrt(head_dim),
+    ).transpose(1, 2)
+    torch.testing.assert_close(
+        output.float(),
+        reference[:, shard],
+        rtol=2e-2,
+        atol=2e-2,
+        msg=f"Rank {rank}: Attention2D(cuDNN) output differs from standard attention",
+    )
+
+
 def _logic_attn2d_replicated_kv_unequal_lengths_impl(
     rank, world_size, use_fa4, row_size=1, col_size=2, unequal_lengths=True
 ):
@@ -663,6 +709,12 @@ def test_attn2d_vs_standard_attention(self):
         """Attention2DAttention output matches standard full-sequence SDPA (2x2 mesh)."""
         run_test_in_distributed(world_size=4, test_fn=_logic_attn2d_vs_standard, use_cuda=True)
 
+    def test_attn2d_cudnn_vs_standard_attention(self):
+        """Attention2D with the cuDNN inner backend matches standard SDPA (2x2 mesh)."""
+        run_test_in_distributed(
+            world_size=4, test_fn=_logic_attn2d_cudnn_vs_standard, use_cuda=True
+        )
+
     def test_attn2d_replicated_kv_unequal_lengths(self):
         """A 1x2 mesh matches full attention with unequal context lengths."""
         run_test_in_distributed(

tests/unittest/_torch/visual_gen/test_attention_cudnn.py

@@ -13,6 +13,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+from unittest.mock import patch
+
 import pytest
 import torch
 import torch.nn.functional as F
@@ -76,7 +78,7 @@ def _reference(q, k, v, is_causal):
         seq_q, seq_kv = q.shape[2], k.shape[2]
         causal_mask = torch.ones(seq_q, seq_kv, device=q.device, dtype=torch.bool)
         logits = logits.masked_fill(causal_mask.triu(seq_kv - seq_q + 1), float("-inf"))
-    return out.transpose(1, 2), torch.logsumexp(logits, dim=-1).transpose(1, 2)
+    return out.transpose(1, 2), torch.logsumexp(logits, dim=-1)
 
 
 @pytest.mark.parametrize("quant_dtype", list(QUANT_CONFIGS), ids=["unquantized", "fp8", "mxfp8"])
@@ -105,7 +107,7 @@ def test_cudnn_attention(quant_dtype, shape, is_causal):
     ref_out, ref_lse = _reference(q, k, v, is_causal)
     assert output.shape == (batch, seq_q, num_heads, head_dim)
     assert output.dtype == torch.bfloat16
-    assert lse.shape == (batch, seq_q, num_heads)
+    assert lse.shape == (batch, num_heads, seq_q)
 
     cosine = F.cosine_similarity(output.float().flatten(), ref_out.flatten(), dim=0).item()
     assert cosine > MIN_COSINE[quant_dtype], (
@@ -165,6 +167,35 @@ def dequant_e8m0(scale_factors):
     assert rel_err_v < 0.05, f"V MXFP8 round-trip error {rel_err_v} exceeds e4m3 block noise"
 
 
+_SM_VERSION = "tensorrt_llm._torch.visual_gen.attention_backend.cudnn._device_sm_version"
+
+
+@pytest.mark.parametrize("sm_version", [100, 103, 107])
+@pytest.mark.parametrize("quant_dtype", ["fp8", "mxfp8"])
+def test_cudnn_quantized_attention_allowed_on_supported_sm(sm_version, quant_dtype):
+    """Quantized attention passes the hardware gate on Blackwell and Rubin (SM107)."""
+    with patch(_SM_VERSION, return_value=sm_version):
+        CuDNNAttention.check_hardware_compatibility(torch.device("cuda:0"), quant_dtype)
+
+
+@pytest.mark.parametrize("sm_version", [90, 120])
+def test_cudnn_quantized_attention_rejected_on_unsupported_sm(sm_version):
+    """Quantized attention fails fast on other SMs; unquantized attention is not gated."""
+    with patch(_SM_VERSION, return_value=sm_version):
+        with pytest.raises(RuntimeError, match="SM 100, 103 or 107"):
+            CuDNNAttention.check_hardware_compatibility(torch.device("cuda:0"), "fp8")
+        CuDNNAttention.check_hardware_compatibility(torch.device("cuda:0"), None)
+
+
+def test_cudnn_hardware_gate_checks_target_device():
+    """On a mixed-GPU host the gate follows the tensors' device, not CUDA device 0."""
+    sm_by_index = {0: 90, 1: 107}
+    with patch(_SM_VERSION, side_effect=lambda device: sm_by_index[device.index]):
+        CuDNNAttention.check_hardware_compatibility(torch.device("cuda:1"), "fp8")
+        with pytest.raises(RuntimeError, match="SM 100, 103 or 107"):
+            CuDNNAttention.check_hardware_compatibility(torch.device("cuda:0"), "fp8")
+
+
 def test_cudnn_graph_cache_reuses_plans():
     """Repeated calls with the same geometry reuse one compiled graph."""
     _require_cudnn(None)