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, sliding_window=None)

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

Reserving at the worst case (max KV length -> 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). Like forward_mqa, the KV length is clamped to the layer's sliding window.

Source code in vllm/v1/attention/backends/mla/triton_mla.py 
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def _reserve_attn_logits_workspace(
    vllm_config: VllmConfig,
    max_decode_rows: int,
    q_num_heads: int,
    kv_lora_rank: int,
    sliding_window: int | None = None,
) -> None:
    """Pre-size the shared workspace for the decode split-KV attn logits.

    Reserving at the worst case (max KV length -> 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). Like
    ``forward_mqa``, the KV length is clamped to the layer's sliding window.
    """
    if not is_workspace_manager_initialized():
        return
    max_kv_len = vllm_config.model_config.max_model_len
    if sliding_window is not None:
        max_kv_len = min(max_kv_len, sliding_window)
    max_splits = _compute_num_kv_splits(
        max_kv_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),
    )