[TRTLLM-13445][test] Add SageAttention backend-routing and SAGE recip… (NVIDIA#15746)

Signed-off-by: Ying Guo <244492186+yingguo-trt@users.noreply.github.com>

Author: yingguo-trt

tests/unittest/_torch/visual_gen/test_attention_integration.py

@@ -26,6 +26,7 @@
     RingAttention,
     UlyssesAttention,
 )
+from tensorrt_llm._torch.visual_gen.attention_backend.trtllm import TrtllmAttention
 from tensorrt_llm._torch.visual_gen.attention_backend.vanilla import VanillaAttention
 from tensorrt_llm._torch.visual_gen.config import (
     DiffusionModelConfig,
@@ -292,6 +293,43 @@ def test_attn2d_with_sequence_parallel_disabled_allowed(self):
         assert isinstance(attn.attn, VanillaAttention)
 
 
+def _build_sage_routed_attention(qkv_mode: QKVMode):
+    """Build a TRTLLM-SAGE-configured Attention for backend-routing checks."""
+    quant_cfg = QuantAttentionConfig(
+        qk_dtype="int8", q_block_size=1, k_block_size=16, v_block_size=1
+    )
+    config = create_model_config(
+        hidden_size=512,
+        num_heads=4,
+        head_dim=128,
+        attn_backend="TRTLLM",
+        quant_attention_config=quant_cfg,
+        skip_create_weights_in_init=True,
+    )
+    attn = Attention(
+        hidden_size=512,
+        num_attention_heads=4,
+        head_dim=128,
+        qkv_mode=qkv_mode,
+        config=config,
+    )
+    return attn, quant_cfg
+
+
+class TestSageAttentionBackendRouting:
+    def test_self_attention_uses_trtllm_sage_backend(self):
+        attn, quant_cfg = _build_sage_routed_attention(QKVMode.FUSE_QKV)
+        assert attn.attn_backend == "TRTLLM"
+        assert isinstance(attn.attn, TrtllmAttention)
+        assert attn.attn.quant_attention_config == quant_cfg
+        assert not attn.attn.support_fused_qkv()
+
+    def test_cross_attention_with_sage_config_falls_back_to_vanilla(self):
+        attn, _ = _build_sage_routed_attention(QKVMode.SEPARATE_QKV)
+        assert attn.attn_backend == "VANILLA"
+        assert isinstance(attn.attn, VanillaAttention)
+
+
 # ============================================================================
 # Test functions
 # ============================================================================

tests/unittest/_torch/visual_gen/test_visual_gen_args.py

@@ -89,11 +89,24 @@ def test_quant_config_rejected_when_unsupported(self):
                 ),
             )
 
-    def test_supported_quant_config_sage(self):
+    @pytest.mark.parametrize(
+        ("qk_dtype", "q_block_size", "k_block_size", "v_block_size"),
+        [
+            ("int8", 1, 1, 1),
+            ("int8", 1, 4, 1),
+            ("int8", 1, 16, 1),
+            ("fp8", 1, 1, 1),
+            ("fp8", 1, 4, 1),
+        ],
+    )
+    def test_supported_quant_config_sage(self, qk_dtype, q_block_size, k_block_size, v_block_size):
         attention = AttentionConfig(
             backend="TRTLLM",
             quant_attention_config=QuantAttentionConfig(
-                qk_dtype="int8", q_block_size=1, k_block_size=16, v_block_size=1
+                qk_dtype=qk_dtype,
+                q_block_size=q_block_size,
+                k_block_size=k_block_size,
+                v_block_size=v_block_size,
             ),
         )