A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 3, 01:45 UTC
Correcting the mock classes for MM GC tests (#43321)
Signed-off-by: Weida Hong <wdhongtw@google.com>
Weida Hong committed
May 22, 2026 at 15:21 UTC
6bb8753db1076f498c240fffdd88b1ab983b7f40
1 file changed
+3
-4
tests/v1/cudagraph/test_encoder_cudagraph.py
+3
-4
@@ -17,6 +17,7 @@ from typing import Any
17
import pytest
18
import torch
19
20
+from vllm.model_executor.models.interfaces import SupportsEncoderCudaGraph
21
from vllm.platforms import current_platform
22
from vllm.v1.worker.encoder_cudagraph import (
23
EncoderCudaGraphManager,
@@ -75,7 +76,7 @@ class _MockVllmConfig:
76
self.parallel_config = _MockParallelConfig()
77
78
78
-class _MockModel:
79
+class _MockModel(SupportsEncoderCudaGraph):
80
"""Minimal mock implementing SupportsEncoderCudaGraph for __init__."""
81
82
def __init__(self, min_budget: int = 4, max_budget: int = 128):
@@ -250,7 +251,7 @@ def _count_output_tokens(
251
return sum(t * (h // m) * (w // m) for t, h, w in grid_thw_list)
252
253
253
-class SimpleMockViTModel(torch.nn.Module):
254
+class SimpleMockViTModel(torch.nn.Module, SupportsEncoderCudaGraph):
255
"""Minimal ViT model for CUDA graph tests.
256
257
Implements the SupportsEncoderCudaGraph protocol by providing
@@ -258,8 +259,6 @@ class SimpleMockViTModel(torch.nn.Module):
259
simulates spatial merge by averaging groups of m^2 patches.
260
"""
261
261
- supports_encoder_cudagraph = True
262
-
262
def __init__(self):
263
super().__init__()
264
self.proj = torch.nn.Linear(_FLAT, _HIDDEN)