vllm.v1.attention.backends.mla.flashinfer_mla ¶

def _get_multi_ctas_kv_counter_buffer(
    min_bytes: int, device: torch.device
) -> torch.Tensor:
    """Persistent, zero-initialized trtllm-gen multi-CTA-KV counter buffer.

    trtllm-gen's multi-CTA-KV MLA decode kernel resets these semaphores to zero
    at the end of every launch, so the buffer only needs zeroing once. The
    public ``trtllm_batch_decode_with_kv_cache_mla`` entry point builds a fresh
    runner per call, so without a caller-owned buffer it re-allocates and
    re-zeros this counter on every decode step (a tiny ``FillFunctor<uint8>``
    launch right before the FMHA). Owning it here and passing it in removes that
    per-step launch. ``min_bytes`` is sized to the worst-case batch so the
    buffer is allocated once and never reallocated after CUDA-graph capture.
    """
    global _fi_multi_ctas_kv_counter
    if (
        _fi_multi_ctas_kv_counter is None
        or _fi_multi_ctas_kv_counter.numel() < min_bytes
    ):
        _fi_multi_ctas_kv_counter = torch.zeros(
            min_bytes, dtype=torch.uint8, device=device
        )
    return _fi_multi_ctas_kv_counter

def _select_mla_decode_backend(num_heads: int) -> str | None:
    """cute-dsl for head counts trtllm-gen cannot tile, else None (=> auto)."""
    if not _trtllm_gen_mla_decode_supports_num_heads(num_heads):
        logger.warning_once(
            "trtllm-gen MLA decode does not support num_heads=%d "
            "(query/kv head ratio); falling back to the cute-dsl backend.",
            num_heads,
        )
        return "cute-dsl"
    return None

def _trtllm_gen_mla_decode_supports_num_heads(num_heads: int) -> bool:
    """True if trtllm-gen's MLA decode kernel supports this query head count.

    The kernel groups Q heads into CTAs of ``min(num_heads, tileSizeQ)`` and
    requires ``num_heads`` divisible by that, else raises "The
    numHeadsQ/numHeadsKv is not supported" (flashinfer fmhaKernels.cuh).
    ``tileSizeQ`` is 8/16 for ``num_heads <= 8``/``<= 32`` (SwapsMmaAb) else 64;
    treat ``> 32`` as tile 64 (safe bound). E.g. 96/24 -> False, 48/64/128 -> True.
    """
    if num_heads <= 8:
        tile = 8
    elif num_heads <= 32:
        tile = 16
    else:
        tile = 64
    return num_heads % min(num_heads, tile) == 0