Motivation.

Problem Statement

Supporting hybrid models that combine FullAttention (FA) and Mamba-style
SSM layers introduces several challenges for the KV connector:

• FA and Mamba layers use different internal state layouts
  (K/V vs Conv/SSM).
• Kernel constraints may require physical block sizes that differ
  from the logical block abstraction used by the block manager.

As a result, FA and Mamba layers must be able to index the same underlying KV cache tensor while using different block descriptor layouts.

Proposed Design

We introduce two logical descriptor views over the same registered memory regions:

• Current descriptor view (used by non-Mamba layers)

  • Descriptors correspond to K/V blocks

• Mamba descriptor view (used by Mamba layers)

  • Descriptors correspond to Conv and SSM state blocks

Both descriptor sets reference the same underlying tensor but use different offsets and sizes.

The descriptor lists are stored continuously, allowing the existing block_id → desc_id mapping logic to be extended with a simple index shift.

All changes proposed here are not meant to modify the existing workflow for "regular" models, but rather extend it for this specific case.

PR here https://github.com/vllm-project/vllm/pull/36687.

Proposed Change.

Prerequisite

This PR builds on the recently introduced HMA interface for NIXL https://github.com/vllm-project/vllm/pull/35758.

HMA

Hybrid Memory allocator enables efficient kv cache management for hybrid models, as in any type of model that combines multiple attention types (FA/SW/Linear..). HMA groups layers with the same attention type together. This leverages the fact that many models follow a repeating pattern such as 1:n ratios (e.g., three sliding-window layers followed by one full-attention layer)

Some invariants

• all groups have the same size
  • padding layers are added to match
• KVCache tensors can be viewed as 2d num_blocks x page_size, where page_size=block_size*num_heads*head_dim
• HMA ensures that all groups have equal page_size (in bytes)
  • one way to ensure that when dims do not match is to "bump" the block_size dims to a common multiple
  • can only do that if max_page_size % layer_page_size

HMA uses memory pooling

HMA implements memory pooling across groups:

• A single kv cache tensor is shared across groups!
  • all layers at position layer0 in each group share tensor0 This has several implications:
• only tensors referenced by group0 need to be registered in NIXL
• other tensors contain replicated addresses and can be skipped

Therefore we move from registering num_layers → num_regions

Consequences:

• fewer regions are registered
• each region typically contains more blocks

Importantly:

• Block IDs across groups never overlap, since they refer to different offsets in the shared tensor.
  • eg block_ids=[[1, 2], [3]]->blocks [1, 2] need to be read across all regions, same for [3].

from kv_cache_utils (group_size = max(len(group.layer_names) for group in kv_cache_groups)

# General Case:
# We will have group_size memory pools, each is shared by one layer from
# each group. As layers of different groups have different block table,
# they will use different parts of the shared Tensor.

# The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2),
# (sw.1, padding) will be: (group_size = 2)

# full.0, sw.0, sw.1: share a Tensor with size=available_memory//2
# full.1, sw.2: share another Tensor with size=available_memory//2

SSM

block_size is inherently one, there's no concept of tokens as the state is collapsed to a fixed-sized representation. This makes it tricky to operate with FA that reason in terms of blocks. The way this is done is in practice is to assume block_size for SSM layers to be 1. This results in a page_size which is much bigger than what we usually have for FA layers.

To address that, block_size for FA layers is scaled until its size is bigger than that of Conv+temporal state. SSM is then padded accordingly to satisfy the constant page_size invariant.

Example shape

# SSM padding not visible through tensor view

torch.Size([18525, 3, 3328])      # conv_state
torch.Size([18525, 48, 64, 128])  # ssm_state

torch.Size([18525, 400, 4, 128])  # FA KV, note the block_size 400 dim

The diagram below summarizes the layout (padding in SSM is not shown here):

![[Pasted image 20260310181323.png]]

Note that unlike K and V, SSM and Conv states are of different sizes (and may include additional padding).

Also note that the FI layout is assumed to be used for all FullAttention layers, as that matches with how conv|ssm are layed out and ensures no data corruption when overwriting blocks (Mamba blocks can be re-used by FA at any time).

Nixl work

The NixlConnector workflow in terms of handling blocks can be summarized to the following:

• a. register memory regions
• b. register local xfer handles and block descriptors (our unit of transfer is a block not a region)
• c. register remote xfer handles and block descriptors (during handshake)
• d. map block ids to descriptor ids (_get_block_descs_ids)

Registering two descriptors "views"

Going through the workflow, registering regions (a.) can be done similarly to what we already do for dense models, registering the whole shared memory pooled KVCache tensor once. Mind that the actual tensors passed to register_kv_caches is a tuple (ssm, conv), each indexing the same shared tensor at offseted positions. To avoid inefficient registration of overlapping areas, we use the base address pointer provided by ssm

(ssm conv)
  |   |          Tensor
  +-------->+-------------+
      |     |             |
      +----------...      |

Note that this registration process is very much akin to what we do for FI-like layouts, where we register the whole tensor and then split descriptors on K/V dim. The main difference is that K/V have identical sizes, while conv/ssm do not.

Also, more importantly, both FA and SSM layers will need to be able to index into these regions with descriptors. To address that (moving to step b. and specular c.), we propose creating two separate sets of block descriptors for the same tensor: one for FullAttn layers, the other to be used by LinearAttn layers, like in the diagram below:

Assuming: M regions, N num_blocks

+------------------------------------------------------+<----nixl_xfer_handle
|  # FA descriptors                                    |
|                                                      |
|  Region 0                                            |
|    FA_block_desc_K0                                  |
|    ...                                               |
|    FA_block_desc_Kn                                  |
|    FA_block_desc_V0                                  |
|    ...                                               |
|    FA_block_desc_Vn                                  |
|                                                      |
|  ...                                                 |
|                                                      |
|  Region M                                            |
|    FA_block_desc_K0                                  |
|    ...                                               |
|    FA_block_desc_Kn                                  |
|    FA_block_desc_V0                                  | ^
|    ...                                               | | 
|    FA_block_desc_Vn                                  | | num_descs offset
|                                                      | |
|  --------------------------------------------------  | v
|  # Mamba descriptors                                 |
|                                                      |
|  Region 0                                            |
|    Mamba_block_desc_SSM0                             |
|    ...                                               |
|    Mamba_block_desc_SSMn                             |
|    Mamba_block_desc_Conv0                            |
|    ...                                               |
|    Mamba_block_desc_Convn                            |
|                                                      |
|  ...                                                 |
|                                                      |
|  Region M                                            |
|    Mamba_block_desc_SSM0                             |
|    ...                                               |
|    Mamba_block_desc_SSMn                             |
|    Mamba_block_desc_Conv0                            |
|    ...                                               |
|    Mamba_block_desc_Convn                            |
|                                                      |
+------------------------------------------------------+

The rationale is that block_ids coming from manager can be freely allocated to either a FA or Mamba layer. Given the fact that we register K/V block descs separately, and that Conv/SSM have non-matching sizes independent from both K and V, having separate views allows to select the right block descriptor (add, len) based on whether the corresponding block_id belongs to a mamba group or not. The above can be implemented in practice by simply shifting the current block_id->desc_index mapping logic by num_descs positions.

SSM + kernel_block size

Similar problem has been also described here https://github.com/vllm-project/vllm/pull/28677.

Going back to HMA assumptions, we assume every kv cache tensor in storage to be able to be laid out as num_blocks * page_size . Due to kernel block_size requirements (such as https://github.com/vllm-project/vllm/blob/545d18d81bf11761e51c2b11a006573c2ae366c1/vllm/v1/attention/backends/flashinfer.py#L304) of backends, resulting physical tensor may actually be represented with a different number of blocks wrt what the logical manager sees. This creates two different num_blocks: a physical and a logical. Practically, in order to maintain kernel_block_size requirements, num_blocks may be scaled by a ratio, that is ratio=logical_block_size/physical_block_size, like so:

- FA tensor: `[num_logical * ratio, 2, kernel_block_size, heads, head_dim]` -- e.g. `[780390, 2, 16, 2, 128]`
    
- SSM tensor: `[num_logical, state_dim, hidden]` -- e.g. `[2990, 3, 6144]`

Where effectively one logical block includes ratio physical contiguous ones.

Therefore, to address this we allow for 2 separate num_blocks in code depending on whether we're dealing with a SSM region. This affects both block descs registration (b/c) leading to the introduction of N1 and N2 in the diagram above, but more importantly it also affects step d., as the mapping now has to take into account that different regions may use different num_blocks value.

Why do we need things in physical plane to begin with? That is mainly because for heterogeneous setup we need to index into the num_head (H) dimension of the kv cache layout directly. Using logical sizes only would prevent this, as one logical block (our unit) is encompassing ratio physical_blocks, so inner splitting would be overcomplicated.

Feedback Period.

No response

CC List.

@tlrmchlsmth @roikoren755 @robertgshaw2-redhat @ZhanqiuHu @orozery @tdoublep

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.