Hybrid SSM / Mamba models and non-KV state

vLLM's hybrid-model startup path contains an arithmetic loop that no pure-attention model ever runs. It computes the byte size of one Mamba state page, computes the byte size of one token of KV, and then inflates the attention block size until a block of KV is at least as large as a state:

            attn_block_size = chunk_size * cdiv(attn_tokens_per_mamba_state, chunk_size)
            cache_config.mamba_block_size = attn_block_size
        else:
            # Without prefix caching, use minimum block size that satisfies
            # both backend alignment and mamba page size compatibility
            attn_block_size = kernel_block_alignment_size * cdiv(
                mamba_page_size,
                kernel_block_alignment_size * attn_page_size_1_token,
            )

        if cache_config.block_size < attn_block_size:
            cache_config.block_size = attn_block_size
            logger.info(
                "Setting attention block size to %d tokens "
                "to ensure that attention page size is >= mamba page size.",
                attn_block_size,

The consequence is not cosmetic. §2.3 established that prefix-cache hit length is quantised to the block: a 1279-token shared system prompt in a 1280-token block hits nothing. A hybrid model does not merely make prefix caching harder in principle; it makes the granularity of prefix caching two orders of magnitude coarser, for the attention layers too, because both pools are forced onto one page size. vLLM has an open bug for exactly this failure mode — issue #45238, "Hybrid-model prefix caching silently drops to 0% when the align-mode Mamba checkpoint lands in request-unique tokens."

Two qualifiers, owed up front. The coarse granularity is a default, not a law: --prefix-match-unit decouples the hash unit from the physical block, so a hit can land inside a block — "It controls matching granularity only, not how often states are stored" (vllm/config/cache.py:L68-L79). Its default is None, and it only bites in "align" mode; unset, the hash unit is the GCD of the two groups' block sizes, which for a hybrid is the block itself (vllm/v1/core/kv_cache_utils.py:L699-L709). And prefix reuse is not gone: both engines checkpoint the recurrent state and match to the nearest checkpoint, and both do it by default — §6 establishes exactly what is on and what is not.

Everything in this chapter follows from one substitution: replace a per-token cache with a per-sequence summary and the memory curve flattens, but the addressing model dies.

An attention layer at step $t$ reads every key and value it has ever written. To do that it must keep them, so its memory grows as $O(s)$ in context length $s$. A selective state-space layer instead maintains a recurrence: a state tensor $S_t$ that absorbs token $t$ and is then overwritten. Nothing per-token survives. Memory is $O(1)$ in $s$ — a fixed slot per sequence, the same size at token 8 and token 128,000.

The catch is the constant. "Constant" is not "small": the state has to carry enough capacity to be a useful summary, so it is sized like a fat activation, not like one token of KV. The interesting question is therefore not "is it smaller" but "past what context is it smaller".

At 131,072 tokens, Llama's KV is 16 GiB. The hybrid uses 2 GiB of attention KV plus about 97.4 MiB of recurrent state, about 2.095 GiB total: roughly 7.64 times smaller, not 168 times. The 779-token marked crossover compares Llama KV with the state component only; the total-hybrid crossover is about 891 tokens. A 256-slot pool reserves about 24.4 GiB of recurrent state, before its attention cache and other memory.

vLLM computes Mamba2 state shapes in one function. Read it rather than the paper — the layout below is what actually gets allocated:

        # if n_groups is not divisible by world_size, need to extend the shards
        # to ensure all groups needed by a head is sharded along with it
        n_groups = n_groups + cls.extra_groups_for_head_shards(n_groups, tp_world_size)
        # heads and n_groups are TP-ed
        conv_dim = intermediate_size + 2 * n_groups * state_size

        conv_state_shape = cls._orient_conv_shape(
            divide(conv_dim, tp_world_size), conv_kernel - 1 + num_spec
        )

        # These are not TP-ed as they depend on A, dt_bias, D
        # - they are typically small
        #   e.g., (h_heads, head_dim, state_size) = (128, 64, 128)
        temporal_state_shape = (divide(num_heads, tp_world_size), head_dim, state_size)
        return conv_state_shape, temporal_state_shape

There are two objects, not one.

• The temporal (SSM) state, shape $(H_m, d_h^m, d_{\text{state}})$ — one matrix per SSM head. This is the summary the recurrence reads and overwrites every token.
• The short-convolution state, shape $(d_{\text{conv}}, W-1)$ with $W$ the depthwise kernel width and $d_{\text{conv}} = d_{\text{inner}} + 2 g\, d_{\text{state}}$ — a rolling window of the last $W-1$ token activations, not a summary. A tiny FIFO that keeps the convolution causal across step boundaries.

Now a real model. vLLM ships Nemotron-H's config class in-tree, so the shapes are readable without a network:

        ssm_state_size=128,  # mamba_state_size
        mamba_num_heads=128,
        mamba_n_groups=8,  # nemo: mamba_ssm_ngroups = num_heads
        mamba_head_dim=64,
        mamba_d_conv=4,
        mamba_expand=2,

With hidden_size=4096, num_hidden_layers=52, head_dim=128, num_key_value_heads=8, and the layer pattern "M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-" (all from lines L151-L179 of the same file), the model has 24 Mamba layers, 4 attention layers, and 24 MLP layers. Nemotron-H builds its mixer with intermediate_size = mamba_num_heads * mamba_head_dim (vllm/model_executor/models/nemotron_h.py:L371), so $d_{\text{inner}} = 128 \times 64 = 8192$. The arithmetic, at TP=1:

The conv state is 1.4 % of the total. The SSM state is everything, and its dtype is the largest lever available: --mamba-ssm-cache-dtype bfloat16 halves temporal-state bytes, not unchanged convolution state or hybrid KV, at a numerical risk vLLM flags in one line — "Only float32 is known to have no accuracy issues by default."

The crossover, derived. Llama-3-8B carries 128 KiB of KV per token (§2.1), so 8k context is 1 GiB per sequence. Setting $s \cdot 131072 = 102{,}137{,}856$ gives $s = 779$ tokens; with the bf16 override, $s = 395$. These compare the recurrent-state component alone. Add the hybrid's four attention layers before comparing whole-model memory: at 128k it is about 2.095 GiB versus Llama's 16 GiB, as Figure 1's component accounting explains.

§3.4 found that vLLM selects an attention backend by capability: get_attn_backend assembles an AttentionSelectorConfig of fifteen-odd predicates — head size, dtype, KV cache dtype, block size, MLA, sinks, sparsity, sliding window, DCP, PCP, batch invariance — and hands them to the platform, which rejects backends that cannot satisfy them (vllm/v1/attention/selector.py:L154-L190). That finding continues here, because SSM layers do not use any of it. They use a second enum:

class MambaAttentionBackendEnum(Enum, metaclass=_AttentionBackendEnumMeta):
    """Enumeration of all supported mamba attention backends.

    The enum value is the default class path, but this can be overridden
    at runtime using register_backend().

    To get the actual backend class (respecting overrides), use:
        backend.get_class()
    """

    MAMBA1 = "vllm.v1.attention.backends.mamba1_attn.Mamba1AttentionBackend"
    MAMBA2 = "vllm.v1.attention.backends.mamba2_attn.Mamba2AttentionBackend"
    SHORT_CONV = "vllm.v1.attention.backends.short_conv_attn.ShortConvAttentionBackend"
    LINEAR = "vllm.v1.attention.backends.linear_attn.LinearAttentionBackend"
    GDN_ATTN = "vllm.v1.attention.backends.gdn_attn.GDNAttentionBackend"

And a selector with no selection in it:

def get_mamba_attn_backend(
    mamba_type: MambaAttentionBackendEnum,
) -> type[AttentionBackend]:
    """Select which mamba attention backend to use and lazily import it."""
    return _cached_get_mamba_attn_backend(mamba_type)


@cache
def _cached_get_mamba_attn_backend(
    mamba_type: MambaAttentionBackendEnum,
) -> type[AttentionBackend]:
    assert mamba_type and isinstance(mamba_type, MambaAttentionBackendEnum)

    mamba_attn_backend = mamba_type.get_class()

The only argument is mamba_type, and mamba_type is a property the layer declares — MambaMixer2.mamba_type returns MAMBA2, ShortConv.mamba_type returns SHORT_CONV, unconditionally (vllm/model_executor/layers/mamba/mamba_mixer2.py:L1146-L1148, vllm/model_executor/layers/mamba/short_conv.py:L331-L333). Selection is by model architecture, resolved at class-definition time; nothing is validated, nothing is negotiated, and there is no --attention-backend equivalent short of registering an override. The one runtime check in the whole path is a batch-invariance guard that raises rather than falls back.

The split is defensible — every predicate in AttentionSelectorConfig is an attention concept, and an SSM layer has no head size, no KV layout, no sliding window, and exactly one kernel that can execute it. The cost is that the SSM path inherits none of the capability checking, per-kind overrides, or fallback logic of §3.4, so an unsupported combination surfaces as a kernel failure rather than a startup refusal.

The block pool of §2.2 exists to solve one problem: a sequence needs an unknown, growing number of pages. A recurrent state needs exactly one slot, known at admission — so the pool's reason for existing is absent, and both engines end up back at a flat slot array. SGLang is explicit: MambaPool allocates one dense tensor per state object, indexed [layer, slot, ...]:

                conv_state = [
                    torch.zeros(
                        size=(num_mamba_layers, size + 1) + conv_shape,
                        dtype=conv_dtype,
                        device=device,
                    )
                    for conv_shape in conv_state_shape
                ]
# ...
                temporal_state = torch.zeros(
                    size=(num_mamba_layers, size + 1) + temporal_state_shape,
                    dtype=ssm_dtype,
                    device=device,
                )

size is --max-mamba-cache-size, a request-count cap, and the + 1 is the padding slot — the same null-slot idiom the block pool uses for unallocated rows (§2.2). A request holds req.mamba_pool_idx, an integer.

vLLM keeps the state inside the KV cache manager instead, as a MambaSpec group whose page is one state — one allocator for both pools, at the price of pretending the state is paged. The pretence shows in the block-table shapes:

    Get the block table tensor for mamba kernels from the input
    common_attn_metadata.block_table_tensor given different mamba cache modes.

    - "all":   input  (#requests, cdiv(max_model_len, block_size)
                        + num_speculative_blocks);
               output (#requests, cdiv(max_model_len, block_size)
                        + num_speculative_blocks).

    - "none":  input  (#requests, 1 + num_speculative_blocks);
               output (#requests, 1 + num_speculative_blocks).

    - "align": input  (#requests, cdiv(max_model_len, block_size));
               output (#requests, 1 + num_speculative_blocks), which are the last
               1 + num_speculative_blocks of each request.

In "none" mode the block table is one column wide: vLLM has degenerated to SGLang's slot array, and mamba_block_size is set to max_model_len so that one block is one sequence (vllm/model_executor/models/config.py:L646-L647). The residency budget follows the same three-way split:

    def max_memory_usage_bytes(self, vllm_config: VllmConfig) -> int:
        if vllm_config.cache_config.mamba_cache_mode == "all":
            max_model_len = vllm_config.model_config.max_model_len
            return (
                cdiv(max_model_len, self.block_size) + self.num_speculative_blocks
            ) * self.page_size_bytes
        elif vllm_config.cache_config.mamba_cache_mode == "align":
            return self.page_size_bytes * (2 + self.num_speculative_blocks)
        else:
            return self.page_size_bytes * (1 + self.num_speculative_blocks)

One page in "none"; two in "align" (you need the previous step's state alive while you write the next); and in "all", one page per block_size tokens of the whole model length — a checkpoint chain. Both "align" and "all" serve prefix caching; "all" is simply the denser policy, and the more expensive one. Blocks are shared across every layer in the group, so a 128k-context Nemotron-H request in "all" mode budgets $\lceil 131072/1280 \rceil = 103$ state blocks at 4.06 MiB per block per Mamba layer: 9.8 GiB for one request — 12.0 GiB in practice, because the spec is built with page_size_padded and the page is inflated to the 5.00 MiB attention page derived below (vllm/model_executor/layers/mamba/abstract.py:L63-L71). Both figures are derived.

That number is why prefix_cache_retention_interval and its reachable_block_mask exist — and the default is 0, not dense: the field's factory reads the env var and returns 0 when it is unset (vllm/config/cache.py:L40-L46). At 0 the mask marks only the semantic checkpoints — the replay boundary and any detected shared-prefix junction — and every other block is allocated but never hashed into the prefix cache (vllm/v1/core/single_type_kv_cache_manager.py:L1400-L1429). Marconi-style sparse admission is what runs out of the box; --prefix-cache-retention-interval N layers periodic checkpoints on top of the semantic ones, at a cost you can now compute. Note the in-code comment beside that branch still reads "Dense caching (default)" — it describes the function's own default argument, not the configured one. §6 comes back to how both engines thin the chain.

The page-size-alignment problem

A hybrid model produces a KVCacheConfig with two groups of different specs. vLLM's coordinator requires their pages to be the same size so one block index means the same thing in both, which is why the startup loop in §1 exists. Derived, for Nemotron-H-8B at TP=1 with FlashAttention (get_supported_kernel_block_sizes returns MultipleOf(16)):

In "all" mode the block must also be a multiple of the Mamba chunk size (256 here) because the SSD kernels lay out chunks that way — a constraint the source says "can be relaxed fairly easily by changing the way we layout chunks in the mamba2 kernels", which is a fair description of an unpaid debt. The 23.2 % padding is HBM burned purely so the two groups' arithmetic agrees.

Here is the crux. A KV cache is positionally addressable: block $i$ holds tokens $[i\cdot B, (i{+}1)\cdot B)$ and means the same thing regardless of what follows, so a radix tree (§2.4) can split a node at any page boundary. A recurrent state is a summary: $S_{4000}$ is a function of all 4000 tokens and contains no sub-object equal to $S_{2000}$. You can reuse it only if you snapshotted at exactly the position you need. SGLang says this in a six-word comment:

    def _split_node(self, key: RadixKey, child: TreeNode, split_len: int) -> TreeNode:
        # new_node -> child
        new_node = TreeNode()
        new_node.children = {key[split_len:].child_key(self.page_size): child}
        new_node.parent = child.parent
        new_node.mamba_value = None  # mamba cache can not be split

SGLang: one checkpoint per node, matching truncates backwards

MambaRadixCache gives every tree node an optional mamba_value alongside its value (the KV indices), and maintains two LRU lists and two lock counters over the same tree, because a KV hit reuses the whole matched path while a Mamba hit consumes exactly one node's state (python/sglang/srt/mem_cache/mamba_radix_cache.py:L84-L96). Matching walks down as far as the tokens agree, but remembers the deepest node that had a state:

        while len(key) > 0 and child_key in node.children.keys():
            child = node.children[child_key]
            # update best_value_len and best_last_node if needed
            if node.mamba_value is not None:
                best_value_len = len(value)
                best_last_node = node
# ...
        # handle best_value_len and best_last_node, for the case that last node is fully matched
        if node.mamba_value is not None:
            best_value_len = len(value)
            best_last_node = node

        return value, best_last_node, best_value_len

_match_post_processor then truncates the KV hit to that node (value = value[:best_value_len]) — the KV cache is deliberately under-used to keep the two pools consistent — and records mamba_branching_seqlen, the last chunk-aligned position without a state, as a hint about where a checkpoint would pay off. The state itself is copied into a fresh slot rather than shared, because the recurrence updates in place (python/sglang/srt/mem_cache/mamba_radix_cache.py:L1181-L1193).

SGLang then offers something vLLM does not: storing cached checkpoints compressed. Read the default before you believe the docstring — --enable-int8-mamba-checkpoint is False (python/sglang/srt/server_args.py:L2620-L2624), and it is refused alongside --enable-hierarchical-cache or a custom --radix-cache-backend, neither of which is int8-aware (python/sglang/srt/server_args.py:L6280-L6300). Unset, the radix parks a full-precision slot on the node like everything else. What the file describes is the opt-in path:

It decouples the *cached* states (radix-owned, idle, compressed) from the *active*
``MambaPool`` (running requests, full precision, kernel-facing). The radix stores
one cached state per node HERE; on a prefix-cache hit it is dequantized back into
a fresh active slot (copy-on-write).

Per cached slot it holds:
  * the SSM temporal state in **int8** (per-(head,k-channel) symmetric), via the
    embedded ``Int8CheckpointStore`` — ~2x more cached states than bf16,
    quality-safe (quantized once on store, dequantized once on a hit; never
    re-enters the recurrence as a quant->dequant loop).
  * the conv1d window state at its native dtype (tiny, W-1 tokens; not worth

Quantizing once avoids injecting fresh quantization error every recurrent step, but the restored error still propagates. With the same subsequent inputs, $\Delta S_{t+m}=(A_{t+m}\cdots A_{t+1})\Delta S_t$ in a linear recurrence. Products may contract or amplify the initial error; output maps and nonlinear dependencies matter too. Test long continuations and state magnitudes. A source comment saying "quality-safe" is not an accuracy proof or a blanket advantage over KV quantization.

vLLM: three modes, and a search that runs right to left

vLLM exposes the policy as a flag, documented in the config:

    mamba_cache_mode: MambaCacheMode = "none"
    """The cache strategy for Mamba layers:

    - "none": set when prefix caching is disabled.
    - "all": cache the mamba state of all tokens at position i * block_size.
    - "align": only cache the mamba state of the last token of each scheduler step and
      when the token is at position i * block_size. This is the default when prefix
      caching is enabled.

And the lookup, in MambaManager.find_longest_cache_hit, is the mirror image of the attention version. Full attention scans left to right accumulating a run of hits; Mamba scans right to left and takes the first one it finds:

        max_num_blocks = max_length // block_size
        # Search from right to left and early stop when a match is found.
        for i in range(max_num_blocks - 1, -1, -1):
            if cached_block := block_pool.get_cached_block(
                block_hashes[i], kv_cache_group_ids
            ):
# ...
                for computed, cached in zip(computed_blocks, cached_block):
                    # the hit length logic later assumes:
                    #  hit_length = len(hit_blocks_other_attn[0])
                    #               * self.other_block_size
                    # so we insert dummy blocks at the beginning:
                    computed.extend([block_pool.null_block] * i)
                    computed.append(cached)
                hit_length = (i + 1) * block_size
                break  # we just need the last match - early stopping

The null block from §2.2 earns its keep here: the returned list has to be $i{+}1$ entries long so the generic hit-length arithmetic works, but only the last entry is real. A "hit" of 3840 tokens is one 4.06-MiB-per-layer state and 2 null blocks.

So the honest answer to "does prefix caching survive?": yes, but only at snapshot points, and the snapshots are far apart. Both engines converge on the same design — checkpoint the state at chosen positions, match to the nearest checkpoint at or before the divergence, replay the remainder. For a shared system prompt this works well when the prompt is long relative to the block: a 4000-token system prompt with 1280-token blocks yields a state at 1280, 2560, and 3840, so a second request replays at most 160 tokens instead of 4000. It works badly when the prompt is short (a 900-token system prompt gets no checkpoint at all and every request replays it in full) and it works badly when the shared prefix ends at an unaligned position, which is issue #45238.

Is it on by default? In both engines, yes

vLLM: enable_prefix_caching is True in the dataclass (vllm/config/cache.py:L107), and for any Mamba or hybrid architecture MambaModelConfig.verify_and_update_config promotes mamba_cache_mode from "none" to "align" and logs that it did (vllm/model_executor/models/config.py:L612-L620). Checkpoint-based Mamba prefix caching therefore runs unasked. "all" is not the default and is additionally gated on the SupportsMambaPrefixCaching protocol, which seven model files declare at this SHA (mamba.py, mamba2.py, jamba.py, zamba2.py, falcon_h1.py, granitemoehybrid.py, nemotron_h.py); anything else asking for "all" is warned back down to "align".

SGLang: for a registered hybrid architecture the resolution pass declares uses_mamba_radix_cache = True unless --disable-radix-cache is set (python/sglang/srt/arg_groups/overrides.py:L1567-L1580), and NemotronHForCausalLM is one of the seventeen names in that frozen set (python/sglang/srt/arg_groups/overrides.py:L1495-L1515), alongside GraniteMoeHybridForCausalLM when its config actually declares Mamba layers. The pass is invoked unconditionally and self-guards on that set (python/sglang/srt/server_args.py:L5792-L5798), so here too it is on without being asked for — but not free of side effects. --mamba-radix-cache-strategy defaults to auto, which selects extra_buffer only for an architecture in _MAMBA_EXTRA_BUFFER_ARCHS running the Triton linear-attention backend. Nemotron-H qualifies on both counts (python/sglang/srt/arg_groups/overrides.py:L1520-L1546). Anything outside that list drops to no_buffer, which force-sets disable_overlap_schedule (python/sglang/srt/arg_groups/overrides.py:L1580-L1591) — on those models, prefix caching for the recurrent state costs you overlap scheduling.

What is not symmetric is the thinning policy, and both defaults are the opposite of what the flag names suggest. vLLM's prefix_cache_retention_interval defaults to 0 — sparse, semantic-only admission, the policy reachable_block_mask labels "Marconi-style APC" in its own docstring. SGLang's --mamba-max-states-per-path defaults to -1, unlimited (python/sglang/srt/server_args.py:L2571-L2578): the path cap exists but is off, and what runs by default is Marconi's admission half only — mamba_branching_seqlen, checkpoint where sharing was observed — with the mamba LRU doing the retention.

The checkpoint grids differ too, and this is the sharpest consequence of the allocation split in §5. vLLM materialises at block_size, which the page-alignment loop of §1 has already inflated to 1040 tokens (align) or 1280 (all). SGLang keeps state in a separate slot array, never couples the two page sizes, and computes its grid as lcm(max(model mamba chunk size, page_size), tree_page) (python/sglang/srt/runtime_context.py:L1479-L1490, python/sglang/srt/server_args.py:L9177-L9199). For Nemotron-H that is mamba_chunk_size = 256 at any page size that divides it — four to five times finer than vLLM's, and finer without costing the attention layers anything. vLLM's answer to the same problem is --prefix-match-unit, which buys back matching granularity but, as its own docstring says, "controls matching granularity only, not how often states are stored". Setting it buys exactly one extra materialisation point — the scheduler adds a mandatory stop at the prompt's own last hash boundary so a partial-tail entry can be registered (vllm/v1/core/sched/scheduler.py:L324-L334). Everything else still lands on the 1040-token grid.

Preemption changes character. For a KV cache it is annoying but simple: free the blocks, reset num_computed_tokens to zero, and on resume the prefix cache hands most of them back. vLLM does exactly that, with nothing Mamba-specific (vllm/v1/core/sched/scheduler.py:L1347-L1388). The asymmetry is that the attention group's work is re-derivable from a hit, while the state group can reuse only a checkpoint — and in "none" mode there are none, so a preempted 100k-token sequence replays 100k tokens of recurrence.

SGLang's reset_for_retract makes the loss explicit — nine Mamba fields, all cleared:

        self.mamba_pool_idx = None
        self.mamba_ping_pong_track_buffer = None
        self.mamba_next_track_idx = None
        self.mamba_last_track_idx = None
        self.mamba_last_track_seqlen = None
        self.mamba_branching_seqlen = None
        self.mamba_cow_src_index = None
        self.mamba_needs_clear = False
        self.already_computed = 0

It also offers the alternative — keep it, on the host. Req.offload_kv_cache copies both pools down: token_to_kv_pool_allocator.get_cpu_copy(token_indices, mamba_indices=self.mamba_pool_idx) (python/sglang/srt/managers/schedule_batch.py:L1720-L1731). At this SHA that path is gated on disaggregation_mode == "decode". Ordinary preemption may discard active state, but a surviving compatible prefix checkpoint provides a middle option: restore it and replay only the missing suffix. Full replay is necessary only when no usable checkpoint remains.

Chunked prefill is a state handoff. The recurrence must be applied in order, so a chunk boundary is a place where $S$ leaves one kernel launch and enters the next. Mamba2's metadata builder carries has_initial_states_p and prep_initial_states for exactly this: when num_computed_tokens > 0 the chunk seeds the scan from the stored state instead of zero (vllm/v1/attention/backends/mamba_attn.py:L568-L570, vllm/v1/attention/backends/mamba2_attn.py:L145-L156). In "align" mode the boundary is further constrained, because the slot a chunk writes into must correspond to a cacheable position:

        """Clip a prefill chunk so it ends where Mamba state must be cached.

        In "align" cache mode reusable SSM states are materialized at block
        boundaries, plus mandatory early stops (the prompt's partial-tail hash
        boundary, a detected shared-prefix junction). If a block is larger
        than the configured prefill chunk limit, intermediate chunks keep
        private running state until they reach the next cacheable position.
        """

The scheduler is now clipping token budgets to satisfy a kernel's state invariant — a genuinely new coupling. In a pure-attention engine the scheduler picks chunk sizes for latency alone and the kernel does not care where they land.

For stored state $S_t\in\mathbb{R}^{d_h\times d_{\mathrm{state}}}$, a simplified recurrence is $S_t=A_t\odot S_{t-1}+x_tB_t^\top$ and $y_t=S_tC_t$. Mamba's full block also has discretization, convolution, skip and gating operations; this equation isolates the scan kernel's state geometry.

Prefill processes many tokens at once, so the sequential dependency is the enemy. Mamba2's structured-state-space-duality form splits the sequence into chunks, computes each chunk's intra-chunk contribution as a masked matmul on tensor cores, and stitches chunks with a small inter-chunk scan:

            intermediate_states, varlen_state = mamba_chunk_scan_combined(
                hidden_states_p.view(
                    1, num_prefill_tokens, local_num_heads, self.head_dim
                ),
                dt_p.unsqueeze(0),
                self.A,

Decode processes one token per sequence, so there is nothing to parallelise along time and the chunked machinery is pure overhead. It becomes a plain elementwise recurrent update, gathered across the batch by slot index:

                selective_state_update(
                    ssm_state,
                    hidden_states_d,
                    dt_d,
                    A_d,
                    B_d,
                    C_d,
                    D_d,
                    z=None,
                    dt_bias=dt_bias,
                    dt_softplus=True,
                    state_batch_indices=state_indices_tensor_d,
                    out=preallocated_ssm_out_d.view(num_decodes, -1, self.head_dim),
                )

The state update reads and writes each fp32 state element. Counting only one multiply-add gives roughly 2 FLOPs / 8 bytes = 0.25 FLOP/byte, before input injection, output reduction and other traffic. Use the scalar instruction roofline, not tensor-core peak, for scalar updates. Low intensity suggests memory pressure, but achieved bandwidth and small-batch occupancy determine runtime; attention layers still add context-dependent traffic in a hybrid.

Dispatch is a wrapper. SGLang's HybridLinearAttnBackend holds two backends and a set of layer indices, and routes per layer:

    def forward_decode(
        self,
        layer: RadixAttention,
        forward_batch: ForwardBatch,
        save_kv_cache: bool = True,
        q: Optional[torch.Tensor] = None,  # For full attention
        k: Optional[torch.Tensor] = None,  # For full attention
        v: Optional[torch.Tensor] = None,  # For full attention
        mixed_qkv: Optional[torch.Tensor] = None,  # For linear attention
        a: Optional[torch.Tensor] = None,  # For GDN linear attention
        b: Optional[torch.Tensor] = None,  # For GDN linear attention
        **kwargs,
    ):
        if self._is_full_attn(layer, kwargs.get("layer_id")):
            return self.full_attn_backend.forward_decode(
                q, k, v, layer, forward_batch, save_kv_cache, **kwargs
            )
        return self.linear_attn_backend.forward_decode(

The union-typed signature — q/k/v for one branch, mixed_qkv/a/b for the other — is the honest shape of a hybrid: two unrelated operators behind one interface. vLLM's Mamba2AttentionBackend reaches the same place through the layer's get_attn_backend() rather than a wrapper.

Nemotron-H-8B on vLLM, prefix caching on, --mamba-cache-mode all, so block_size = 1280. Assume request A's 4000-token prompt was processed with checkpoints actually saved and retained at 1280, 2560 and 3840. Merely setting all mode and finishing a prompt is not evidence every historical boundary is materialized. Request B adds 200 private tokens to the same prompt.

1. Scheduler.schedule() calls kv_cache_manager.get_computed_blocks(request), which fans out to each group's find_longest_cache_hit.
2. The full-attention group scans left to right and hits blocks 0, 1, 2 — 3840 tokens.
3. The Mamba group enters MambaManager.find_longest_cache_hit (single_type_kv_cache_manager.py:L1295-L1371), walks $i = 2, 1, 0$, and stops at $i = 2$ — the state saved after token 3839. It returns [null, null, state_block] with hit_length = 3840.
4. The coordinator narrows group by group — curr_hit_length = min(curr_hit_length, hit_length_by_group[...]), then passes that as the next finder's max_length (vllm/v1/core/kv_cache_coordinator.py:L800-L845). Here both groups say 3840, because the block sizes were forced equal at startup; without that forcing the smaller hit would win.
5. _mamba_block_aligned_split is skipped (mode is all, not align), so the remaining 360 tokens are scheduled as an ordinary chunk.
6. Mamba2AttentionMetadataBuilder.build sees num_computed_tokens = 3840 > 0, sets has_initial_states_p true and prep_initial_states true (mamba2_attn.py:L145-L156), and calls compute_varlen_chunk_metadata to split the 360 tokens at 256-token physical chunk boundaries.
7. In the forward pass, the SSM layers seed initial_states from the hit state block and run mamba_chunk_scan_combined over 360 tokens; the four attention layers run FlashAttention over the 3840 cached tokens plus 360 new ones.
8. At the block boundary the new state is written back to its own block and, at cache_blocks time, hashed into the pool so request C can find it.

B computes 360 tokens instead of its full 4200-token prompt, reusing 3840. With a 900-token shared prefix plus 200 private tokens and no checkpoint, it computes all 1100. Checkpoint existence, retention and exact prefix identity determine reuse.

Affine transitions compose associatively: applying (a1,b1) then (a2,b2) gives (a2*a1, a2*b1+b2). That enables scans without pretending the recurrence has no dependency. A checkpoint perturbation evolves through the product of later transition factors; one-time quantization can therefore contract or amplify. In a real selective block, output maps, convolution state and input-dependent coefficients must also match.

import numpy as np

def compose(first, second):
    a1, b1 = first
    a2, b2 = second
    return a2*a1, a2*b1+b2

steps = [(0.9, 1.0), (0.8, -0.5), (0.7, 2.0)]
left = compose(compose(steps[0], steps[1]), steps[2])
right = compose(steps[0], compose(steps[1], steps[2]))
np.testing.assert_allclose(left, right)
for factor in (0.9, 1.1):
    exact, restored = 2.0, 2.01
    for _ in range(20):
        exact = factor*exact + 0.1
        restored = factor*restored + 0.1
    assert np.isclose(restored-exact, 0.01*factor**20)
hybrid_mib = 2048 + 97.4
assert 7.6 < 16384/hybrid_mib < 7.7
print("Affine scan composition, propagated error and full hybrid memory pass.")

Start a hybrid model twice and diff the startup logs. No GPU is needed to read the first two of these lines; they are emitted during config verification.

vllm serve nvidia/Nemotron-H-8B-Base-8K --enable-prefix-caching --mamba-cache-mode all 2>&1 | grep -E "attention block size|Padding mamba page|Mamba cache mode"

vllm serve nvidia/Nemotron-H-8B-Base-8K --no-enable-prefix-caching 2>&1 | grep -E "attention block size|Padding mamba page|Mamba cache mode"

python3 -m sglang.launch_server --model-path nvidia/Nemotron-H-8B-Base-8K 2>&1 | grep "Mamba Cache is allocated"

Three things to measure. (1) The reported block size against the derivation in §5 — if it does not match, one of the shapes changed and the table is stale. (2) SGLang prints Mamba Cache is allocated. max_mamba_cache_size: N, conv_state size: X.XXGB, ssm_state size: Y.YYGB (python/sglang/srt/mem_cache/memory_pool.py:L806-L831); divide ssm_state size by max_mamba_cache_size and check it equals the 96 MiB derived in §3. (3) Send the same system prompt at 900, 1300, and 4000 tokens and watch TTFT for the second request — the step change at the block boundary is the whole chapter in one graph.

• Gu & Dao, Mamba (arXiv:2312.00752), and Dao & Gu, Transformers are SSMs (arXiv:2405.21060) — the latter is where the chunked SSD algorithm behind mamba_chunk_scan_combined comes from.
• Waleffe et al., An Empirical Study of Mamba-based Language Models (arXiv:2406.07887). Cited result: at 8B parameters on identical data, pure SSMs match or exceed Transformers on many tasks but lag where copying or in-context learning is needed (5-shot MMLU, Phonebook), while a 43 % Mamba-2 / 7 % attention / 50 % MLP hybrid exceeded the Transformer on all 12 standard tasks evaluated — essentially Nemotron-H's mixing ratio.
• Jelassi et al., Repeat After Me: Transformers are Better than State Space Models at Copying (arXiv:2402.01032) — the theoretical statement of the same limit: a fixed-size latent bounds how much of the context can be reproduced verbatim.
• Pan et al., Marconi: Prefix Caching for the Era of Hybrid LLMs (MLSys 2025, arXiv:2411.19379) — diagnoses the exact constraint this chapter traces, and proposes the admission policy vLLM's reachable_block_mask cites by name.
• vLLM issue #26201 (tracking: prefix caching for hybrid models), PR #45939 (partial prefix cache primitives), PR #46384 (partial hits for hybrid models), issue #45238 (the silent-zero-hit-rate bug).
• NVIDIA, Nemotron-H (arXiv:2504.03624) — the model whose config file supplied every shape in §3.