vllm.v1.attention.backends.mla.compressor_utils ¶

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

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