vllm.v1.attention.ops.triton_merge_attn_states ¶

def mask_empty_context(
    lse: torch.Tensor,
    output: torch.Tensor,
    query_start_loc: torch.Tensor,
    context_start_loc: torch.Tensor,
) -> None:
    """Neutralize context chunks that cover no keys before merging.

    A prefill query whose context chunk is empty attended to no keys, so its
    partial attention is undefined: the backend leaves the output rows as
    uninitialized scratch (which may hold NaN/Inf) even when it reports an LSE
    of -inf. Sanitize both here so ``merge_attn_states`` can stay generic:
    force the LSE to -inf (zero softmax weight) and zero the undefined output
    rows (so a zero weight cannot combine with NaN/Inf). Emptiness is derived
    from the context offsets, not from the -inf LSE, so no merge kernel has to
    reason about undefined partials.

    Args:
        lse: Chunk log-sum-exp, shape [num_heads, num_tokens].
        output: Chunk attention output, shape [num_tokens, num_heads, ...].
        query_start_loc: Prefill query cumulative offsets, shape [num_reqs + 1].
        context_start_loc: Chunk context cumulative offsets,
            shape [num_reqs + 1]; an empty chunk has a zero-length span.
    """
    num_heads, num_tokens = lse.shape
    num_reqs = query_start_loc.shape[0] - 1
    block_size = 128
    # Reserve the worst-case number of request-local blocks.
    num_query_blocks = num_tokens // block_size + num_reqs
    is_empty = torch.zeros(num_tokens, dtype=torch.bool, device=lse.device)
    mask_empty_context_kernel[(num_query_blocks,)](
        lse,
        is_empty,
        query_start_loc,
        context_start_loc,
        lse.stride(0),
        lse.stride(1),
        num_reqs,
        NUM_HEADS=num_heads,
        BLOCK_SIZE=block_size,
        BLOCK_HEADS=8,
        num_warps=8,
    )
    output.masked_fill_(is_empty[:, None, None], 0.0)