Skip to content

vllm.kernels.aiter_ops

Functions:

Attributes:

AITER_SUPPORTED = is_aiter_found() module-attribute

Most kernels in this file are supported if AITER is installed.

aiter_lib = Library('vllm_aiter', 'FRAGMENT') module-attribute

This library holds torch custom ops for wrapped AITER ops. Many AITER ops want to remain invisible to torch.compile even after lowering. They are thus wrapped into torch custom ops inside the IR op implementations.

direct_register_aiter_op = functools.partial(direct_register_custom_op, target_lib=aiter_lib) module-attribute

Syntactic sugar for registering AITER custom ops.

rms_add_no_var_16bit_only = lambda x, x_residual, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) and (weight is None or weight.dtype == x.dtype) and flatten_to_2d_is_free(x) and flatten_to_2d_is_free(x_residual) module-attribute

AITER fused_add_rms_norm only supports 16-bit activations and no var_size override. Requires weight dtype to match x dtype, and flattening both the activation and the residual to 2D to be free.

rms_no_var_16bit_only = lambda x, weight, epsilon, variance_size=None: variance_size is None and x.dtype in (torch.float16, torch.bfloat16) and (weight is None or weight.dtype == x.dtype) and flatten_to_2d_is_free(x) module-attribute

AITER rms_norm only supports float16 and bfloat16 acts, no var_size override, requires weight dtype to match x dtype, and requires flattening the activation to 2D to be free.

flatten_to_2d_is_free(x)

Whether x.reshape(-1, x.shape[-1]) is a view with unit-stride rows.

The AITER norm kernels only take dense 2D inputs, so the wrappers below flatten the leading dims with Tensor.reshape, which silently falls back to contiguous() when the flattened shape is not expressible with the existing strides, adding a whole-tensor device copy that costs more than the kernel saves. A strided last dim flattens without a copy but not into a dense buffer, so it is rejected too.

Source code in vllm/kernels/aiter_ops.py 
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
def flatten_to_2d_is_free(x: Tensor) -> bool:
    """Whether ``x.reshape(-1, x.shape[-1])`` is a view with unit-stride rows.

    The AITER norm kernels only take dense 2D inputs, so the wrappers below
    flatten the leading dims with ``Tensor.reshape``, which silently falls back
    to ``contiguous()`` when the flattened shape is not expressible with the
    existing strides, adding a whole-tensor device copy that costs more than
    the kernel saves. A strided last dim flattens without a copy but not into a
    dense buffer, so it is rejected too.
    """
    if x.dim() <= 2:
        return True
    if x.stride(-1) != 1:
        return False
    expected_stride = x.size(-1)
    for i in range(x.dim() - 2, -1, -1):
        if x.size(i) != 1 and x.stride(i) != expected_stride:
            return False
        expected_stride *= x.size(i)
    return True