diff --git a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py index e78f3b0301fe..32a098caaaea 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py +++ b/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() diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py index e1ece0697ecd..75f12899bf6a 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py +++ b/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}" diff --git a/tests/unittest/_torch/pyexecutor/engine/runners/test_no_kv_cache.py b/tests/unittest/_torch/pyexecutor/engine/runners/test_no_kv_cache.py index 1bb0cfc9bbc7..70c8863b49f0 100644 --- a/tests/unittest/_torch/pyexecutor/engine/runners/test_no_kv_cache.py +++ b/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,