Skip to content

vllm.v1.attention.backends.mla.triton_mla

_reserve_attn_logits_workspace(vllm_config, max_decode_rows, q_num_heads, kv_lora_rank)

Pre-size the shared workspace for the decode split-KV attn logits.

Reserving at the worst case (max_model_len -> max num_kv_splits, max_decode_rows) before warmup/cudagraph capture means the per-call get_simultaneous in forward_mqa never has to grow the buffer at runtime (which would raise once the workspace is locked).

Source code in vllm/v1/attention/backends/mla/triton_mla.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
def _reserve_attn_logits_workspace(
    vllm_config: VllmConfig, max_decode_rows: int, q_num_heads: int, kv_lora_rank: int
) -> None:
    """Pre-size the shared workspace for the decode split-KV attn logits.

    Reserving at the worst case (max_model_len -> max num_kv_splits,
    max_decode_rows) before warmup/cudagraph capture means the per-call
    ``get_simultaneous`` in ``forward_mqa`` never has to grow the buffer at
    runtime (which would raise once the workspace is locked).
    """
    if not is_workspace_manager_initialized():
        return
    max_splits = _compute_num_kv_splits(
        vllm_config.model_config.max_model_len,
        current_platform.num_compute_units(),
    )
    current_workspace_manager().get_simultaneous(
        ((max_decode_rows, q_num_heads, max_splits, kv_lora_rank + 1), torch.float32),
    )