Ollama: Supported Architectures For Flash Attention & KV Cache

Understanding the nuances of memory optimization is crucial when deploying large language models (LLMs) using Ollama, especially on resource-constrained GPUs like the Nvidia L4. Two key features for memory management are Flash Attention and KV Cache Quantization, enabled via the OLLAMA_FLASH_ATTENTION and OLLAMA_KV_CACHE_TYPE environment variables, respectively. However, the relationship between these features and specific model architectures requires clarification.

This article delves into the supported architectures for Flash Attention and KV Cache Quantization within Ollama. We'll explore the importance of explicitly documenting these dependencies to prevent unexpected memory issues and streamline the deployment process.

The Importance of Explicit Architecture Support

When optimizing memory usage for LLM deployments, developers often turn to techniques like Flash Attention and KV Cache Quantization. Flash Attention reduces memory footprint and accelerates computation by optimizing the attention mechanism. KV Cache Quantization further compresses the memory required to store key-value pairs during inference.

The OLLAMA_FLASH_ATTENTION and OLLAMA_KV_CACHE_TYPE environment variables in Ollama provide a convenient way to enable these features. However, not all model architectures are created equal. Some architectures are specifically designed to work with Flash Attention, while others may not be compatible. Attempting to force these optimizations on unsupported architectures can lead to unexpected behavior, such as the server silently falling back to less efficient memory formats like f16. This can result in out-of-memory (OOM) errors or higher VRAM usage than anticipated, negating the intended benefits of memory optimization. Ensuring that developers have a clear understanding of which architectures support these features is critical for successful deployment.

To address this, explicit documentation is essential. A simple table or note in the docs/faq.md or docs/gpu.md files, listing the architectures that currently support Flash Attention (and thus KV Quantization), can save developers significant time and effort in debugging memory usage and planning capacity for specific models. This clarity will enable more efficient resource allocation and a smoother deployment experience.

Current Architecture Support in Ollama

Currently, Ollama's FlashAttention() function in fs/ggml/ggml.go relies on a specific allowlist to determine which architectures support Flash Attention. This allowlist includes:

• gemma3
• gptoss, gpt-oss
• mistral3
• qwen3, qwen3moe
• qwen3vl, qwen3vlmoe

This means that only these architectures will effectively utilize Flash Attention when OLLAMA_FLASH_ATTENTION is enabled. Other architectures, such as command-r or llama3 (standard), will not benefit from this optimization, even if OLLAMA_KV_CACHE_TYPE=q8_0 is set. This discrepancy can lead to confusion and unexpected resource consumption, highlighting the need for clearer documentation.

The Issue: Silent Fallbacks and Unexpected Memory Usage

The core issue lies in the silent fallback mechanism. When a developer attempts to enable OLLAMA_KV_CACHE_TYPE=q8_0 on an architecture not included in the FlashAttention() allowlist, the server doesn't throw an error or warning. Instead, it silently reverts to f16, a less memory-efficient format. This silent fallback can be misleading, as developers might assume that KV Cache Quantization is active when it isn't.

This can lead to several problems:

• Unexpected OOMs: The higher memory footprint of f16 can cause out-of-memory errors, especially on GPUs with limited VRAM.
• Higher VRAM Usage: Even if OOMs don't occur, the increased VRAM usage can impact performance and limit the number of concurrent users the system can support.
• Debugging Challenges: The silent fallback makes it difficult to diagnose memory-related issues, as developers might not realize that KV Cache Quantization is not actually enabled.

To mitigate these issues, it's crucial to provide clear and accessible information about supported architectures. This will empower developers to make informed decisions about model selection and configuration, leading to more predictable and efficient deployments.

Proposed Solution: Documenting Supported Architectures

To address the lack of clarity regarding supported architectures for Flash Attention and KV Cache Quantization, I propose adding a dedicated section to the Ollama documentation. This section could take the form of a table or a detailed note in either docs/faq.md or docs/gpu.md. The documentation should clearly list the architectures that are compatible with Flash Attention and KV Cache Quantization.

Here's an example of how the table could be structured:

In addition to the table, the documentation should also explain the implications of using Flash Attention and KV Cache Quantization on unsupported architectures. This will help developers understand why certain configurations might not work as expected and guide them towards more appropriate solutions.

Future-Proofing the Documentation

To ensure the documentation remains accurate and up-to-date, it's essential to establish a process for maintaining it. Ideally, whenever a new model architecture is added to the FlashAttention function in fs/ggml/ggml.go, the documentation should be updated accordingly. This could be achieved through a combination of automated checks and manual review.

Furthermore, I'm interested in contributing to this effort by maintaining the table in the docs/faq.mdx file whenever a new model gets added to the FlashAttention function. By actively participating in the documentation process, I hope to help ensure that Ollama users have access to the information they need to deploy models effectively.

Conclusion

Clarifying the supported architectures for Flash Attention and KV Cache Quantization is crucial for optimizing memory usage and preventing unexpected issues in Ollama deployments. By providing clear and accessible documentation, we can empower developers to make informed decisions, leading to more efficient resource allocation and a smoother deployment experience. This proactive approach will not only save developers time and effort but also enhance the overall usability and reliability of Ollama.

For more in-depth information on GPU memory management and optimization techniques, consider exploring resources like the NVIDIA Developer Blog for insights and best practices.