vllm.model_executor.layers.utils

Utility methods for model layers.

Functions:

• apply_penalties –

  Applies penalties in place to the logits tensor

• cache_aliasing_pad_elems –

  Row-stride pad, in elements, that stops an [n, k] weight aliasing on gfx11,

• lmhead_skinny_guard_permits –

  Should this weight still take the wvSplitK skinny path at n?

• maybe_pad_weight_avoid_cache_aliasing –

  Offset a dense weight's row stride so it stops aliasing.

• warmup_rocm_skinny_gemm_workspaces –

  Eagerly allocate wvSplitKrc's per-device split-K workspace pool.

• wvsplitkrc_dispatch –

  Pick the K-shard split for wvSplitKrc and say whether the shape fits.

_amd_l3_cache_bytes(cache_info=None)

The GPU's L3 (memory-side) cache size in bytes, or None.

The skinny-GEMM weight cap tracks this cache: past it a weight cannot stay resident, and the skinny kernel's cost at n activation rows (~n weight passes) pays HBM round trips instead. The device-properties API exposes only the far smaller L2 (4 MiB on gfx950), so query the L3 through AMD's amdsmi library: 256 MiB on MI355X, 128 MiB on MI300X-class parts. cache_info (an amdsmi cache list) lets tests drive the parse; the default queries the library.

Source code in vllm/model_executor/layers/utils.py

def _amd_l3_cache_bytes(cache_info: list | None = None) -> int | None:
    """The GPU's L3 (memory-side) cache size in bytes, or None.

    The skinny-GEMM weight cap tracks this cache: past it a weight cannot
    stay resident, and the skinny kernel's cost at n activation rows (~n
    weight passes) pays HBM round trips instead. The device-properties
    API exposes only the far smaller L2 (4 MiB on gfx950), so query the
    L3 through AMD's amdsmi library: 256 MiB on MI355X, 128 MiB on
    MI300X-class parts. `cache_info` (an amdsmi cache list) lets tests
    drive the parse; the default queries the library.
    """
    if cache_info is None:
        try:
            import amdsmi

            amdsmi.amdsmi_init()
            try:
                handles = amdsmi.amdsmi_get_processor_handles()
                cache_info = amdsmi.amdsmi_get_gpu_cache_info(handles[0])["cache"]
            finally:
                amdsmi.amdsmi_shut_down()
        except Exception:  # noqa: BLE001 -- optional lib, any failure -> fallback
            return None
    for cache in cache_info:
        if cache.get("cache_level") == 3 and "DATA_CACHE" in cache.get(
            "cache_properties", []
        ):
            return cache["cache_size"] * 1024  # amdsmi reports KB
    return None

_skinny_gemm_max_weight_bytes() cached

Effective skinny-GEMM weight cap in bytes, evaluated once.

VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES > 0 overrides; 0 (the default) derives the cap from the hardware L3 (memory-side) cache size, falling back to _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES where that cannot be read. lmhead_skinny_guard_permits runs per decode step, so this is computed on the first call only (env access in hot paths: #17067). Tests that change the env var or mock the hardware read must call cache_clear().

Source code in vllm/model_executor/layers/utils.py

@functools.lru_cache(maxsize=1)
def _skinny_gemm_max_weight_bytes() -> int:
    """Effective skinny-GEMM weight cap in bytes, evaluated once.

    VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES > 0 overrides; 0 (the default)
    derives the cap from the hardware L3 (memory-side) cache size,
    falling back to _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES where that cannot
    be read. lmhead_skinny_guard_permits runs per decode step, so this is
    computed on the first call only (env access in hot paths: #17067).
    Tests that change the env var or mock the hardware read must call
    cache_clear().
    """
    env_bytes = envs.VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES
    if env_bytes > 0:
        return env_bytes
    return _amd_l3_cache_bytes() or _SKINNY_GEMM_FALLBACK_WEIGHT_BYTES

apply_penalties(logits, prompt_tokens_tensor, output_tokens_tensor, presence_penalties, frequency_penalties, repetition_penalties)

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, )

Source code in vllm/model_executor/layers/utils.py

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

cache_aliasing_pad_elems(k, element_size)

Row-stride pad, in elements, that stops an [n, k] weight aliasing on gfx11, or 0 when the stride does not alias.

Parameters:

• k

  (int) –

  Number of elements per weight row.

• element_size

  (int) –

  Size of one element in bytes.

Returns:

• int –

  The number of elements to add to the row stride.

Source code in vllm/model_executor/layers/utils.py

def cache_aliasing_pad_elems(k: int, element_size: int) -> int:
    """Row-stride pad, in elements, that stops an ``[n, k]`` weight aliasing on gfx11,
    or 0 when the stride does not alias.

    Args:
        k: Number of elements per weight row.
        element_size: Size of one element in bytes.

    Returns:
        The number of elements to add to the row stride.

    """
    if (k * element_size) % 2048 != 0:
        return 0
    return 128 // element_size

lmhead_skinny_guard_permits(weight, n)

Should this weight still take the wvSplitK skinny path at n?

wvSplitK_hf_sml's cost at n activation rows scales as ~n passes over the weight (MI355X measurements in the PR description), so past the last-level (memory-side) cache its effective bandwidth collapses: every pass beyond the first is an HBM round trip. Weights above VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES (default: the memory-side cache size, read from the hardware once at first use, 256 MiB on MI355X; see envs.py for what the cap controls and how to tune it) therefore skip the skinny path at 1 < n <= 5 and fall through to the generic GEMM.

n == 1 always keeps wvSplitK: the n=1 kernel variant shows no pathology on large weights, and its fall-through would be a numerics change with no speed win, so single-row calls stay on the stock path.

Returns True (keep wvSplitK) when n <= 1 or the weight is under the cap; otherwise returns False after emitting a census WARNING (once per (shape, n) per process) naming the redirected site.

Source code in vllm/model_executor/layers/utils.py

def lmhead_skinny_guard_permits(weight: torch.Tensor, n: int) -> bool:
    """Should this weight still take the wvSplitK skinny path at n?

    wvSplitK_hf_sml's cost at n activation rows scales as ~n passes over
    the weight (MI355X measurements in the PR description), so past the
    last-level (memory-side) cache its effective bandwidth collapses:
    every pass beyond the first is an HBM round trip. Weights above
    VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES (default: the memory-side cache
    size, read from the hardware once at first use, 256 MiB on MI355X;
    see envs.py for what the cap controls and how to tune it) therefore
    skip the skinny path at 1 < n <= 5 and fall through to the generic
    GEMM.

    n == 1 always keeps wvSplitK: the n=1 kernel variant shows no
    pathology on large weights, and its fall-through would be a numerics
    change with no speed win, so single-row calls stay on the stock path.

    Returns True (keep wvSplitK) when n <= 1 or the weight is under the
    cap; otherwise returns False after emitting a census WARNING (once per
    (shape, n) per process) naming the redirected site.
    """
    if n <= 1:
        return True
    max_bytes = _skinny_gemm_max_weight_bytes()
    if weight.nbytes <= max_bytes:
        return True
    key = (tuple(weight.shape), n)
    if key not in _skinny_guard_warned:
        _skinny_guard_warned.add(key)
        logger.warning(
            "lm_head skinny-GEMM guard engaged: weight %s (%d bytes) exceeds "
            "the weight cap %d (VLLM_SKINNY_GEMM_MAX_WEIGHT_BYTES, or the "
            "hardware memory-side cache size when unset) at n=%d; skipping "
            "wvSplitK for the generic GEMM fall-through",
            tuple(weight.shape),
            weight.nbytes,
            max_bytes,
            n,
        )
    return False

maybe_pad_weight_avoid_cache_aliasing(weight)

Offset a dense weight's row stride so it stops aliasing.

Returns a view with the original shape whose row stride is one cache line wider, or weight unchanged when the pad does not apply.

Source code in vllm/model_executor/layers/utils.py

def maybe_pad_weight_avoid_cache_aliasing(weight: torch.Tensor) -> torch.Tensor:
    """Offset a dense weight's row stride so it stops aliasing.

    Returns a view with the original shape whose row stride is one cache line
    wider, or ``weight`` unchanged when the pad does not apply.
    """
    from vllm.platforms.rocm import on_gfx11

    if not (current_platform.is_rocm() and on_gfx11()):
        return weight
    if weight.dim() != 2 or not weight.is_contiguous():
        return weight
    pad = cache_aliasing_pad_elems(weight.shape[1], weight.element_size())
    if pad == 0:
        return weight
    return torch.nn.functional.pad(weight, (0, pad))[..., :-pad]

warmup_rocm_skinny_gemm_workspaces(device) cached

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.

Source code in vllm/model_executor/layers/utils.py

@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)

wvsplitkrc_dispatch(n, k, m, cu_count)

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:

• int –

  The CHUNKK the kernel will dispatch with, and whether the CU budget and

• bool –

  split-K workspace admit the shape at all.

Source code in vllm/model_executor/layers/utils.py

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