nano-vllm-prefillonly: An Inference Optimization Solution for Multimodal Large Models in Discriminative Tasks

A prefill-only optimization framework based on nano-vllm, which achieves up to 10x memory savings and 2x inference speedup by eliminating KV cache overhead, designed specifically for industrial-grade multimodal discriminative tasks.

In real-world industrial scenarios, many large language model applications fall into discriminative tasks, which only require the model to output a single token to complete the judgment. Typical application scenarios include:

• Reranking: Determine the relevance between documents and queries
• Retrieval/Embedding: Generate vector representations for semantic search
• Classification Tasks: Binary or multi-class classification
• Visual Question Answering: Answer yes/no questions about images
• Spatial Reasoning: Compare object sizes, positions, or relationships
• Attribute Recognition: Identify visual attributes like color and shape

With the rise of multimodal large models, these tasks are gradually shifting from traditional visual models to multimodal LLMs. For example: "Is there a dog in this picture?" "Which sign is the most eye-catching?" "Which picture best represents traditional Chinese architecture?"

However, when dealing with hundreds of millions of images, traditional vLLM solutions face severe challenges—KV cache becomes a performance bottleneck.

Traditional vLLM allocates KV cache for all models, even for embedding and reranking models that don't need it at all. On an H20 GPU with 96GB of memory, the KV cache alone can occupy 82-85GB of memory, which means:

• Each GPU can only serve one embedding/reranking model
• A large amount of memory is wasted on unused cache
• Extremely low memory efficiency in high-throughput scenarios

This design is necessary for generative tasks, but it's completely over-engineered for discriminative tasks.

This project implements nano-vllm based on the lightweight vLLM, with specialized optimizations for the prefill phase:

For single-token generation tasks, KV cache is unnecessary. This project directly skips cache allocation and uses only model weights for inference.

For multimodal models, it directly processes pixel_values without using vision cache, further reducing memory overhead.

Eliminate cache management overhead and focus on the core computation of model forward propagation.

Metric	Transformers	Prefill-Only	Original nano-vllm
Average Inference Time	1.211s	0.577s	0.571s
Peak Memory Usage	4459MB	4892MB	49680MB
Memory Savings vs Original Solution	-	10.15x	-

Key Finding: The Prefill-Only mode uses only 10% of the memory of the original solution and is only 1% slower, making it an ideal choice for memory-sensitive scenarios.