Skip to content

vllm.kernels.aiter_ops ¶

AITER_SUPPORTED `module-attribute` ¶

AITER_SUPPORTED = is_aiter_found()

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

aiter_lib `module-attribute` ¶

aiter_lib = Library('vllm_aiter', 'FRAGMENT')

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 `module-attribute` ¶

direct_register_aiter_op = partial(
    direct_register_custom_op, target_lib=aiter_lib
)

Syntactic sugar for registering AITER custom ops.

rms_no_var_16bit_only `module-attribute` ¶

rms_no_var_16bit_only = (
    lambda x, weight, epsilon, variance_size=None: (
        variance_size is None
        and dtype in (float16, bfloat16)
    )
)

AITER rms_norm only supports float16 and bfloat16 acts and no var_size override.