vllm.v1.attention.backends.mla.compressor_utils

Functions:

• get_compressed_slot_mapping –

  Slot mapping for writing the compressed states of num_tokens tokens.

• get_dspark_swa_index_width –

  Return the padded width of non-causal DSpark SWA indices.

get_compressed_slot_mapping(num_tokens, slot_mapping, query_start_loc, seq_lens, block_table, block_size, compress_ratio, out=None)

Slot mapping for writing the compressed states of num_tokens tokens.

Every compress_ratio tokens share one compressed state, written by the last of them: that token maps to the state's slot, the others to PAD. A token whose own slot_mapping entry is PAD maps to PAD too: SWA bounded replay recomputes tokens whose KV is cached already, and their compressed states must not be rewritten either.

Source code in vllm/v1/attention/backends/mla/compressor_utils.py

def get_compressed_slot_mapping(
    num_tokens: int,
    slot_mapping: torch.Tensor,
    query_start_loc: torch.Tensor,
    seq_lens: torch.Tensor,
    block_table: torch.Tensor,
    block_size: int,
    compress_ratio: int,
    out: torch.Tensor | None = None,
) -> torch.Tensor:
    """Slot mapping for writing the compressed states of ``num_tokens`` tokens.

    Every ``compress_ratio`` tokens share one compressed state, written by the
    last of them: that token maps to the state's slot, the others to PAD. A
    token whose own ``slot_mapping`` entry is PAD maps to PAD too: SWA bounded
    replay recomputes tokens whose KV is cached already, and their compressed
    states must not be rewritten either.
    """
    if out is not None:
        # Guard: for padded / invalid sequences.
        # Negative positions produce bogus block indices that lead to illegal memory
        # accesses inside the block_table load.
        # NOTE: Fill -1 to the whole tensor, not just the first `num_tokens`.
        out.fill_(-1)
        compressed_slot_mapping = out[:num_tokens]
    else:
        compressed_slot_mapping = torch.full(
            (num_tokens,), -1, dtype=torch.int64, device=query_start_loc.device
        )

    _COMPRESSED_SLOT_MAPPING_KERNEL(
        compressed_slot_mapping,
        slot_mapping,
        query_start_loc,
        seq_lens,
        block_table,
        block_size,
        compress_ratio,
    )
    return compressed_slot_mapping

get_dspark_swa_index_width(window_size, num_speculative_tokens)

Return the padded width of non-causal DSpark SWA indices.

Source code in vllm/v1/attention/backends/mla/compressor_utils.py

def get_dspark_swa_index_width(
    window_size: int,
    num_speculative_tokens: int,
) -> int:
    """Return the padded width of non-causal DSpark SWA indices."""
    width = max(int(window_size), 0) + max(int(num_speculative_tokens), 0)
    return cdiv(width, _DSPARK_SWA_INDEX_ALIGNMENT) * _DSPARK_SWA_INDEX_ALIGNMENT