[https://nvbugs/6858845][fix] Let no-KV-cache runners use FlashInfer attention

tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py

@@ -15,7 +15,6 @@
     AttentionRuntimeFeatures,
 )
 from tensorrt_llm._torch.attention.backends.trtllm import TrtllmAttentionMetadata
-from tensorrt_llm._torch.attention.backends.vanilla import VanillaAttentionMetadata
 from tensorrt_llm._torch.distributed import Distributed
 from tensorrt_llm._torch.models.modeling_multimodal_mixin import _build_request_multimodal_input
 from tensorrt_llm._torch.moe.fused_moe.moe_load_balancer import (
@@ -221,13 +220,6 @@ def _prepare_inputs(
 
         # This is for no-KV-cache attention, not for dummy attention.
         if attn_metadata.kv_cache_manager is None:
-            assert isinstance(
-                attn_metadata,
-                (VanillaAttentionMetadata, TrtllmAttentionMetadata),
-            ), (
-                "Only vanilla and trtllm attention metadata are supported for "
-                "no-KV-cache attention for now"
-            )
             attn_metadata.max_seq_len = runner_config.max_seq_len
             attn_metadata.request_ids = request_ids
             attn_metadata.prepare()

tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py

@@ -46,6 +46,7 @@
 _QWEN_3_VL_DIR = llm_models_root() / "Qwen3" / "Qwen3-VL-2B-Instruct"
 _QWEN_3_VL_30B_A3B_FP8_DIR = llm_models_root(
 ) / "Qwen3" / "Qwen3-VL-30B-A3B-Instruct-FP8"
+_QWEN_2_5_VL_3B_DIR = llm_models_root() / "Qwen2.5-VL-3B-Instruct"
 
 _FAKE_QWEN3_VL_30B_A3B_FP8_SENTINEL = "qwen3_vl_30b_a3b_fp8_fake"
 _FAKE_CHECKPOINT_MARKER = ".tllm_fake_checkpoint"
@@ -1348,3 +1349,46 @@ def test_pd_disagg_multimodal_no_reuse_when_disabled():
             assert len(decode_outputs[0].outputs) > 0
             text = decode_outputs[0].outputs[0].text
             assert len(text) > 0, f"[req {req_num}] empty output"
+
+
+def _encode_images(model_dir: Path, attn_backend: str,
+                   media: list[list[str]]) -> list[torch.Tensor]:
+    encoder = MultimodalEncoder(model=model_dir,
+                                attn_backend=attn_backend,
+                                max_batch_size=len(media))
+    with encoder:
+        inputs = _load_inputs(encoder, ["Describe the image."] * len(media),
+                              media)
+        outputs = encoder.generate(inputs)
+        embeddings = []
+        for output in outputs:
+            handles = output.disaggregated_params.multimodal_embedding_handles
+            assert handles is not None
+            embeddings.append(
+                torch.cat([
+                    SharedTensorContainer.from_dict(
+                        handle).get_local_view().float().cpu()
+                    for handle in handles
+                ]))
+    return embeddings
+
+
+@pytest.mark.threadleak(enabled=False)
+def test_encoder_flashinfer_attention_matches_trtllm():
+    """The encoder-only engine runs with FLASHINFER attention and matches TRTLLM."""
+    # get_attention_backend() falls back to TRTLLM when FlashInfer cannot be
+    # imported, which would turn this test into a TRTLLM self-comparison.
+    pytest.importorskip("flashinfer")
+    if not _QWEN_2_5_VL_3B_DIR.is_dir():
+        pytest.skip(f"Checkpoint not found: {_QWEN_2_5_VL_3B_DIR}")
+    media = [[image] for image in example_images]
+    reference = _encode_images(_QWEN_2_5_VL_3B_DIR, "TRTLLM", media)
+    flashinfer = _encode_images(_QWEN_2_5_VL_3B_DIR, "FLASHINFER", media)
+
+    assert len(flashinfer) == len(reference) == len(media)
+    for ref, test in zip(reference, flashinfer):
+        assert test.shape == ref.shape
+        cos = torch.nn.functional.cosine_similarity(ref.flatten(),
+                                                    test.flatten(),
+                                                    dim=0)
+        assert cos.item() > 0.99, f"cosine similarity {cos.item():.6f}"

tests/unittest/_torch/pyexecutor/engine/runners/test_no_kv_cache.py

@@ -84,7 +84,6 @@ def _prepare(
     torch.Tensor,
 ]:
     monkeypatch.setattr(no_kv_cache_module, "prefer_pinned", lambda: False)
-    monkeypatch.setattr(no_kv_cache_module, "VanillaAttentionMetadata", _AttentionMetadata)
     attn_metadata = _AttentionMetadata(sum(len(request.get_tokens(0)) for request in requests))
     scheduled_requests = SimpleNamespace(
         context_requests=requests,