| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Inkling kernels (NVIDIA). |
| 4 | |
| 5 | ``rmsnorm`` / ``sconv`` import eagerly. The SwiGLU kernels and the FA4 |
| 6 | relative-attention wrapper are exposed lazily to keep this package's |
| 7 | import path lightweight. |
| 8 | """ |
| 9 | |
| 10 | from typing import TYPE_CHECKING |
| 11 | |
| 12 | from .norm import add_rmsnorm, rmsnorm |
| 13 | from .sconv import fused_sconv |
| 14 | |
| 15 | _LAZY_EXPORTS = { |
| 16 | "silu_and_mul_triton": "silu_and_mul", |
| 17 | "sink_silu_mul_epilogue": "silu_and_mul", |
| 18 | "inkling_fa4_rel_attention": "fa4_rel_attention", |
| 19 | "inkling_rel_attention_split_kv_decode": "rel_attention_decode", |
| 20 | } |
| 21 | |
| 22 | if TYPE_CHECKING: |
| 23 | from .fa4_rel_attention import inkling_fa4_rel_attention # noqa: F401 |
| 24 | from .silu_and_mul import ( # noqa: F401 |
| 25 | silu_and_mul_triton, |
| 26 | sink_silu_mul_epilogue, |
| 27 | ) |
| 28 | |
| 29 | __all__ = [ |
| 30 | "add_rmsnorm", |
| 31 | "rmsnorm", |
| 32 | "fused_sconv", |
| 33 | *sorted(_LAZY_EXPORTS), |
| 34 | ] |
| 35 | |
| 36 | |
| 37 | def __getattr__(name: str): |
| 38 | module = _LAZY_EXPORTS.get(name) |
| 39 | if module is not None: |
| 40 | import importlib |
| 41 | |
| 42 | mod = importlib.import_module(f".{module}", __name__) |
| 43 | return getattr(mod, name) |
| 44 | raise AttributeError(f"module {__name__!r} has no attribute {name!r}") |