vllm.v1.attention.backends.utils

Classes:

• PerLayerParameters –

  Currently, FlashInfer backend only support models in which all layers share

Functions:

• compute_mm_prefix_range_tensor –

  Convert mm_prefix_range dict to padded tensor for Triton kernel.

• compute_mm_prefix_ranges –

  Per-request PrefixLM bidirectional ranges for image/video tokens.

• fill_mm_prefix_query_ranges –

  Map each scheduled query token to the mm_prefix range containing it.

• get_dcp_local_seq_lens –

  While using dcp, kv_cache size stored on each rank may be different,

• get_flashinfer_layout_string –

  Return the layout name in FlashInfer's convention (NHD/HND).

• get_num_attention_heads_from_layers –

  Per-TP-rank num_heads shared by the named Attention layers.

• get_per_layer_parameters –

  Scan layers in layer_names and determine some hyperparameters

• get_supported_kv_cache_layouts –

  Layouts every one of the worker's backends supports, most preferred first.

• infer_global_hyperparameters –

  Currently, FlashInfer backend other than trtllm-gen

• log2_lse_to_ln –

  Convert a base-2 log-sum-exp tensor to natural-log units.

• mamba_get_block_table_tensor –

  Get the block table tensor for mamba kernels from the input

• record_kv_cache_layout –

  Adopt a layout resolved elsewhere (the engine core) in this process.

• reorder_batch_to_split_decodes_and_prefills –

  Reorders the batch to split into prefill and decode requests; places all

• reshape_attn_output_for_spec_decode –

  Reshapes the attention output tensor, so that

• reshape_query_for_spec_decode –

  Reshapes the query tensor for the specified batch size, so that

• resolve_kv_cache_layout –

  Resolve one KV cache layout for the whole model.

• set_triton_ptr_range –

  Bound the bytes Triton kernels may access through views.

• split_decodes_and_prefills –

  Assuming a reordered batch, finds the boundary between prefill and decode

• split_decodes_prefills_and_extends –

  Assuming a reordered batch, finds the boundary between prefill and decode

• split_prefill_chunks –

  Split the prefill requests into chunks such that the total sequence length

PerLayerParameters dataclass

Currently, FlashInfer backend only support models in which all layers share the same values for the following hyperparameters. Should not be used for trtllm-gen backend since it supports different values for the following hyperparameters.

Source code in vllm/v1/attention/backends/utils.py

@dataclass
class PerLayerParameters:
    """Currently, FlashInfer backend only support models in which all layers share
    the same values for the following hyperparameters. Should not be used for
    trtllm-gen backend since it supports different values for the following
    hyperparameters.
    """

    window_left: int
    logits_soft_cap: float | None
    sm_scale: float
    has_sinks: bool = False
    # has same params for all layers
    has_same_window_lefts: bool | None = field(default=None, compare=False)
    has_same_all_params: bool | None = field(default=None, compare=False)

_mamba_align_block_table_kernel(out_ptr, block_table_ptr, block_table_stride_0, block_table_stride_1, seq_lens_ptr, seq_lens_stride, block_size, n_out, BLOCK)

Gather the last n_out block-table entries of each request.

One program per request. This replaces a five-statement PyTorch expression that expanded to seven kernels over a (num_reqs,) tensor -- one element in a single-request decode -- so almost all of its cost was launch overhead rather than work.

Both block-table strides are passed explicitly. The expression this replaces was torch.gather(block_table, 1, ...), which honours any layout, and this is a shared utility -- gdn_attn, mamba_attn, short_conv_attn and linear_attn all call it -- so assuming a column stride of 1 would silently read the wrong slots for a caller holding a non-contiguous view. kimi_k3/nvidia/kda_metadata.py passes both strides for the same reason.

Two annotations above are load-bearing and easy to drop by accident:

block_table_stride_0: tl.int64 Without it req_idx * block_table_stride_0 is computed in i32 -- program_id is i32 and Triton specializes a stride below 2^31 to i32 -- so the product wraps once a row offset reaches 2^31 elements (a 4096 x 625000 table faults from row 3436 on MI355X). torch.gather, which this replaces, indexes in 64-bit and has no such limit, so narrowing it would be a regression. mamba_utils.py:43 annotates its row stride the same way.

do_not_specialize_on_alignment on the three pointers Triton puts pointer alignment in the specialization key, and the MRV1 DBO path hands this kernel seq_lens[request_slice] / block_table_tensor[request_slice] VIEWS (ubatch_utils.py:278/316). When a ubatch starts at a request index that is not a multiple of 4 the view is differently aligned and a fresh variant JIT-compiles mid-step -- 25-28 ms cold, ~1.5 ms warm, and an error under jit_monitor_mode=error. The torch expression this replaced never compiled anything, so that recompile would be a regression introduced by this change.

Source code in vllm/v1/attention/backends/utils.py

@triton.jit(
    do_not_specialize_on_alignment=["out_ptr", "block_table_ptr", "seq_lens_ptr"]
)
def _mamba_align_block_table_kernel(
    out_ptr,
    block_table_ptr,
    block_table_stride_0: tl.int64,
    block_table_stride_1,
    seq_lens_ptr,
    seq_lens_stride,
    block_size,
    n_out,
    BLOCK: tl.constexpr,
):
    """Gather the last ``n_out`` block-table entries of each request.

    One program per request. This replaces a five-statement PyTorch expression
    that expanded to seven kernels over a ``(num_reqs,)`` tensor -- one element
    in a single-request decode -- so almost all of its cost was launch
    overhead rather than work.

    Both block-table strides are passed explicitly. The expression this replaces
    was ``torch.gather(block_table, 1, ...)``, which honours any layout, and
    this is a shared utility -- ``gdn_attn``, ``mamba_attn``, ``short_conv_attn``
    and ``linear_attn`` all call it -- so assuming a column stride of 1 would
    silently read the wrong slots for a caller holding a non-contiguous view.
    ``kimi_k3/nvidia/kda_metadata.py`` passes both strides for the same reason.

    Two annotations above are load-bearing and easy to drop by accident:

    ``block_table_stride_0: tl.int64``
        Without it ``req_idx * block_table_stride_0`` is computed in i32 --
        ``program_id`` is i32 and Triton specializes a stride below 2^31 to
        i32 -- so the product wraps once a row offset reaches 2^31 elements
        (a 4096 x 625000 table faults from row 3436 on MI355X).
        ``torch.gather``, which this replaces, indexes in 64-bit and has no
        such limit, so narrowing it would be a regression.
        ``mamba_utils.py:43`` annotates its row stride the same way.

    ``do_not_specialize_on_alignment`` on the three pointers
        Triton puts pointer alignment in the specialization key, and the MRV1
        DBO path hands this kernel ``seq_lens[request_slice]`` /
        ``block_table_tensor[request_slice]`` VIEWS (ubatch_utils.py:278/316).
        When a ubatch starts at a request index that is not a multiple of 4
        the view is differently aligned and a fresh variant JIT-compiles
        mid-step -- 25-28 ms cold, ~1.5 ms warm, and an error under
        ``jit_monitor_mode=error``. The torch expression this replaced never
        compiled anything, so that recompile would be a regression introduced
        by this change.
    """
    req_idx = tl.program_id(0)
    seq_len = tl.load(seq_lens_ptr + req_idx * seq_lens_stride)
    # A zero-length request in a CUDA graph has an invalid block table; clamp
    # to 0 so it reads a valid row, exactly as the previous clamp_ did.
    start = tl.maximum((seq_len - 1) // block_size, 0)
    offs = tl.arange(0, BLOCK)
    mask = offs < n_out
    vals = tl.load(
        block_table_ptr
        + req_idx * block_table_stride_0
        + (start + offs) * block_table_stride_1,
        mask=mask,
        other=0,
    )
    # `out` is freshly allocated by the caller and therefore contiguous, so its
    # row stride is exactly n_out.
    tl.store(out_ptr + req_idx * n_out + offs, vals, mask=mask)

_merge_supported_layouts(supported_layouts_lists)

Layouts in every supported set, merged as in get_supported_kv_cache_layouts.

Source code in vllm/v1/attention/backends/utils.py

def _merge_supported_layouts(
    supported_layouts_lists: Sequence[Sequence[KVCacheLayout]],
) -> list[KVCacheLayout]:
    """Layouts in every supported set, merged as in
    ``get_supported_kv_cache_layouts``."""
    first = supported_layouts_lists[0]
    if all(layouts == first for layouts in supported_layouts_lists[1:]):
        return list(first)

    priorities: dict[KVCacheLayout, int] = defaultdict(int)
    for preferred_layout, *_ in supported_layouts_lists:
        priorities[preferred_layout] += 1
    supported_layouts = set.intersection(*map(set, supported_layouts_lists))
    candidates = sorted(
        (layout for layout in KVCacheLayout if layout in supported_layouts),
        key=lambda layout: priorities[layout],
        reverse=True,
    )
    if not candidates:
        raise ValueError(
            "No KV cache layout satisfies every supported set: "
            f"{list(map(_layout_names, supported_layouts_lists))}."
        )
    return candidates

compute_mm_prefix_range_tensor(mm_prefix_range, num_seqs, device)

Convert mm_prefix_range dict to padded tensor for Triton kernel.

Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges. Empty ranges have start==end==0, which kernel skips via is_valid check.

Source code in vllm/v1/attention/backends/utils.py

def compute_mm_prefix_range_tensor(
    mm_prefix_range: dict[int, list[tuple[int, int]]] | None,
    num_seqs: int,
    device: torch.device,
) -> torch.Tensor | None:
    """Convert mm_prefix_range dict to padded tensor for Triton kernel.

    Returns shape: (num_seqs, max_ranges, 2) with 0-padding for empty ranges.
    Empty ranges have start==end==0, which kernel skips via is_valid check.
    """
    if mm_prefix_range is None:
        return None

    range_lists = [
        mm_prefix_range.get(i, [(0, 0)]) or [(0, 0)] for i in range(num_seqs)
    ]

    if all(r == [(0, 0)] for r in range_lists):
        return None

    max_ranges = max(len(r) for r in range_lists)
    padded = []
    for r in range_lists:
        padded_r = list(r) + [(0, 0)] * (max_ranges - len(r))
        padded.append(padded_r)
    padded = async_tensor_h2d(padded, dtype=torch.int32, device=device)
    return padded.view(num_seqs, max_ranges, 2)

compute_mm_prefix_ranges(mm_features_per_req, sliding_window=None, clamp_sliding_window=False, span_pad=0)

Per-request PrefixLM bidirectional ranges for image/video tokens.

Source code in vllm/v1/attention/backends/utils.py

def compute_mm_prefix_ranges(
    mm_features_per_req: Sequence[Sequence[MultiModalFeatureSpec]],
    sliding_window: int | None = None,
    clamp_sliding_window: bool = False,
    span_pad: int = 0,
) -> dict[int, list[tuple[int, int]]]:
    """Per-request PrefixLM bidirectional ranges for image/video tokens."""
    req_doc_ranges: dict[int, list[tuple[int, int]]] = {}
    for req_idx, mm_features in enumerate(mm_features_per_req):
        image_doc_ranges = []
        for mm_feature in mm_features:
            if mm_feature.modality not in ("image", "video"):
                continue
            pos_info = mm_feature.mm_position
            if span_pad:
                start = round_down(pos_info.offset, span_pad) + span_pad - 1
                feature_ranges = [(start, pos_info.offset + pos_info.length - 1)]
            else:
                feature_ranges = pos_info.extract_embeds_range()
            for r in feature_ranges:
                if (
                    not clamp_sliding_window
                    and sliding_window is not None
                    and (r[1] - r[0] + 1) > sliding_window
                ):
                    continue
                image_doc_ranges.append(r)
        req_doc_ranges[req_idx] = image_doc_ranges
    return req_doc_ranges

fill_mm_prefix_query_ranges(out, mm_prefix_range, query_start_loc_cpu, seq_lens_cpu)

Map each scheduled query token to the mm_prefix range containing it.

Writes into out, a caller-owned (max_num_batched_tokens, 2) int32 staging buffer, and returns the number of rows written (0 if no range covers any scheduled query token, in which case out is untouched and the caller should skip the mask_mod entirely). Row i holds the absolute [start, end] bounds of the bidirectional range that query token i belongs to, or (-1, -1) when it is outside every range.

mm_prefix ranges never overlap, so "query and key share a range" is equivalent to "the key lies inside the query's own range". The kernel therefore needs no key-side lookup, and this metadata is sized by scheduled query tokens rather than by context length -- bounded by max_num_batched_tokens instead of num_seqs * max_seq_len.

Ranges are absolute prompt positions and may extend past the tokens scheduled so far under chunked prefill; the portion outside the current chunk is simply not recorded. Degenerate ranges (start >= end) are skipped to match the Triton path's start < end validity check.

seq_lens_cpu only needs to be exact for prefill rows, since mm_prefix ranges cover prompt tokens: an over-estimate on a decode row shifts that row's query position further past every range, which still matches nothing.

Source code in vllm/v1/attention/backends/utils.py

def fill_mm_prefix_query_ranges(
    out: np.ndarray,
    mm_prefix_range: dict[int, list[tuple[int, int]]] | None,
    query_start_loc_cpu: torch.Tensor,
    seq_lens_cpu: torch.Tensor,
) -> int:
    """Map each scheduled query token to the mm_prefix range containing it.

    Writes into ``out``, a caller-owned ``(max_num_batched_tokens, 2)`` int32
    staging buffer, and returns the number of rows written (0 if no range
    covers any scheduled query token, in which case ``out`` is untouched and
    the caller should skip the mask_mod entirely). Row ``i`` holds the absolute
    ``[start, end]`` bounds of the bidirectional range that query token ``i``
    belongs to, or ``(-1, -1)`` when it is outside every range.

    mm_prefix ranges never overlap, so "query and key share a range" is
    equivalent to "the key lies inside the query's own range". The kernel
    therefore needs no key-side lookup, and this metadata is sized by scheduled
    query tokens rather than by context length -- bounded by
    ``max_num_batched_tokens`` instead of ``num_seqs * max_seq_len``.

    Ranges are absolute prompt positions and may extend past the tokens
    scheduled so far under chunked prefill; the portion outside the current
    chunk is simply not recorded. Degenerate ranges (``start >= end``) are
    skipped to match the Triton path's ``start < end`` validity check.

    ``seq_lens_cpu`` only needs to be exact for prefill rows, since mm_prefix
    ranges cover prompt tokens: an over-estimate on a decode row shifts that
    row's query position further past every range, which still matches nothing.
    """
    if mm_prefix_range is None:
        return 0

    query_start_loc = query_start_loc_cpu.numpy()
    num_actual_tokens = int(query_start_loc[-1])
    if num_actual_tokens <= 0:
        return 0
    assert num_actual_tokens <= out.shape[0], (
        f"mm_prefix staging buffer holds {out.shape[0]} tokens, got {num_actual_tokens}"
    )

    # Resolve every span before touching `out`, so batches whose ranges all
    # fall outside the scheduled tokens skip the fill entirely.
    spans: list[tuple[int, int, int, int]] = []
    for req_idx, req_ranges in mm_prefix_range.items():
        if not req_ranges:
            continue
        token_start = int(query_start_loc[req_idx])
        query_len = int(query_start_loc[req_idx + 1]) - token_start
        if query_len <= 0:
            continue
        # Absolute position of this request's first scheduled query token.
        context_len = int(seq_lens_cpu[req_idx]) - query_len
        for start, end in req_ranges:
            if start >= end:
                continue
            first = max(start - context_len, 0)
            last = min(end - context_len, query_len - 1)
            if first > last:
                continue
            spans.append((token_start + first, token_start + last + 1, start, end))

    if not spans:
        return 0

    out[:num_actual_tokens] = -1
    for row_start, row_end, start, end in spans:
        out[row_start:row_end] = (start, end)
    return num_actual_tokens

get_dcp_local_seq_lens(seq_lens, dcp_size=1, dcp_rank=None, cp_kv_cache_interleave_size=1)

While using dcp, kv_cache size stored on each rank may be different, use this function to calculate split decode seq_lens of each dcp rank. Only consider dcp now, we can extend the case of cp based on this.

Source code in vllm/v1/attention/backends/utils.py

def get_dcp_local_seq_lens(
    seq_lens: torch.Tensor,
    dcp_size: int = 1,
    dcp_rank: int | torch.Tensor | None = None,
    cp_kv_cache_interleave_size: int = 1,
) -> torch.Tensor:
    """While using dcp, kv_cache size stored on each rank may be different,
    use this function to calculate split decode seq_lens of each dcp rank.
    Only consider dcp now, we can extend the case of cp based on this.
    """
    seq_lens_i32 = seq_lens.to(torch.int32)
    rank_offsets: int | torch.Tensor
    if dcp_rank is None:
        rank_offsets = torch.arange(
            dcp_size,
            dtype=torch.int32,
            device=seq_lens.device,
        ).view(
            *((1,) * seq_lens_i32.dim()),
            dcp_size,
        )
        seq_lens_tiled = seq_lens_i32.unsqueeze(-1)
    elif isinstance(dcp_rank, torch.Tensor):
        assert dcp_rank.dtype == torch.int32
        assert dcp_rank.device == seq_lens.device
        assert dcp_rank.numel() == 1
        rank_offsets = dcp_rank
        seq_lens_tiled = seq_lens_i32
    else:
        # Use the Python scalar directly to avoid a synchronizing H2D copy.
        rank_offsets = dcp_rank
        seq_lens_tiled = seq_lens_i32
    base = (
        seq_lens_tiled
        // cp_kv_cache_interleave_size
        // dcp_size
        * cp_kv_cache_interleave_size
    )
    remainder = seq_lens_tiled - base * dcp_size
    remainder = torch.clip(
        remainder - rank_offsets * cp_kv_cache_interleave_size,
        0,
        cp_kv_cache_interleave_size,
    )
    dcp_local_seq_lens = base + remainder
    return dcp_local_seq_lens

get_flashinfer_layout_string(layout)

Return the layout name in FlashInfer's convention (NHD/HND).

Source code in vllm/v1/attention/backends/utils.py

def get_flashinfer_layout_string(layout: KVCacheLayout) -> str:
    """Return the layout name in FlashInfer's convention (NHD/HND)."""
    assert layout.name in _FLASHINFER_LAYOUT_NAMES, (
        f"KV cache layout {layout.name} has no FlashInfer equivalent; FlashInfer "
        "rejects it in supported_kv_cache_layouts"
    )
    return _FLASHINFER_LAYOUT_NAMES[layout.name]

get_num_attention_heads_from_layers(vllm_config, layer_names)

Per-TP-rank num_heads shared by the named Attention layers.

Use in metadata builders whose plan-time allocations depend on the head count: the model-wide get_num_attention_heads() is wrong for models with non-uniform per-layer head counts. All layers in one attention group must agree on num_heads; this is asserted. Returns None when no matching Attention layer is found.

Source code in vllm/v1/attention/backends/utils.py

def get_num_attention_heads_from_layers(
    vllm_config: VllmConfig, layer_names: list[str]
) -> int | None:
    """Per-TP-rank ``num_heads`` shared by the named Attention layers.

    Use in metadata builders whose plan-time allocations depend on the
    head count: the model-wide ``get_num_attention_heads()`` is wrong
    for models with non-uniform per-layer head counts. All layers in
    one attention group must agree on ``num_heads``; this is asserted.
    Returns ``None`` when no matching Attention layer is found.
    """
    attn_layers = get_layers_from_vllm_config(
        vllm_config,
        AttentionLayerBase,  # type: ignore[type-abstract]
        layer_names,
    )
    if not attn_layers:
        return None
    heads = {
        cast(Any, getattr(layer, "impl", layer)).num_heads
        for layer in attn_layers.values()
    }
    assert len(heads) == 1, (
        f"All layers in one attention group must share num_heads; "
        f"got {heads} for {layer_names}."
    )
    return heads.pop()

get_per_layer_parameters(vllm_config, layer_names, cls_)

Scan layers in layer_names and determine some hyperparameters to use during plan.

Source code in vllm/v1/attention/backends/utils.py

def get_per_layer_parameters(
    vllm_config: VllmConfig, layer_names: list[str], cls_: type["AttentionImpl"]
) -> dict[str, PerLayerParameters]:
    """Scan layers in `layer_names` and determine some hyperparameters
    to use during `plan`.
    """
    layers = get_layers_from_vllm_config(
        vllm_config,
        AttentionLayerBase,  # type: ignore[type-abstract]
        layer_names,
    )
    per_layer_params: dict[str, PerLayerParameters] = {}

    for key, layer in layers.items():
        impl = find_attention_impl_variant(layer.impl, cls_)
        assert impl is not None

        # Infer hyperparameters from the attention layer
        window_size = getattr(impl, "sliding_window", None)
        window_left = window_size[0] if window_size is not None else -1
        logits_soft_cap = getattr(impl, "logits_soft_cap", None)
        sm_scale = impl.scale
        has_sinks = getattr(impl, "sinks", None) is not None

        per_layer_params[key] = PerLayerParameters(
            window_left, logits_soft_cap, sm_scale, has_sinks
        )

    return per_layer_params

get_supported_kv_cache_layouts(backends)

Layouts every one of the worker's backends supports, most preferred first.

Every backend declares the layouts its kernels support, most preferred first (supported_kv_cache_layouts), or None when any layout works; workers where nothing declares follow the default preference. Identical declarations keep their order; otherwise the layout the most backends put first wins, ties keeping the enum order. An empty intersection is a hard error.

Source code in vllm/v1/attention/backends/utils.py

def get_supported_kv_cache_layouts(
    backends: Iterable[type[AttentionBackend]],
) -> list[KVCacheLayout]:
    """Layouts every one of the worker's backends supports, most preferred first.

    Every backend declares the layouts its kernels support, most preferred first
    (``supported_kv_cache_layouts``), or None when any layout works; workers where
    nothing declares follow the default preference. Identical declarations keep
    their order; otherwise the layout the most backends put first wins, ties
    keeping the enum order. An empty intersection is a hard error.
    """
    return _merge_supported_layouts(
        [
            layouts
            for backend in backends
            if (layouts := backend.supported_kv_cache_layouts()) is not None
        ]
        or [_DEFAULT_LAYOUT_PREFERENCE]
    )

infer_global_hyperparameters(per_layer_params)

Currently, FlashInfer backend other than trtllm-gen only support models in which all layers share the same values for the following hyperparameters: - window_left - logits_soft_cap - sm_scale

So this function asserts that all layers share the same values for these hyperparameters and returns the global values.

Source code in vllm/v1/attention/backends/utils.py

def infer_global_hyperparameters(
    per_layer_params: dict[str, PerLayerParameters],
) -> PerLayerParameters:
    """Currently, FlashInfer backend other than trtllm-gen
    only support models in which all layers share
    the same values for the following hyperparameters:
    - `window_left`
    - `logits_soft_cap`
    - `sm_scale`

    So this function asserts that all layers share the same values for these
    hyperparameters and returns the global values.
    """
    assert len(per_layer_params) > 0, "No attention layers found in the model."

    param_sets = list(per_layer_params.values())
    global_params = param_sets[0]

    global_params.has_same_window_lefts = all(
        params.window_left == global_params.window_left for params in param_sets
    )
    global_params.has_same_all_params = all(
        params == global_params for params in param_sets
    )

    return global_params

log2_lse_to_ln(lse)

Convert a base-2 log-sum-exp tensor to natural-log units.

Source code in vllm/v1/attention/backends/utils.py

def log2_lse_to_ln(lse: torch.Tensor) -> torch.Tensor:
    """Convert a base-2 log-sum-exp tensor to natural-log units."""
    return lse * _LN_2

mamba_get_block_table_tensor(block_table, seq_lens, kv_cache_spec, mamba_cache_mode)

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

• "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.

Source code in vllm/v1/attention/backends/utils.py

def mamba_get_block_table_tensor(
    block_table: torch.Tensor,
    seq_lens: torch.Tensor,
    kv_cache_spec: KVCacheSpec,
    mamba_cache_mode: str,
) -> torch.Tensor:
    """Get the block table tensor for mamba kernels from the input
    common_attn_metadata.block_table_tensor given different mamba cache modes.

    - "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.
    """
    if mamba_cache_mode == "none":
        return block_table
    else:
        assert isinstance(kv_cache_spec, MambaSpec)
        # The arithmetic here is trivial -- a divide, a clamp, an arange, an
        # add, a cast and a gather -- but as PyTorch ops it was seven kernel
        # launches over a (num_reqs,) tensor, which holds a single element in a
        # one-request decode. One Triton kernel does the same work in one
        # launch. Zero-length requests in a CUDA graph still clamp to row 0.
        num_reqs = block_table.shape[0]
        n_out = 1 + kv_cache_spec.num_speculative_blocks
        # The expression this replaced took its row count from `seq_lens` (a
        # `torch.gather` whose index tensor is built from it), so a mismatch
        # raised and a cross-device pair raised. Taking `num_reqs` from
        # `block_table` instead silently changed all of that: with 4 block-table
        # rows and 2 `seq_lens`, the old code returned 2 correct rows and the new
        # kernel returns 4, two of them read from memory past the end of
        # `seq_lens`; a pinned-CPU `seq_lens` with a GPU `block_table` used to
        # raise and would now be read zero-copy from host memory.
        #
        # No in-tree caller does any of that today, but this helper is shared by
        # gdn_attn, mamba_attn, short_conv_attn and linear_attn. Both are host
        # metadata, so this costs no synchronisation, and it keeps the CPU
        # branch below and the GPU kernel agreeing on the same inputs.
        assert seq_lens.shape[0] == num_reqs, (
            f"seq_lens has {seq_lens.shape[0]} rows, block_table has {num_reqs}"
        )
        assert seq_lens.device == block_table.device, (
            f"seq_lens is on {seq_lens.device}, block_table on {block_table.device}"
        )
        if not block_table.is_cuda:
            # Triton runs on the accelerator only, and this helper IS called
            # with CPU tensors: the GDN and linear-attention metadata builders
            # are unit-tested on CPU (`DEVICE = torch.device("cpu")`). The
            # expression this replaced was plain torch and handled that, so
            # keep it for the CPU path rather than narrowing the contract.
            # kimi_k3/nvidia/kda_metadata.py asserts `is_cuda` instead; a
            # shared utility cannot, because its callers legitimately pass CPU.
            start_indices = (seq_lens - 1) // kv_cache_spec.block_size
            start_indices = start_indices.clamp(min=0)
            offsets = torch.arange(n_out, device=block_table.device, dtype=torch.int32)
            return torch.gather(
                block_table, 1, (start_indices.unsqueeze(1) + offsets).to(torch.int64)
            )
        out = block_table.new_empty((num_reqs, n_out))
        if num_reqs == 0:
            return out
        _mamba_align_block_table_kernel[(num_reqs,)](
            out,
            block_table,
            block_table.stride(0),
            block_table.stride(1),
            seq_lens,
            seq_lens.stride(0),
            kv_cache_spec.block_size,
            n_out,
            BLOCK=triton.next_power_of_2(n_out),
        )
        return out

record_kv_cache_layout(cache_config, layout_name)

Adopt a layout resolved elsewhere (the engine core) in this process.

Source code in vllm/v1/attention/backends/utils.py

def record_kv_cache_layout(cache_config: CacheConfig, layout_name: str) -> None:
    """Adopt a layout resolved elsewhere (the engine core) in this process."""
    layout = _layout_from_name(layout_name)
    existing = cache_config.kv_cache_layout
    if existing is not None and existing != layout.name:
        raise ValueError(
            f"KV cache layout is already resolved to {existing}; "
            f"cannot change it to {layout.name}."
        )
    cache_config.kv_cache_layout = layout.name

reorder_batch_to_split_decodes_and_prefills(input_batch, scheduler_output, decode_threshold=1)

Reorders the batch to split into prefill and decode requests; places all requests with <= decode_threshold tokens at the front of the batch.

The batch is reordered into 4 regions

decode: (num_scheduled <= threshold AND is not prefilling) short_extend: (num_scheduled <= threshold AND is chunked prefilling) long_extend: (num_scheduled > threshold AND is chunked prefilling) prefill: (num_computed == 0) # First chunks

Returns:

• bool –

  True if the batch was modified, False otherwise.

Source code in vllm/v1/attention/backends/utils.py

def reorder_batch_to_split_decodes_and_prefills(
    input_batch: "InputBatch",
    scheduler_output: "SchedulerOutput",
    decode_threshold: int = 1,
) -> bool:
    """Reorders the batch to split into prefill and decode requests; places all
    requests with <= decode_threshold tokens at the front of the batch.

    The batch is reordered into 4 regions:
        decode:        (num_scheduled <= threshold AND is not prefilling)
        short_extend:  (num_scheduled <= threshold AND is chunked prefilling)
        long_extend:   (num_scheduled > threshold AND is chunked prefilling)
        prefill:       (num_computed == 0)   # First chunks

    Returns:
        True if the batch was modified, False otherwise.

    """
    num_reqs = len(input_batch.req_ids)
    num_scheduled_tokens = [
        scheduler_output.num_scheduled_tokens[id] for id in input_batch.req_ids
    ]
    num_scheduled_tokens_np = np.array(num_scheduled_tokens)
    num_computed_tokens_np = input_batch.num_computed_tokens_cpu[:num_reqs]
    num_prompt_tokens_np = input_batch.num_prompt_tokens[:num_reqs]

    has_context = num_computed_tokens_np > 0
    is_below_threshold = num_scheduled_tokens_np <= decode_threshold
    done_prefilling = num_computed_tokens_np >= num_prompt_tokens_np

    # Mutually exclusive categories (exactly one True per request):
    # 1. No context yet -> prefill
    # 2. Has context, above threshold -> long_extend
    # 3. Has context, below threshold, still prefilling -> short_extend
    # 4. Has context, below threshold, done prefilling -> decode
    is_pure_prefill = ~has_context
    is_long_extend = has_context & ~is_below_threshold
    is_short_extend = has_context & is_below_threshold & ~done_prefilling
    is_decode = has_context & is_below_threshold & done_prefilling

    # Desired order: decode → short_extend → long_extend → prefill
    req_regions = np.zeros(num_reqs, dtype=np.int32)  # 0 = decode by default
    req_regions[is_short_extend] = 1
    req_regions[is_long_extend] = 2
    req_regions[is_pure_prefill] = 3

    num_decodes = int(is_decode.sum())
    num_short_extends = int(is_short_extend.sum())
    num_long_extends = int(is_long_extend.sum())
    num_prefills = int(is_pure_prefill.sum())

    target_regions = np.repeat(
        [0, 1, 2, 3],
        [num_decodes, num_short_extends, num_long_extends, num_prefills],
    ).astype(np.int32)

    needs_swap = req_regions != target_regions

    if not needs_swap.any():
        return False

    # Extract indices that need swapping and sort by target region
    orig_indices = np.where(needs_swap)[0]
    sorted_order = np.argsort(req_regions[needs_swap], kind="stable")
    src_indices = orig_indices[sorted_order]

    src_dest_map = {int(src): int(dst) for src, dst in zip(src_indices, orig_indices)}

    for src in src_dest_map:
        dst = src_dest_map[src]
        while src != dst:
            input_batch.swap_states(src, dst)
            # Mark dst as done by updating its destination to itself
            next_dst = src_dest_map.get(dst, dst)
            src_dest_map[dst] = dst
            dst = next_dst

    return True

reshape_attn_output_for_spec_decode(attn_output)

Reshapes the attention output tensor, so that the batch_size and seq_len dimensions are combined.

Source code in vllm/v1/attention/backends/utils.py

def reshape_attn_output_for_spec_decode(attn_output: torch.Tensor) -> torch.Tensor:
    """Reshapes the attention output tensor, so that
    the batch_size and seq_len dimensions are combined.
    """
    if attn_output.dim() == 3:
        # Already in the correct shape
        return attn_output
    assert attn_output.dim() == 4, f"attn_output must be 4D, got {attn_output.dim()}D"
    total_tokens = attn_output.shape[0] * attn_output.shape[1]
    return attn_output.view(total_tokens, attn_output.shape[2], attn_output.shape[3])

reshape_query_for_spec_decode(query, batch_size)

Reshapes the query tensor for the specified batch size, so that it has shape (batch_size, seq_len, num_heads, head_dim).

Source code in vllm/v1/attention/backends/utils.py

def reshape_query_for_spec_decode(query: torch.Tensor, batch_size: int) -> torch.Tensor:
    """Reshapes the query tensor for the specified batch size, so that
    it has shape (batch_size, seq_len, num_heads, head_dim).
    """
    assert query.dim() == 3, f"query must be 3D, got {query.dim()}D"
    total_tokens = query.shape[0]
    num_heads = query.shape[1]
    head_dim = query.shape[2]
    assert total_tokens % batch_size == 0, (
        f"{total_tokens=} is not divisible by {batch_size=}"
    )
    seq_len = total_tokens // batch_size
    return query.view(batch_size, seq_len, num_heads, head_dim)

resolve_kv_cache_layout(vllm_config, supported_layouts, kv_cache_specs=None)

Resolve one KV cache layout for the whole model.

Runs once in the engine core. Every worker reports the layouts its backends support, most preferred first (get_supported_kv_cache_layouts). Pipeline stages hold different layers, so their lists may differ; the candidates are the layouts every worker supports, and the layout the most workers put first wins. Workers whose backends declare nothing do not vote. Specs mixing HNC shapes narrow the candidates to block-compact layouts. An explicit VLLM_KV_CACHE_LAYOUT must be one of the candidates or resolution fails, with the legacy NHD/HND names as aliases for LBNHC/LBHNC; the connector's preference is used when compatible and dropped with a warning otherwise. A layout already present on cache_config wins outright, and the result is recorded there (see CacheConfig.kv_cache_layout); it reaches workers through the set_kv_cache_layout RPC and KVCacheConfig.kv_cache_layout.

Source code in vllm/v1/attention/backends/utils.py

def resolve_kv_cache_layout(
    vllm_config: VllmConfig,
    supported_layouts: list[list[str]],
    kv_cache_specs: Iterable[KVCacheSpec] | None = None,
) -> KVCacheLayout:
    """Resolve one KV cache layout for the whole model.

    Runs once in the engine core. Every worker reports the layouts its backends
    support, most preferred first (``get_supported_kv_cache_layouts``). Pipeline
    stages hold different layers, so their lists may differ; the candidates are
    the layouts every worker supports, and the layout the most workers put first
    wins. Workers whose backends declare nothing do not vote. Specs mixing HNC
    shapes narrow the candidates to block-compact layouts. An explicit
    ``VLLM_KV_CACHE_LAYOUT`` must be one of the candidates or resolution fails,
    with the legacy ``NHD``/``HND`` names as aliases for ``LBNHC``/``LBHNC``; the
    connector's preference is used when compatible and dropped with a warning
    otherwise. A layout already present on ``cache_config`` wins outright, and
    the result is recorded there (see ``CacheConfig.kv_cache_layout``); it
    reaches workers through the ``set_kv_cache_layout`` RPC and
    ``KVCacheConfig.kv_cache_layout``.
    """
    cache_config = vllm_config.cache_config
    if cache_config.kv_cache_layout is not None:
        return cache_config.get_resolved_kv_cache_layout()

    assert supported_layouts and all(supported_layouts), (
        "No worker reported supported KV cache layouts."
    )
    reported = [
        tuple(_layout_from_name(name) for name in names) for names in supported_layouts
    ]
    # Workers where no backend declares a set report the default preference;
    # like undeclared backends of one worker, they don't vote.
    candidates = _merge_supported_layouts(
        [layouts for layouts in reported if layouts != _DEFAULT_LAYOUT_PREFERENCE]
        or reported
    )

    # A block-compact layout means the block is densely packed in memory, so any mix of
    # specs can re-interpret HNC with different sizes as long as the total number of
    # bytes is the same. If not block-compact, each spec must agree on HNC to alias
    # the same page (this aliasing is done by the Hybrid Memory Allocator, HMA).
    # Specs without per-layer views lay out their own raw backing tensor.
    kv_cache_specs = tuple(
        spec for spec in kv_cache_specs or () if spec.has_layer_views
    )
    hnc_shapes = {
        (spec.num_heads, spec.num_states, spec.page_size_bytes)
        for spec in kv_cache_specs
    }
    if len(hnc_shapes) > 1:
        candidates = [m for m in candidates if m.is_block_compact]
        if not candidates:
            raise ValueError(
                "Specs with mixed HNC shapes need a block-compact layout, but "
                f"none is in every supported set: {supported_layouts}."
            )

    # Self-addressed per-request rings (e.g. the kpool tail) are replicated
    # state like Mamba and don't need a separate draft group.
    dcp_sharding = {
        spec.dcp_sharded
        for spec in kv_cache_specs
        if isinstance(spec, AttentionSpec) and spec.uses_slot_mapping
    }
    page_sizes = {spec.page_size_bytes for spec in kv_cache_specs}
    if len(dcp_sharding) > 1 and len(page_sizes) > 1:
        # Sharded target and replicated draft caches need independent groups.
        # Block-outer layouts can pack those groups without equalizing pages.
        candidates = [layout for layout in candidates if layout.is_block_outermost]
        if not candidates:
            raise ValueError(
                "DCP with a replicated draft and mixed KV page sizes requires "
                f"a block-outer KV cache layout; supported sets: {supported_layouts}."
            )

    if (requested := envs.VLLM_KV_CACHE_LAYOUT) is not None:
        layout = _layout_from_name(requested)
        if layout not in candidates:
            raise ValueError(
                f"VLLM_KV_CACHE_LAYOUT={requested} does not satisfy every "
                f"supported set; valid layouts: {_layout_names(candidates)}."
            )
    elif (connector := get_kv_connector_cache_layout(vllm_config)) is not None:
        layout = _layout_from_name(connector)
        if layout not in candidates:
            logger.warning_once(
                f"KV connector cache layout {connector} does not satisfy every "
                f"supported set; valid layouts: {_layout_names(candidates)}. "
                f"Using {candidates[0].name} instead."
            )
            layout = candidates[0]
    else:
        layout = candidates[0]

    logger.info_once("Using %s KV cache layout.", layout.name)
    cache_config.kv_cache_layout = layout.name
    return layout

set_triton_ptr_range(region, *views)

Bound the bytes Triton kernels may access through views.

Triton's AMD backend uses 32-bit buffer addressing only for pointer arguments within 2 GiB. It checks ptr_range() first and otherwise the size of the whole storage, so views into one large allocation (such as a layer of the KV cache) lose it. Each view gets the distance from its data_ptr() to the end of region; kernels must not access past it.

Source code in vllm/v1/attention/backends/utils.py

def set_triton_ptr_range(region: torch.Tensor, *views: torch.Tensor) -> None:
    """Bound the bytes Triton kernels may access through ``views``.

    Triton's AMD backend uses 32-bit buffer addressing only for pointer
    arguments within 2 GiB. It checks ``ptr_range()`` first and otherwise the
    size of the whole storage, so views into one large allocation (such as a
    layer of the KV cache) lose it. Each view gets the distance from its
    ``data_ptr()`` to the end of ``region``; kernels must not access past it.
    """
    if region.numel() == 0:
        return
    last = sum((n - 1) * stride for n, stride in zip(region.shape, region.stride()))
    region_end = region.data_ptr() + (last + 1) * region.element_size()
    for view in views:
        view.ptr_range = functools.partial(int, region_end - view.data_ptr())

split_decodes_and_prefills(common_attn_metadata, decode_threshold=1, require_uniform=False, treat_short_extends_as_decodes=True)

Assuming a reordered batch, finds the boundary between prefill and decode requests.

The batch is expected to be ordered as

decode → short_extend → long_extend → prefill

Parameters:

• common_attn_metadata

  (CommonAttentionMetadata) –

  CommonAttentionMetadata object containing the batch metadata.

• decode_threshold

  (int, default: 1 ) –

  The maximum query length to be considered a decode.

• require_uniform

  (bool, default: False ) –

  If True, requires that all decode requests have the same query length. When set, some queries may be considered prefills even if they are <= decode_threshold, in order to ensure uniformity.

• treat_short_extends_as_decodes

  (bool, default: True ) –

  If True (default), short extends (query_len <= threshold but still prefilling) are counted as decodes. If False, they are counted as prefills.

Returns:

• num_decodes ( int ) –

  The number of decode requests.

• num_prefills ( int ) –

  The number of prefill requests.

• num_decode_tokens ( int ) –

  The number of tokens in the decode requests.

• num_prefill_tokens ( int ) –

  The number of tokens in the prefill requests.

Source code in vllm/v1/attention/backends/utils.py

def split_decodes_and_prefills(
    common_attn_metadata: CommonAttentionMetadata,
    decode_threshold: int = 1,
    require_uniform: bool = False,
    treat_short_extends_as_decodes: bool = True,
) -> tuple[int, int, int, int]:
    """Assuming a reordered batch, finds the boundary between prefill and decode
    requests.

    The batch is expected to be ordered as:
        decode → short_extend → long_extend → prefill

    Args:
        common_attn_metadata: CommonAttentionMetadata object containing the
            batch metadata.
        decode_threshold: The maximum query length to be considered a decode.
        require_uniform: If True, requires that all decode requests have the
            same query length. When set, some queries may be considered prefills
            even if they are <= decode_threshold, in order to ensure uniformity.
        treat_short_extends_as_decodes: If True (default), short extends
            (query_len <= threshold but still prefilling) are counted as
            decodes. If False, they are counted as prefills.

    Returns:
        num_decodes: The number of decode requests.
        num_prefills: The number of prefill requests.
        num_decode_tokens: The number of tokens in the decode requests.
        num_prefill_tokens: The number of tokens in the prefill requests.

    """
    max_query_len = common_attn_metadata.max_query_len
    num_reqs = common_attn_metadata.num_reqs
    num_tokens = common_attn_metadata.num_actual_tokens
    query_start_loc = common_attn_metadata.query_start_loc_cpu

    if (
        max_query_len <= decode_threshold
        and (not require_uniform or decode_threshold <= 1)
        and treat_short_extends_as_decodes
    ):
        return num_reqs, 0, num_tokens, 0

    query_lens = query_start_loc[1:] - query_start_loc[:-1]
    if query_lens[0].item() > decode_threshold:
        # first request is not decode, so no decode requests
        return 0, num_reqs, 0, num_tokens

    if require_uniform:
        # check if we are in a padded uniform batch; this is used for full-CGs, some
        # requests may have a query length of 0 but since they are padding its fine
        # to treat them as decodes (ensures num_decodes matches the captured size)
        if treat_short_extends_as_decodes and torch.all(
            (query_lens == query_lens[0]) | (query_lens == 0)
        ):
            return num_reqs, 0, num_tokens, 0  # all decodes
        is_prefill = (query_lens != query_lens[0]) & (query_lens != 0)
    else:
        is_prefill = query_lens > decode_threshold

    if not treat_short_extends_as_decodes:
        assert common_attn_metadata.is_prefilling is not None
        is_prefill |= common_attn_metadata.is_prefilling

    if not torch.any(is_prefill):
        return num_reqs, 0, num_tokens, 0

    first_prefill = is_prefill.int().argmax(dim=-1).item()
    num_decodes = first_prefill
    num_prefills = num_reqs - num_decodes
    num_decode_tokens = query_start_loc[first_prefill].item()
    num_prefill_tokens = num_tokens - num_decode_tokens
    return (num_decodes, num_prefills, num_decode_tokens, num_prefill_tokens)

split_decodes_prefills_and_extends(common_attn_metadata, decode_threshold=1)

Assuming a reordered batch, finds the boundary between prefill and decode requests.

Parameters:

• common_attn_metadata

  (CommonAttentionMetadata) –

  CommonAttentionMetadata object containing the batch metadata.

• decode_threshold

  (int, default: 1 ) –

  The maximum query length to be considered a decode.

Returns:

• num_decodes ( int ) –

  The number of decode requests.

• num_extends ( int ) –

  The number of extend requests.

• num_prefills ( int ) –

  The number of prefill requests.

• num_decode_tokens ( int ) –

  The number of tokens in the decode requests.

• num_extend_tokens ( int ) –

  The number of tokens in the extend requests.

• num_prefill_tokens ( int ) –

  The number of tokens in the prefill requests.

Source code in vllm/v1/attention/backends/utils.py

def split_decodes_prefills_and_extends(
    common_attn_metadata: CommonAttentionMetadata,
    decode_threshold: int = 1,
) -> tuple[int, int, int, int, int, int]:
    """Assuming a reordered batch, finds the boundary between prefill and decode
    requests.

    Args:
        common_attn_metadata: CommonAttentionMetadata object containing the
            batch metadata.
        decode_threshold: The maximum query length to be considered a decode.

    Returns:
        num_decodes: The number of decode requests.
        num_extends: The number of extend requests.
        num_prefills: The number of prefill requests.
        num_decode_tokens: The number of tokens in the decode requests.
        num_extend_tokens: The number of tokens in the extend requests.
        num_prefill_tokens: The number of tokens in the prefill requests.

    """
    max_query_len = common_attn_metadata.max_query_len
    num_reqs = common_attn_metadata.num_reqs
    num_tokens = common_attn_metadata.num_actual_tokens
    query_start_loc = common_attn_metadata.query_start_loc_cpu

    if max_query_len <= decode_threshold:
        return num_reqs, 0, 0, num_tokens, 0, 0

    # Upper bound is exact for prefill rows; decode rows still satisfy
    # seq_len > query_len under the optimistic bound, so `seq_lens ==
    # query_lens` identifies prefills correctly either way.
    assert common_attn_metadata.seq_lens_cpu_upper_bound is not None
    seq_lens = common_attn_metadata.seq_lens_cpu_upper_bound

    query_lens = query_start_loc[1:] - query_start_loc[:-1]
    is_prefill_or_extend = query_lens > decode_threshold
    is_prefill = (seq_lens == query_lens) & is_prefill_or_extend
    first_extend = is_prefill_or_extend.int().argmax(dim=-1).item()
    first_prefill = is_prefill.int().argmax(dim=-1).item()
    num_decodes = first_extend
    num_decode_tokens = query_start_loc[first_extend].item()
    if not torch.any(is_prefill_or_extend):
        return (num_decodes, 0, 0, num_decode_tokens, 0, 0)

    num_prefills_or_extends = num_reqs - num_decodes
    num_prefill_or_extend_tokens = num_tokens - num_decode_tokens
    if not torch.any(is_prefill):
        return (
            num_decodes,
            num_prefills_or_extends,
            0,
            num_decode_tokens,
            num_prefill_or_extend_tokens,
            0,
        )

    num_extends = first_prefill - num_decodes
    num_prefills = num_reqs - first_prefill

    num_prefill_tokens = num_tokens - query_start_loc[first_prefill]
    num_extend_tokens = num_prefill_or_extend_tokens - num_prefill_tokens
    return (
        num_decodes,
        num_extends,
        num_prefills,
        num_decode_tokens,
        num_extend_tokens,
        num_prefill_tokens,
    )

split_prefill_chunks(seq_lens_cpu, workspace_size, request_offset=0)

Split the prefill requests into chunks such that the total sequence length of each chunk is less than or equal to the workspace size.

Parameters:

• seq_lens_cpu

  (Tensor) –

  The sequence lengths of the prefill requests on CPU.

• workspace_size

  (int) –

  The maximum workspace size (in tokens) per chunk.

• request_offset

  (int, default: 0 ) –

  The offset to add to the request indices.

Returns:

• list[tuple[int, int]] –

  A list of tuples of (reqs_start, reqs_end) representing chunk boundaries.

Source code in vllm/v1/attention/backends/utils.py

def split_prefill_chunks(
    seq_lens_cpu: torch.Tensor, workspace_size: int, request_offset: int = 0
) -> list[tuple[int, int]]:
    """Split the prefill requests into chunks such that the total sequence length
    of each chunk is less than or equal to the workspace size.

    Args:
        seq_lens_cpu: The sequence lengths of the prefill requests on CPU.
        workspace_size: The maximum workspace size (in tokens) per chunk.
        request_offset: The offset to add to the request indices.

    Returns:
        A list of tuples of (reqs_start, reqs_end) representing chunk boundaries.

    """
    chunk_bounds = []
    i, n = 0, len(seq_lens_cpu)
    assert torch.all(seq_lens_cpu <= workspace_size).item()

    while i < n:
        start, chunk_total = i, 0
        while i < n and (chunk_total + (s := seq_lens_cpu[i].item())) <= workspace_size:
            chunk_total += s
            i += 1
        chunk_bounds.append((start + request_offset, i + request_offset))
    return chunk_bounds