[RFC]: Support Layer-Specific KV-Cache Dtype & Layout in V1 Engine for Hybrid Models

[zhiyuan1i]

Motivation.

Problem Statement

Currently, V1’s ConstantCacheManager enforces:

Uniform KV-cache dtype across all layers (e.g., cannot mix BF16 for Linear attention’s shortconv + FP32 for recurrent-cache).
FlashInfer-only layout ([num_blocks, 2, ...]), which breaks MLA (block_size=64) unless FLASHINFER backend is used.
This prevents hybrid models with layer-specific requirements from running efficiently.

Motivation

We are developing an internal HybridModel that combines a custom Linear-attention variant (with a short convolution path) with Multi-head Latent Attention (MLA). And we successfully integrate our model to v0 engine.
After merging Minimax’s V1-engine updates from last week, we discovered that the current cache manager is too rigid for this topology.

Specifically:

We need BF16 KV-cache for the Linear-attention layers (driven by a short-convolution kernel), but FP32 KV-cache for the recurrent / MLA layers. The V1 engine presently allocates a single, uniform dtype for the entire model, forcing us to over-allocate memory when we default everything to FP32.
The engine also hard-codes the FlashInfer layout [num_blocks, 2, …] and the same block size for every layer. As a result, MLA (which we run with block_size=64) is incompatible unless we pin the backend to FLASHINFER. Skipping the safety checks produces corrupted outputs.
These constraints make it impossible to run the model without either (a) redundant memory usage or (b) a backend lock-in that we would prefer to avoid.

Proposed Change.

We’d like to open a design discussion on how V1 could allow layer-aware KV-cache policies. At a high level:

Decouple the Linear-attention and standard-attention cache pools so that each can specify its own dtype and layout.
Lift the single-block-size requirement so that different attention mechanisms can select page sizes that match their kernels (e.g., 64 for MLA vs. 1008 for Linear).
Avoid mandating FLASHINFER for all layers when only a subset actually needs it. For example, we use cutlass kernel and triton kernel.
We are intentionally not proposing a concrete patch at this stage; instead, we’d like the community’s input on:

Whether a per-layer metadata mechanism (or another approach) is the right direction.
How to maintain backward compatibility while introducing this flexibility.
Any hidden assumptions in the current block manager that would complicate heterogeneous allocations.
Feedback and alternative ideas are very welcome.

Feedback Period.

No response

CC List.

No response

Any Other Things.

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.

[heheda12345]

In short, your need is MLA + Linear attn? Can your need be resolved if we can support (num_blocks, 2, ***) layout for MLA (it is WIP) and pad the linear attn layer's state size to a multiple of MLA page size?

[heheda12345]

Also CC @tdoublep

[zhiyuan1i]

@heheda12345 Thank you for your prompt reply.

Your solution make sense. However, two fundamental issues remain:

1. Incompatible cache dtypes

Padding + unified layout would allow MLA and Linear attention to share the same physical pool, but we still lose the ability to store the short-conv cache in BF16 while leaving the recurrent cache(used by linear attention kernel) in FP32. This forces the entire pool to be allocated at the higher precision, eliminating the memory savings we’re targeting.

2. Backend flexibility vs. padding

Padding the Linear layer’s state to match MLA’s page size does not address backend heterogeneity. We currently rely on FlashMLA (block_size=64) and our in-house Triton kernels (block_size=64/256). Padded allocations tied to FlashInfer would degrade performance and lock us into a single backend, sacrificing both speed and flexibility.

Proposed Resolution

Decouple the cache pools into three independent, backend-optimized groups:

1. Short-conv pool – BF16 (Triton kernel, shared by either Linear-attention or MLA layers that use the short-conv path).
2. Linear-attention KV pool – FP32, page_size chosen by the Linear kernel (no alignment restrictions).
3. MLA KV pool – BF16 or remains unchange, block_size = 64 (FlashMLA, or remain unchanged if another MLA backend is selected).

This cleanly separates dtype and layout requirements, eliminates padding, and lets each kernel run in its optimal configuration without global constraints.

[heheda12345]

For your internal use case, you can try to make 2 BF16 layers share the memory pool with 1 FP32 layer. Then it fall back to the case with full attention + mamba. vLLM doesn't have strict constraint on data type. The assertions related to data type should be easy to fix. The only constraint is all page sizes should be same.

[zhiyuan1i]

Thanks for your reply.

Do you mean splitting one layer into multiple layers just to bypass the dtype implementation? Doesn’t that feel a bit inelegant?

Also, with the progress of hybrid models—including intra-layer mixing—should we revisit the restriction that “the only constraint is that all page sizes must be the same”?

[heheda12345]

For dtype, I think the current implementation doesn't have dtype constraint on layers (maybe there are somr assertion but the design doesn't enforce it). Layer split won't be too complex.