vllm.model_executor.layers.utils ¶

def apply_penalties(
    logits: torch.Tensor,
    prompt_tokens_tensor: torch.Tensor,
    output_tokens_tensor: torch.Tensor,
    presence_penalties: torch.Tensor,
    frequency_penalties: torch.Tensor,
    repetition_penalties: torch.Tensor,
) -> torch.Tensor:
    """
    Applies penalties in place to the logits tensor
    logits : The input logits tensor of shape [num_seqs, vocab_size]
    prompt_tokens_tensor: A tensor containing the prompt tokens. The prompts
        are padded to the maximum prompt length within the batch using
        `vocab_size` as the padding value. The value `vocab_size` is used
        for padding because it does not correspond to any valid token ID
        in the vocabulary.
    output_tokens_tensor: The output tokens tensor.
    presence_penalties: The presence penalties of shape (num_seqs, )
    frequency_penalties: The frequency penalties of shape (num_seqs, )
    repetition_penalties: The repetition penalties of shape (num_seqs, )
    """
    num_seqs, vocab_size = logits.shape
    _, prompt_mask = get_token_bin_counts_and_mask(
        prompt_tokens_tensor, vocab_size, num_seqs
    )
    output_bin_counts, output_mask = get_token_bin_counts_and_mask(
        output_tokens_tensor, vocab_size, num_seqs
    )

    # Apply repetition penalties as a custom op
    from vllm._custom_ops import apply_repetition_penalties

    apply_repetition_penalties(logits, prompt_mask, output_mask, repetition_penalties)

    # We follow the definition in OpenAI API.
    # Refer to https://platform.openai.com/docs/api-reference/parameter-details
    logits -= frequency_penalties.unsqueeze(dim=1) * output_bin_counts
    logits -= presence_penalties.unsqueeze(dim=1) * output_mask
    return logits

@functools.cache
def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None:
    """Eagerly allocate wvSplitKrc's per-device split-K workspace pool.

    wvSplitKrc partitions one per-device allocation into ``kWvSlots`` slots
    (csrc/rocm/skinny_gemms.cu) and hands each stream one on first use, so that
    two streams never share the split-K partials and counters.

    The pool is otherwise created lazily on the first qualifying GEMM
    (csrc/rocm/skinny_gemms.cu), which can be the first real request — after
    the KV cache backing buffer exists. If it landed in that segment's rounding
    tail, it would pin the entire segment at engine shutdown; it could also land
    inside a cudagraph capture, where it would be taken from the graph's private
    pool and its zero-fill would become a replayed graph node.
    """
    from vllm.platforms.rocm import on_gfx950

    if not on_gfx950():
        return
    try:
        x = torch.zeros(16, 1024, dtype=torch.bfloat16, device=device)
        weight = torch.zeros(32, 1024, dtype=torch.bfloat16, device=device)
        ops.wvSplitKrc(x, weight, num_compute_units())
    except Exception:
        logger.debug("wvSplitKrc workspace warmup failed", exc_info=True)

def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, bool]:
    """Pick the K-shard split for wvSplitKrc and say whether the shape fits.

    Mirrors wvSplitKrc() in csrc/rocm/skinny_gemms.cu, which is also where the
    shard cap is explained. Both must pick the same chunkk or the workspace
    check here bounds the wrong k_rnd.

    Returns:
        The CHUNKK the kernel will dispatch with, and whether the CU budget and
        split-K workspace admit the shape at all.
    """
    # Next ^2 of n
    N_p2 = 1 << (n - 1).bit_length()
    # How many of 4 waves in a group can work on same 16 Ms at same time?
    # This reduces the Ms each group works on, i.e. increasing the CUs needed.
    GrpsShrB = min(N_p2 // 16, 4)
    # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), and each
    # working on a 512-shard of K, how many CUs would we need?
    CuNeeded = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) * GrpsShrB

    CHUNKK2_MAX_SHARDS = 11
    shards_chunkk2 = (k + 256 - 1) // 256  # 256-wide shards
    chunkk = (
        2
        if (
            N_p2 != 16
            and CuNeeded * 2 <= cu_count
            and shards_chunkk2 <= CHUNKK2_MAX_SHARDS
        )
        else 1
    )

    # Deterministic reduction stores one fp32 partial per (M, N, k-shard); all
    # of them must fit the split-K workspace.
    k_rnd = (k + 512 // chunkk - 1) // (512 // chunkk)
    fits = N_p2 * m * k_rnd <= 128 * 1024 * 12 and CuNeeded <= cu_count
    return chunkk, fits