From 127ea504c8fbdcd24a3443849f5bcaea099a1ebb Mon Sep 17 00:00:00 2001 From: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> Date: Fri, 23 Oct 2026 02:30:35 -0700 Subject: [PATCH 1/2] [https://nvbugs/6858845][fix] Let no-KV-cache runners use FlashInfer attention NoKVCacheRunner.prepare_inputs asserted that the attention metadata was Vanilla or TRTLLM before calling prepare() without a KV cache manager. FlashInferAttentionMetadata.prepare() has its own kv_cache_manager=None branch for context-only batches, and the no-KV- cache runner only schedules context requests. Because setup_attn_metadata always builds the metadata with kv_cache_manager=None, prepare_inputs raised on every call when attn_backend was FLASHINFER. Remove the runner-side type assertion, so each backend's prepare() handles the no-KV-cache case itself. Tests: - test_mm_encoder_standalone.py: new test_encoder_flashinfer_attention_matches_trtllm (Qwen2.5-VL-3B encoder-only, FLASHINFER vs TRTLLM embeddings, cosine > 0.99). - test_no_kv_cache.py: drop the monkeypatch that made the stub metadata pass the removed assertion. Signed-off-by: trtllm-agent <296075020+trtllm-agent@users.noreply.github.com> --- .../pyexecutor/engine/runners/no_kv_cache.py | 8 ---- .../executor/engine/test_no_kv_cache.py | 1 - .../multimodal/test_mm_encoder_standalone.py | 39 +++++++++++++++++++ 3 files changed, 39 insertions(+), 9 deletions(-) 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 674b57281447..7f68838b0eb9 100644 --- a/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py +++ b/tensorrt_llm/_torch/pyexecutor/engine/runners/no_kv_cache.py @@ -12,7 +12,6 @@ from tensorrt_llm._torch.attention.backends.interface import AttentionMetadata 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 ( @@ -222,13 +221,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/executor/engine/test_no_kv_cache.py b/tests/unittest/_torch/executor/engine/test_no_kv_cache.py index e4121271172a..636934381b68 100644 --- a/tests/unittest/_torch/executor/engine/test_no_kv_cache.py +++ b/tests/unittest/_torch/executor/engine/test_no_kv_cache.py @@ -70,7 +70,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, diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py index 5df86946432d..579c53e45845 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" @@ -1330,3 +1331,41 @@ 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.""" + 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}" From aec89aec0f7f9c7850be205e512803e5b263a33d Mon Sep 17 00:00:00 2001 From: berryz Date: Fri, 30 Oct 2026 16:11:25 +0800 Subject: [PATCH 2/2] [https://nvbugs/6858845][test] Skip the FlashInfer encoder test when FlashInfer or the model is missing get_attention_backend() falls back to TRTLLM when FlashInfer cannot be imported, so test_encoder_flashinfer_attention_matches_trtllm would compare TRTLLM against itself and pass without covering the FLASHINFER path. Skip the test in that case, and when the Qwen2.5-VL-3B-Instruct checkpoint is not present under the model root, instead of failing during encoder setup. Signed-off-by: berryz --- .../unittest/_torch/multimodal/test_mm_encoder_standalone.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py index 579c53e45845..46d24ecfcd1e 100644 --- a/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py +++ b/tests/unittest/_torch/multimodal/test_mm_encoder_standalone.py @@ -1358,6 +1358,11 @@ def _encode_images(model_dir: Path, attn_backend: str, @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)